Add --remote flag for remote connection

.changeset/healthy-taxis-applaud.md

@@ -0,0 +1,5 @@
+---
+"@astrojs/db": minor
+---
+
+Introduce `astro build --remote` to build with a remote database connection. Running `astro build` plain will use a local database file, and `--remote` will authenticate with a studio app token.

packages/db/src/core/cli/commands/execute/index.ts

@@ -2,7 +2,10 @@ import { existsSync } from 'node:fs';
 import type { AstroConfig } from 'astro';
 import type { Arguments } from 'yargs-parser';
 import { FILE_NOT_FOUND_ERROR, MISSING_EXECUTE_PATH_ERROR } from '../../../errors.js';
-import { getStudioVirtualModContents } from '../../../integration/vite-plugin-db.js';
+import {
+	getLocalVirtualModContents,
+	getStudioVirtualModContents,
+} from '../../../integration/vite-plugin-db.js';
 import { bundleFile, importBundledFile } from '../../../load-file.js';
 import { getManagedAppTokenOrExit } from '../../../tokens.js';
 import { type DBConfig } from '../../../types.js';
@@ -28,12 +31,20 @@ export async function cmd({
 		process.exit(1);
 	}
 
-	const appToken = await getManagedAppTokenOrExit(flags.token);
-
-	const virtualModContents = getStudioVirtualModContents({
-		tables: dbConfig.tables ?? {},
-		appToken: appToken.token,
-	});
+	let virtualModContents: string;
+	if (flags.remote) {
+		const appToken = await getManagedAppTokenOrExit(flags.token);
+		virtualModContents = getStudioVirtualModContents({
+			tables: dbConfig.tables ?? {},
+			appToken: appToken.token,
+		});
+	} else {
+		virtualModContents = getLocalVirtualModContents({
+			tables: dbConfig.tables ?? {},
+			root: astroConfig.root,
+			shouldSeed: false,
+		});
+	}
 	const { code } = await bundleFile({ virtualModContents, root: astroConfig.root, fileUrl });
 	// Executable files use top-level await. Importing will run the file.
 	await importBundledFile({ code, root: astroConfig.root });

packages/db/src/core/cli/commands/shell/index.ts

@@ -1,23 +1,38 @@
 import type { AstroConfig } from 'astro';
 import { sql } from 'drizzle-orm';
 import type { Arguments } from 'yargs-parser';
-import { createRemoteDatabaseClient } from '../../../../runtime/db-client.js';
+import {
+	createRemoteDatabaseClient,
+	createLocalDatabaseClient,
+} from '../../../../runtime/db-client.js';
 import { getManagedAppTokenOrExit } from '../../../tokens.js';
 import type { DBConfigInput } from '../../../types.js';
 import { getRemoteDatabaseUrl } from '../../../utils.js';
+import { DB_PATH } from '../../../consts.js';
+import { SHELL_QUERY_MISSING_ERROR } from '../../../errors.js';
 
 export async function cmd({
 	flags,
+	astroConfig,
 }: {
 	dbConfig: DBConfigInput;
 	astroConfig: AstroConfig;
 	flags: Arguments;
 }) {
 	const query = flags.query;
-	const appToken = await getManagedAppTokenOrExit(flags.token);
-	const db = createRemoteDatabaseClient(appToken.token, getRemoteDatabaseUrl());
-	// Temporary: create the migration table just in case it doesn't exist
-	const result = await db.run(sql.raw(query));
-	await appToken.destroy();
-	console.log(result);
+	if (!query) {
+		console.error(SHELL_QUERY_MISSING_ERROR);
+		process.exit(1);
+	}
+	if (flags.remote) {
+		const appToken = await getManagedAppTokenOrExit(flags.token);
+		const db = createRemoteDatabaseClient(appToken.token, getRemoteDatabaseUrl());
+		const result = await db.run(sql.raw(query));
+		await appToken.destroy();
+		console.log(result);
+	} else {
+		const db = createLocalDatabaseClient({ dbUrl: new URL(DB_PATH, astroConfig.root).href });
+		const result = await db.run(sql.raw(query));
+		console.log(result);
+	}
 }

packages/db/src/core/consts.ts

@@ -1,3 +1,4 @@
+import { randomUUID } from 'node:crypto';
 import { readFileSync } from 'node:fs';
 
 export const PACKAGE_NAME = JSON.parse(
@@ -11,7 +12,9 @@ export const DB_TYPES_FILE = 'db-types.d.ts';
 
 export const VIRTUAL_MODULE_ID = 'astro:db';
 
-export const DB_PATH = '.astro/content.db';
+export const DB_PATH = `.astro/${
+	process.env.ASTRO_TEST_RANDOM_DB_ID ? randomUUID() : 'content.db'
+}`;
 
 export const CONFIG_FILE_NAMES = ['config.ts', 'config.js', 'config.mts', 'config.mjs'];
 

packages/db/src/core/errors.ts

@@ -1,4 +1,4 @@
-import { bold, cyan, green, red, yellow } from 'kleur/colors';
+import { bold, cyan, red } from 'kleur/colors';
 
 export const MISSING_SESSION_ID_ERROR = `${red('▶ Login required!')}
 
@@ -33,6 +33,10 @@ export const RENAME_COLUMN_ERROR = (oldSelector: string, newSelector: string) =>
 export const FILE_NOT_FOUND_ERROR = (path: string) =>
 	`${red('▶ File not found:')} ${bold(path)}\n`;
 
+export const SHELL_QUERY_MISSING_ERROR = `${red(
+	'▶ Please provide a query to execute using the --query flag.'
+)}\n`;
+
 export const SEED_ERROR = (error: string) => {
 	return `${red(`Error while seeding database:`)}\n\n${error}`;
 };

packages/db/src/core/integration/index.ts

@@ -14,6 +14,7 @@ import { fileURLIntegration } from './file-url.js';
 import { typegen } from './typegen.js';
 import { type LateTables, vitePluginDb } from './vite-plugin-db.js';
 import { vitePluginInjectEnvTs } from './vite-plugin-inject-env-ts.js';
+import parseArgs from 'yargs-parser';
 
 function astroDBIntegration(): AstroIntegration {
 	let connectToStudio = false;
@@ -40,7 +41,8 @@ function astroDBIntegration(): AstroIntegration {
 				if (command === 'preview') return;
 
 				let dbPlugin: VitePlugin | undefined = undefined;
-				connectToStudio = command === 'build';
+				const args = parseArgs(process.argv.slice(3));
+				connectToStudio = args['remote'];
 
 				if (connectToStudio) {
 					appToken = await getManagedAppTokenOrExit();
@@ -68,6 +70,8 @@ function astroDBIntegration(): AstroIntegration {
 				});
 			},
 			'astro:config:done': async ({ config }) => {
+				if (command === 'preview') return;
+
 				// TODO: refine where we load tables
 				// @matthewp: may want to load tables by path at runtime
 				const { mod, dependencies } = await loadDbConfigFile(config.root);
@@ -78,7 +82,7 @@ function astroDBIntegration(): AstroIntegration {
 				// TODO: resolve integrations here?
 				tables.get = () => dbConfig.tables ?? {};
 
-				if (!connectToStudio && !process.env.TEST_IN_MEMORY_DB) {
+				if (!connectToStudio) {
 					const dbUrl = new URL(DB_PATH, config.root);
 					if (existsSync(dbUrl)) {
 						await rm(dbUrl);

packages/db/src/core/integration/vite-plugin-db.ts

@@ -1,15 +1,24 @@
 import { fileURLToPath } from 'node:url';
 import { normalizePath } from 'vite';
-import { SEED_DEV_FILE_NAME } from '../../runtime/queries.js';
+import {
+	SEED_DEV_FILE_NAME,
+	getCreateIndexQueries,
+	getCreateTableQuery,
+} from '../../runtime/queries.js';
 import { DB_PATH, RUNTIME_CONFIG_IMPORT, RUNTIME_IMPORT, VIRTUAL_MODULE_ID } from '../consts.js';
 import type { DBTables } from '../types.js';
 import { type VitePlugin, getDbDirectoryUrl, getRemoteDatabaseUrl } from '../utils.js';
+import { createLocalDatabaseClient } from '../../runtime/db-client.js';
+import { type SQL, sql } from 'drizzle-orm';
+import type { SqliteDB } from '../../runtime/index.js';
+import { SQLiteAsyncDialect } from 'drizzle-orm/sqlite-core';
 
-const LOCAL_DB_VIRTUAL_MODULE_ID = 'astro:local';
+const WITH_SEED_VIRTUAL_MODULE_ID = 'astro:db:seed';
 
-const resolvedVirtualModuleId = '\0' + VIRTUAL_MODULE_ID;
-const resolvedLocalDbVirtualModuleId = LOCAL_DB_VIRTUAL_MODULE_ID + '/local-db';
-const resolvedSeedVirtualModuleId = '\0' + VIRTUAL_MODULE_ID + '?shouldSeed';
+const resolved = {
+	virtual: '\0' + VIRTUAL_MODULE_ID,
+	seedVirtual: '\0' + WITH_SEED_VIRTUAL_MODULE_ID,
+};
 
 export type LateTables = {
 	get: () => DBTables;
@@ -32,34 +41,36 @@ type VitePluginDBParams =
 
 export function vitePluginDb(params: VitePluginDBParams): VitePlugin {
 	const srcDirPath = normalizePath(fileURLToPath(params.srcDir));
+	const seedFilePaths = SEED_DEV_FILE_NAME.map((name) =>
+		normalizePath(fileURLToPath(new URL(name, getDbDirectoryUrl(params.root))))
+	);
 	return {
 		name: 'astro:db',
 		enforce: 'pre',
 		async resolveId(id, rawImporter) {
-			if (id === LOCAL_DB_VIRTUAL_MODULE_ID) return resolvedLocalDbVirtualModuleId;
 			if (id !== VIRTUAL_MODULE_ID) return;
-			if (params.connectToStudio) return resolvedVirtualModuleId;
+			if (params.connectToStudio) return resolved.virtual;
 
 			const importer = rawImporter ? await this.resolve(rawImporter) : null;
-			if (!importer) return resolvedVirtualModuleId;
+			if (!importer) return resolved.virtual;
 
 			if (importer.id.startsWith(srcDirPath)) {
 				// Seed only if the importer is in the src directory.
 				// Otherwise, we may get recursive seed calls (ex. import from db/seed.ts).
-				return resolvedSeedVirtualModuleId;
+				return resolved.seedVirtual;
 			}
-			return resolvedVirtualModuleId;
+			return resolved.virtual;
 		},
-		load(id) {
-			if (id === resolvedLocalDbVirtualModuleId) {
-				const dbUrl = new URL(DB_PATH, params.root);
-				return `import { createLocalDatabaseClient } from ${RUNTIME_IMPORT};
-				const dbUrl = ${JSON.stringify(dbUrl)};
-
-				export const db = createLocalDatabaseClient({ dbUrl });`;
+		async load(id) {
+			// Recreate tables whenever a seed file is loaded.
+			if (seedFilePaths.some((f) => id === f)) {
+				await recreateTables({
+					db: createLocalDatabaseClient({ dbUrl: new URL(DB_PATH, params.root).href }),
+					tables: params.tables.get(),
+				});
 			}
 
-			if (id !== resolvedVirtualModuleId && id !== resolvedSeedVirtualModuleId) return;
+			if (id !== resolved.virtual && id !== resolved.seedVirtual) return;
 
 			if (params.connectToStudio) {
 				return getStudioVirtualModContents({
@@ -70,7 +81,7 @@ export function vitePluginDb(params: VitePluginDBParams): VitePlugin {
 			return getLocalVirtualModContents({
 				root: params.root,
 				tables: params.tables.get(),
-				shouldSeed: id === resolvedSeedVirtualModuleId,
+				shouldSeed: id === resolved.seedVirtual,
 			});
 		},
 	};
@@ -82,6 +93,7 @@ export function getConfigVirtualModContents() {
 
 export function getLocalVirtualModContents({
 	tables,
+	root,
 	shouldSeed,
 }: {
 	tables: DBTables;
@@ -94,19 +106,19 @@ export function getLocalVirtualModContents({
 		(name) => new URL(name, getDbDirectoryUrl('file:///')).pathname
 	);
 
+	const dbUrl = new URL(DB_PATH, root);
 	return `
-import { asDrizzleTable, seedLocal } from ${RUNTIME_IMPORT};
-import { db as _db } from ${JSON.stringify(LOCAL_DB_VIRTUAL_MODULE_ID)};
+import { asDrizzleTable, createLocalDatabaseClient } from ${RUNTIME_IMPORT};
+${shouldSeed ? `import { seedLocal } from ${RUNTIME_IMPORT};` : ''}
 
-export const db = _db;
+const dbUrl = ${JSON.stringify(dbUrl)};
+export const db = createLocalDatabaseClient({ dbUrl });
 
 ${
 	shouldSeed
 		? `await seedLocal({
-	db: _db,
-	tables: ${JSON.stringify(tables)},
-	fileGlob: import.meta.glob(${JSON.stringify(seedFilePaths)}),
-})`
+	fileGlob: import.meta.glob(${JSON.stringify(seedFilePaths)}, { eager: true }),
+});`
 		: ''
 }
 
@@ -146,3 +158,19 @@ function getStringifiedCollectionExports(tables: DBTables) {
 		)
 		.join('\n');
 }
+
+const sqlite = new SQLiteAsyncDialect();
+
+async function recreateTables({ db, tables }: { db: SqliteDB; tables: DBTables }) {
+	const setupQueries: SQL[] = [];
+	for (const [name, table] of Object.entries(tables)) {
+		const dropQuery = sql.raw(`DROP TABLE IF EXISTS ${sqlite.escapeName(name)}`);
+		const createQuery = sql.raw(getCreateTableQuery(name, table));
+		const indexQueries = getCreateIndexQueries(name, table);
+		setupQueries.push(dropQuery, createQuery, ...indexQueries.map((s) => sql.raw(s)));
+	}
+	await db.batch([
+		db.run(sql`pragma defer_foreign_keys=true;`),
+		...setupQueries.map((q) => db.run(q)),
+	]);
+}

packages/db/src/runtime/db-client.ts

@@ -9,7 +9,7 @@ const isWebContainer = !!process.versions?.webcontainer;
 
 export function createLocalDatabaseClient({ dbUrl }: { dbUrl: string }): LibSQLDatabase {
 	const url = isWebContainer ? 'file:content.db' : dbUrl;
-	const client = createClient({ url: process.env.TEST_IN_MEMORY_DB ? ':memory:' : url });
+	const client = createClient({ url });
 	const db = drizzleLibsql(client);
 
 	return db;

packages/db/src/runtime/index.ts

@@ -11,12 +11,36 @@ import {
 } from 'drizzle-orm/sqlite-core';
 import { type DBColumn, type DBTable } from '../core/types.js';
 import { type SerializedSQL, isSerializedSQL } from './types.js';
+import { SEED_DEFAULT_EXPORT_ERROR, SEED_ERROR } from '../core/errors.js';
+import { LibsqlError } from '@libsql/client';
 
 export { sql };
 export type SqliteDB = LibSQLDatabase;
 export type { Table } from './types.js';
 export { createRemoteDatabaseClient, createLocalDatabaseClient } from './db-client.js';
-export { seedLocal } from './queries.js';
+
+export async function seedLocal({
+	// Glob all potential seed files to catch renames and deletions.
+	fileGlob,
+}: {
+	fileGlob: Record<string, { default?: () => Promise<void> }>;
+}) {
+	const seedFilePath = Object.keys(fileGlob)[0];
+	if (!seedFilePath) return;
+	const mod = fileGlob[seedFilePath];
+
+	if (!mod.default) {
+		throw new Error(SEED_DEFAULT_EXPORT_ERROR(seedFilePath));
+	}
+	try {
+		await mod.default();
+	} catch (e) {
+		if (e instanceof LibsqlError) {
+			throw new Error(SEED_ERROR(e.message));
+		}
+		throw e;
+	}
+}
 
 export function hasPrimaryKey(column: DBColumn) {
 	return 'primaryKey' in column.schema && !!column.schema.primaryKey;