Add support for LibSQL remote (#11385)

* Add support for remote LibSQL * Add support for local memory DB * Add some tests * Add push support * Fix switch cascading * Update .changeset/healthy-boxes-poke.md Co-authored-by: Sarah Rainsberger <[email protected]> * Update packages/db/src/runtime/db-client.ts [skip ci] Co-authored-by: Emanuele Stoppa <[email protected]> * Use independent env vars for LibSQL and Studio backends * Expand comment regarding missing table * Apply suggestions from code review Co-authored-by: Sarah Rainsberger <[email protected]> --------- Co-authored-by: Sarah Rainsberger <[email protected]> Co-authored-by: Emanuele Stoppa <[email protected]>
withastro · Aug 28, 2024 · d6611e8 · d6611e8
1 parent 3c0ca8d
commit d6611e8
Show file tree

Hide file tree

Showing 9 changed files with 220 additions and 27 deletions.
diff --git a/.changeset/healthy-boxes-poke.md b/.changeset/healthy-boxes-poke.md
@@ -0,0 +1,14 @@
+---
+'@astrojs/db': minor
+---
+
+Adds support for connecting Astro DB to any remote LibSQL server. This allows Astro DB to be used with self-hosting and air-gapped deployments.
+
+To connect Astro DB to a remote LibSQL server instead of Studio, set the following environment variables:
+
+- `ASTRO_DB_REMOTE_URL`: the connection URL to your LibSQL server
+- `ASTRO_DB_APP_TOKEN`: the auth token to your LibSQL server
+
+Details of the LibSQL connection can be configured using the connection URL. For example, `memory:?syncUrl=libsql%3A%2F%2Fdb-server.example.com` would create an in-memory embedded replica for the LibSQL DB on `libsql://db-server.example.com`.
+
+For more details, please visit [the Astro DB documentation](https://docs.astro.build/en/guides/astro-db/#libsql)
diff --git a/packages/db/src/core/cli/commands/push/index.ts b/packages/db/src/core/cli/commands/push/index.ts
@@ -1,11 +1,13 @@
 import { getManagedAppTokenOrExit } from '@astrojs/studio';
 import type { AstroConfig } from 'astro';
+import { sql } from 'drizzle-orm';
 import prompts from 'prompts';
 import type { Arguments } from 'yargs-parser';
+import { createRemoteDatabaseClient } from '../../../../runtime/index.js';
 import { safeFetch } from '../../../../runtime/utils.js';
 import { MIGRATION_VERSION } from '../../../consts.js';
 import type { DBConfig, DBSnapshot } from '../../../types.js';
-import { type Result, getRemoteDatabaseUrl } from '../../../utils.js';
+import { type Result, getRemoteDatabaseInfo } from '../../../utils.js';
 import {
 	createCurrentSnapshot,
 	createEmptySnapshot,
@@ -87,7 +89,7 @@ async function pushSchema({
 	isDryRun: boolean;
 	currentSnapshot: DBSnapshot;
 }) {
-	const requestBody = {
+	const requestBody: RequestBody = {
 		snapshot: currentSnapshot,
 		sql: statements,
 		version: MIGRATION_VERSION,
@@ -96,7 +98,47 @@ async function pushSchema({
 		console.info('[DRY RUN] Batch query:', JSON.stringify(requestBody, null, 2));
 		return new Response(null, { status: 200 });
 	}
-	const url = new URL('/db/push', getRemoteDatabaseUrl());
+
+	const dbInfo = getRemoteDatabaseInfo();
+
+	return dbInfo.type === 'studio'
+		? pushToStudio(requestBody, appToken, dbInfo.url)
+		: pushToDb(requestBody, appToken, dbInfo.url);
+}
+
+type RequestBody = {
+	snapshot: DBSnapshot;
+	sql: string[];
+	version: string;
+};
+
+async function pushToDb(requestBody: RequestBody, appToken: string, remoteUrl: string) {
+	const client = createRemoteDatabaseClient({
+		dbType: 'libsql',
+		appToken,
+		remoteUrl,
+	});
+
+	await client.run(sql`create table if not exists _astro_db_snapshot (
+		id INTEGER PRIMARY KEY AUTOINCREMENT,
+		version TEXT,
+		snapshot BLOB
+	);`);
+
+	await client.transaction(async (tx) => {
+		for (const stmt of requestBody.sql) {
+			await tx.run(sql.raw(stmt));
+		}
+
+		await tx.run(sql`insert into _astro_db_snapshot (version, snapshot) values (
+			${requestBody.version},
+			${JSON.stringify(requestBody.snapshot)}
+		)`);
+	});
+}
+
+async function pushToStudio(requestBody: RequestBody, appToken: string, remoteUrl: string) {
+	const url = new URL('/db/push', remoteUrl);
 	const response = await safeFetch(
 		url,
 		{

diff --git a/packages/db/src/core/cli/commands/shell/index.ts b/packages/db/src/core/cli/commands/shell/index.ts
@@ -10,7 +10,7 @@ import { normalizeDatabaseUrl } from '../../../../runtime/index.js';
 import { DB_PATH } from '../../../consts.js';
 import { SHELL_QUERY_MISSING_ERROR } from '../../../errors.js';
 import type { DBConfigInput } from '../../../types.js';
-import { getAstroEnv, getRemoteDatabaseUrl } from '../../../utils.js';
+import { getAstroEnv, getRemoteDatabaseInfo } from '../../../utils.js';
 
 export async function cmd({
 	flags,
@@ -25,9 +25,14 @@ export async function cmd({
 		console.error(SHELL_QUERY_MISSING_ERROR);
 		process.exit(1);
 	}
+	const dbInfo = getRemoteDatabaseInfo();
 	if (flags.remote) {
 		const appToken = await getManagedAppTokenOrExit(flags.token);
-		const db = createRemoteDatabaseClient(appToken.token, getRemoteDatabaseUrl());
+		const db = createRemoteDatabaseClient({
+			dbType: dbInfo.type,
+			remoteUrl: dbInfo.url,
+			appToken: appToken.token,
+		});
 		const result = await db.run(sql.raw(query));
 		await appToken.destroy();
 		console.log(result);
@@ -37,7 +42,7 @@ export async function cmd({
 			ASTRO_DATABASE_FILE,
 			new URL(DB_PATH, astroConfig.root).href,
 		);
-		const db = createLocalDatabaseClient({ dbUrl });
+		const db = createLocalDatabaseClient({ dbUrl, enableTransations: dbInfo.type === 'libsql' });
 		const result = await db.run(sql.raw(query));
 		console.log(result);
 	}

diff --git a/packages/db/src/core/cli/migration-queries.ts b/packages/db/src/core/cli/migration-queries.ts
@@ -1,9 +1,11 @@
 import deepDiff from 'deep-diff';
+import { sql } from 'drizzle-orm';
 import { SQLiteAsyncDialect } from 'drizzle-orm/sqlite-core';
 import * as color from 'kleur/colors';
 import { customAlphabet } from 'nanoid';
 import stripAnsi from 'strip-ansi';
 import { hasPrimaryKey } from '../../runtime/index.js';
+import { createRemoteDatabaseClient } from '../../runtime/index.js';
 import { isSerializedSQL } from '../../runtime/types.js';
 import { safeFetch } from '../../runtime/utils.js';
 import { MIGRATION_VERSION } from '../consts.js';
@@ -33,7 +35,7 @@ import type {
 	ResolvedIndexes,
 	TextColumn,
 } from '../types.js';
-import { type Result, getRemoteDatabaseUrl } from '../utils.js';
+import { type Result, getRemoteDatabaseInfo } from '../utils.js';
 
 const sqlite = new SQLiteAsyncDialect();
 const genTempTableName = customAlphabet('abcdefghijklmnopqrstuvwxyz', 10);
@@ -422,12 +424,49 @@ function hasRuntimeDefault(column: DBColumn): column is DBColumnWithDefault {
 	return !!(column.schema.default && isSerializedSQL(column.schema.default));
 }
 
-export async function getProductionCurrentSnapshot({
-	appToken,
-}: {
+export function getProductionCurrentSnapshot(options: {
 	appToken: string;
 }): Promise<DBSnapshot | undefined> {
-	const url = new URL('/db/schema', getRemoteDatabaseUrl());
+	const dbInfo = getRemoteDatabaseInfo();
+
+	return dbInfo.type === 'studio'
+		? getStudioCurrentSnapshot(options.appToken, dbInfo.url)
+		: getDbCurrentSnapshot(options.appToken, dbInfo.url);
+}
+
+async function getDbCurrentSnapshot(
+	appToken: string,
+	remoteUrl: string
+): Promise<DBSnapshot | undefined> {
+	const client = createRemoteDatabaseClient({
+		dbType: 'libsql',
+		appToken,
+		remoteUrl,
+	});
+
+	try {
+		const res = await client.get<{ snapshot: string }>(
+			// Latest snapshot
+			sql`select snapshot from _astro_db_snapshot order by id desc limit 1;`
+		);
+
+		return JSON.parse(res.snapshot);
+	} catch (error: any) {
+		if (error.code === 'SQLITE_UNKNOWN') {
+			// If the schema was never pushed to the database yet the table won't exist.
+			// Treat a missing snapshot table as an empty table.
+			return;
+		}
+
+		throw error;
+	}
+}
+
+async function getStudioCurrentSnapshot(
+	appToken: string,
+	remoteUrl: string
+): Promise<DBSnapshot | undefined> {
+	const url = new URL('/db/schema', remoteUrl);
 
 	const response = await safeFetch(
 		url,

diff --git a/packages/db/src/core/integration/vite-plugin-db.ts b/packages/db/src/core/integration/vite-plugin-db.ts
@@ -9,7 +9,7 @@ import { DB_PATH, RUNTIME_IMPORT, RUNTIME_VIRTUAL_IMPORT, VIRTUAL_MODULE_ID } fr
 import { getResolvedFileUrl } from '../load-file.js';
 import { SEED_DEV_FILE_NAME, getCreateIndexQueries, getCreateTableQuery } from '../queries.js';
 import type { DBTables } from '../types.js';
-import { type VitePlugin, getAstroEnv, getDbDirectoryUrl, getRemoteDatabaseUrl } from '../utils.js';
+import { type VitePlugin, getAstroEnv, getDbDirectoryUrl, getRemoteDatabaseInfo } from '../utils.js';
 
 export const resolved = {
 	module: '\0' + VIRTUAL_MODULE_ID,
@@ -119,12 +119,13 @@ export function getLocalVirtualModContents({
 	tables: DBTables;
 	root: URL;
 }) {
+	const dbInfo = getRemoteDatabaseInfo();
 	const dbUrl = new URL(DB_PATH, root);
 	return `
 import { asDrizzleTable, createLocalDatabaseClient, normalizeDatabaseUrl } from ${RUNTIME_IMPORT};
 
 const dbUrl = normalizeDatabaseUrl(import.meta.env.ASTRO_DATABASE_FILE, ${JSON.stringify(dbUrl)});
-export const db = createLocalDatabaseClient({ dbUrl });
+export const db = createLocalDatabaseClient({ dbUrl, enableTransactions: ${dbInfo.url === 'libsql'} });
 
 export * from ${RUNTIME_VIRTUAL_IMPORT};
 
@@ -142,30 +143,40 @@ export function getStudioVirtualModContents({
 	isBuild: boolean;
 	output: AstroConfig['output'];
 }) {
+	const dbInfo = getRemoteDatabaseInfo();
+
 	function appTokenArg() {
 		if (isBuild) {
+			const envPrefix = dbInfo.type === 'studio' ? 'ASTRO_STUDIO' : 'ASTRO_DB';
 			if (output === 'server') {
 				// In production build, always read the runtime environment variable.
-				return 'process.env.ASTRO_STUDIO_APP_TOKEN';
+				return `process.env.${envPrefix}_APP_TOKEN`;
 			} else {
 				// Static mode or prerendering needs the local app token.
-				return `process.env.ASTRO_STUDIO_APP_TOKEN ?? ${JSON.stringify(appToken)}`;
+				return `process.env.${envPrefix}_APP_TOKEN ?? ${JSON.stringify(appToken)}`;
 			}
 		} else {
 			return JSON.stringify(appToken);
 		}
 	}
 
 	function dbUrlArg() {
-		const dbStr = JSON.stringify(getRemoteDatabaseUrl());
+		const dbStr = JSON.stringify(dbInfo.url);
+
 		// Allow overriding, mostly for testing
-		return `import.meta.env.ASTRO_STUDIO_REMOTE_DB_URL ?? ${dbStr}`;
+		return dbInfo.type === 'studio'
+			? `import.meta.env.ASTRO_STUDIO_REMOTE_DB_URL ?? ${dbStr}`
+			: `import.meta.env.ASTRO_DB_REMOTE_URL ?? ${dbStr}`;
 	}
 
 	return `
 import {asDrizzleTable, createRemoteDatabaseClient} from ${RUNTIME_IMPORT};
 
-export const db = await createRemoteDatabaseClient(${appTokenArg()}, ${dbUrlArg()});
+export const db = await createRemoteDatabaseClient({
+  dbType: ${JSON.stringify(dbInfo.type)},
+  remoteUrl: ${dbUrlArg()},
+  appToken: ${appTokenArg()},
+});
 
 export * from ${RUNTIME_VIRTUAL_IMPORT};
 
@@ -187,9 +198,10 @@ function getStringifiedTableExports(tables: DBTables) {
 const sqlite = new SQLiteAsyncDialect();
 
 async function recreateTables({ tables, root }: { tables: LateTables; root: URL }) {
+	const dbInfo = getRemoteDatabaseInfo();
 	const { ASTRO_DATABASE_FILE } = getAstroEnv();
 	const dbUrl = normalizeDatabaseUrl(ASTRO_DATABASE_FILE, new URL(DB_PATH, root).href);
-	const db = createLocalDatabaseClient({ dbUrl });
+	const db = createLocalDatabaseClient({ dbUrl, enableTransations: dbInfo.type === 'libsql' });
 	const setupQueries: SQL[] = [];
 	for (const [name, table] of Object.entries(tables.get() ?? {})) {
 		const dropQuery = sql.raw(`DROP TABLE IF EXISTS ${sqlite.escapeName(name)}`);

diff --git a/packages/db/src/core/load-file.ts b/packages/db/src/core/load-file.ts
@@ -146,6 +146,7 @@ export async function bundleFile({
 		metafile: true,
 		define: {
 			'import.meta.env.ASTRO_STUDIO_REMOTE_DB_URL': 'undefined',
+			'import.meta.env.ASTRO_DB_REMOTE_DB_URL': 'undefined',
 			'import.meta.env.ASTRO_DATABASE_FILE': JSON.stringify(ASTRO_DATABASE_FILE ?? ''),
 		},
 		plugins: [

diff --git a/packages/db/src/core/utils.ts b/packages/db/src/core/utils.ts
@@ -10,9 +10,29 @@ export function getAstroEnv(envMode = ''): Record<`ASTRO_${string}`, string> {
 	return env;
 }
 
-export function getRemoteDatabaseUrl(): string {
-	const env = getAstroStudioEnv();
-	return env.ASTRO_STUDIO_REMOTE_DB_URL || 'https://db.services.astro.build';
+export type RemoteDatabaseInfo = {
+	type: 'libsql' | 'studio';
+	url: string;
+};
+
+export function getRemoteDatabaseInfo(): RemoteDatabaseInfo {
+	const astroEnv = getAstroEnv();
+	const studioEnv = getAstroStudioEnv();
+
+	if (studioEnv.ASTRO_STUDIO_REMOTE_DB_URL) return {
+		type: 'studio',
+		url: studioEnv.ASTRO_STUDIO_REMOTE_DB_URL,
+	};
+
+	if (astroEnv.ASTRO_DB_REMOTE_URL) return {
+		type: 'libsql',
+		url: astroEnv.ASTRO_DB_REMOTE_URL,
+	};
+
+	return {
+		type: 'studio',
+		url: 'https://db.services.astro.build',
+	};
 }
 
 export function getDbDirectoryUrl(root: URL | string) {

diff --git a/packages/db/src/runtime/db-client.ts b/packages/db/src/runtime/db-client.ts
@@ -1,5 +1,5 @@
 import type { InStatement } from '@libsql/client';
-import { createClient } from '@libsql/client';
+import { type Config as LibSQLConfig, createClient } from '@libsql/client';
 import type { LibSQLDatabase } from 'drizzle-orm/libsql';
 import { drizzle as drizzleLibsql } from 'drizzle-orm/libsql';
 import { type SqliteRemoteDatabase, drizzle as drizzleProxy } from 'drizzle-orm/sqlite-proxy';
@@ -18,12 +18,19 @@ function applyTransactionNotSupported(db: SqliteRemoteDatabase) {
 	});
 }
 
-export function createLocalDatabaseClient({ dbUrl }: { dbUrl: string }): LibSQLDatabase {
-	const url = isWebContainer ? 'file:content.db' : dbUrl;
+type LocalDbClientOptions = {
+	dbUrl: string;
+	enableTransations: boolean;
+};
+
+export function createLocalDatabaseClient(options: LocalDbClientOptions): LibSQLDatabase {
+	const url = isWebContainer ? 'file:content.db' : options.dbUrl;
 	const client = createClient({ url });
 	const db = drizzleLibsql(client);
 
-	applyTransactionNotSupported(db);
+	if (!options.enableTransations) {
+		applyTransactionNotSupported(db);
+	}
 	return db;
 }
 
@@ -35,7 +42,33 @@ const remoteResultSchema = z.object({
 	lastInsertRowid: z.unknown().optional(),
 });
 
-export function createRemoteDatabaseClient(appToken: string, remoteDbURL: string) {
+type RemoteDbClientOptions = {
+	dbType: 'studio' | 'libsql',
+	appToken: string,
+	remoteUrl: string | URL,
+}
+
+export function createRemoteDatabaseClient(options: RemoteDbClientOptions) {
+	const remoteUrl = new URL(options.remoteUrl);
+
+	return options.dbType === 'studio'
+		? createStudioDatabaseClient(options.appToken, remoteUrl)
+		: createRemoteLibSQLClient(options.appToken, remoteUrl);
+}
+
+function createRemoteLibSQLClient(appToken: string, remoteDbURL: URL) {
+	const options: Partial<LibSQLConfig> = Object.fromEntries(remoteDbURL.searchParams.entries());
+	remoteDbURL.search = '';
+
+	const client = createClient({
+		...options,
+		authToken: appToken,
+		url: remoteDbURL.protocol === 'memory:' ? ':memory:' : remoteDbURL.toString(),
+	});
+  return drizzleLibsql(client);
+}
+
+function createStudioDatabaseClient(appToken: string, remoteDbURL: URL) {
 	if (appToken == null) {
 		throw new Error(`Cannot create a remote client: missing app token.`);
 	}