Add batch support (#10361)

* deps: drizzle preview

* feat: db.batch and method run handling

* refactor: use db.batch in test fixture

* deps: bump to drizzle 0.29.5

* chore: changeset

* fix: unpin drizzle version

* fix: db execute should uh... execute

.changeset/weak-weeks-pump.md

@@ -0,0 +1,5 @@
+---
+"@astrojs/db": minor
+---
+
+Add support for batch queries with `db.batch()`. This includes an internal bump to Drizzle v0.29.

packages/db/package.json

@@ -62,7 +62,7 @@
     "@libsql/client": "^0.4.3",
     "async-listen": "^3.0.1",
     "deep-diff": "^1.0.2",
-    "drizzle-orm": "^0.28.6",
+    "drizzle-orm": "^0.29.5",
     "kleur": "^4.1.5",
     "nanoid": "^5.0.1",
     "open": "^10.0.3",

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

@@ -1,7 +1,11 @@
 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 {
+	FILE_NOT_FOUND_ERROR,
+	MISSING_EXECUTE_PATH_ERROR,
+	SEED_DEFAULT_EXPORT_ERROR,
+} from '../../../errors.js';
 import {
 	getLocalVirtualModContents,
 	getStudioVirtualModContents,
@@ -47,6 +51,11 @@ export async function cmd({
 		});
 	}
 	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 });
+
+	const mod = await importBundledFile({ code, root: astroConfig.root });
+	if (typeof mod.default !== 'function') {
+		console.error(SEED_DEFAULT_EXPORT_ERROR);
+		process.exit(1);
+	}
+	await mod.default();
 }

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

@@ -15,58 +15,114 @@ export function createLocalDatabaseClient({ dbUrl }: { dbUrl: string }): LibSQLD
 	return db;
 }
 
+const remoteResultSchema = z.object({
+	columns: z.array(z.string()),
+	columnTypes: z.array(z.string()),
+	rows: z.array(z.array(z.unknown())),
+	rowsAffected: z.number(),
+	lastInsertRowid: z.unknown().optional(),
+});
+
 export function createRemoteDatabaseClient(appToken: string, remoteDbURL: string) {
 	const url = new URL('/db/query', remoteDbURL);
 
-	const db = drizzleProxy(async (sql, parameters, method) => {
-		const requestBody: InStatement = { sql, args: parameters };
-		const res = await fetch(url, {
-			method: 'POST',
-			headers: {
-				Authorization: `Bearer ${appToken}`,
-				'Content-Type': 'application/json',
-			},
-			body: JSON.stringify(requestBody),
-		});
-		if (!res.ok) {
-			throw new Error(
-				`Failed to execute query.\nQuery: ${sql}\nFull error: ${res.status} ${await res.text()}}`
-			);
-		}
-
-		const queryResultSchema = z.object({
-			rows: z.array(z.unknown()),
-		});
-		let rows: unknown[];
-		try {
-			const json = await res.json();
-			rows = queryResultSchema.parse(json).rows;
-		} catch (e) {
-			throw new Error(
-				`Failed to execute query.\nQuery: ${sql}\nFull error: Unexpected JSON response. ${
-					e instanceof Error ? e.message : String(e)
-				}`
-			);
-		}
-
-		// Drizzle expects each row as an array of its values
-		const rowValues: unknown[][] = [];
-
-		for (const row of rows) {
-			if (row != null && typeof row === 'object') {
-				rowValues.push(Object.values(row));
+	const db = drizzleProxy(
+		async (sql, parameters, method) => {
+			const requestBody: InStatement = { sql, args: parameters };
+			const res = await fetch(url, {
+				method: 'POST',
+				headers: {
+					Authorization: `Bearer ${appToken}`,
+					'Content-Type': 'application/json',
+				},
+				body: JSON.stringify(requestBody),
+			});
+			if (!res.ok) {
+				throw new Error(
+					`Failed to execute query.\nQuery: ${sql}\nFull error: ${res.status} ${await res.text()}}`
+				);
 			}
+
+			let remoteResult: z.infer<typeof remoteResultSchema>;
+			try {
+				const json = await res.json();
+				remoteResult = remoteResultSchema.parse(json);
+			} catch (e) {
+				throw new Error(
+					`Failed to execute query.\nQuery: ${sql}\nFull error: Unexpected JSON response. ${
+						e instanceof Error ? e.message : String(e)
+					}`
+				);
+			}
+
+			if (method === 'run') return remoteResult;
+
+			// Drizzle expects each row as an array of its values
+			const rowValues: unknown[][] = [];
+
+			for (const row of remoteResult.rows) {
+				if (row != null && typeof row === 'object') {
+					rowValues.push(Object.values(row));
+				}
+			}
+
+			if (method === 'get') {
+				return { rows: rowValues[0] };
+			}
+
+			return { rows: rowValues };
+		},
+		async (queries) => {
+			const stmts: InStatement[] = queries.map(({ sql, params }) => ({ sql, args: params }));
+			const res = await fetch(url, {
+				method: 'POST',
+				headers: {
+					Authorization: `Bearer ${appToken}`,
+					'Content-Type': 'application/json',
+				},
+				body: JSON.stringify(stmts),
+			});
+			if (!res.ok) {
+				throw new Error(
+					`Failed to execute batch queries.\nFull error: ${res.status} ${await res.text()}}`
+				);
+			}
+
+			let remoteResults: z.infer<typeof remoteResultSchema>[];
+			try {
+				const json = await res.json();
+				remoteResults = z.array(remoteResultSchema).parse(json);
+			} catch (e) {
+				throw new Error(
+					`Failed to execute batch queries.\nFull error: Unexpected JSON response. ${
+						e instanceof Error ? e.message : String(e)
+					}`
+				);
+			}
+			let results: any[] = [];
+			for (const [idx, rawResult] of remoteResults.entries()) {
+				if (queries[idx]?.method === 'run') {
+					results.push(rawResult);
+					continue;
+				}
+
+				// Drizzle expects each row as an array of its values
+				const rowValues: unknown[][] = [];
+
+				for (const row of rawResult.rows) {
+					if (row != null && typeof row === 'object') {
+						rowValues.push(Object.values(row));
+					}
+				}
+
+				if (queries[idx]?.method === 'get') {
+					results.push({ rows: rowValues[0] });
+				}
+
+				results.push({ rows: rowValues });
+			}
+			return results;
 		}
-
-		if (method === 'get') {
-			return { rows: rowValues[0] };
-		}
-
-		return { rows: rowValues };
-	});
-
-	(db as any).batch = (_drizzleQueries: Array<Promise<unknown>>) => {
-		throw new Error('db.batch() is not currently supported.');
-	};
+	);
 	return db;
 }

packages/db/test/fixtures/basics/db/seed.ts

@@ -2,20 +2,21 @@ import { asDrizzleTable } from '@astrojs/db/utils';
 import { Themes as ThemesConfig } from './theme';
 import { Author, db } from 'astro:db';
 
+const Themes = asDrizzleTable('Themes', ThemesConfig);
 export default async function () {
-	const Themes = asDrizzleTable('Themes', ThemesConfig);
-
-	await db
-		.insert(Themes)
-		.values([{ name: 'dracula' }, { name: 'monokai', added: new Date() }])
-		.returning({ name: Themes.name });
-	await db
-		.insert(Author)
-		.values([
-			{ name: 'Ben' },
-			{ name: 'Nate' },
-			{ name: 'Erika' },
-			{ name: 'Bjorn' },
-			{ name: 'Sarah' },
-		]);
+	await db.batch([
+		db
+			.insert(Themes)
+			.values([{ name: 'dracula' }, { name: 'monokai', added: new Date() }])
+			.returning({ name: Themes.name }),
+		db
+			.insert(Author)
+			.values([
+				{ name: 'Ben' },
+				{ name: 'Nate' },
+				{ name: 'Erika' },
+				{ name: 'Bjorn' },
+				{ name: 'Sarah' },
+			]),
+	]);
 }

pnpm-lock.yaml

@@ -3824,8 +3824,8 @@ importers:
         specifier: ^1.0.2
         version: 1.0.2
       drizzle-orm:
-        specifier: ^0.28.6
-        version: 0.28.6(@libsql/client@0.4.3)
+        specifier: ^0.29.5
+        version: 0.29.5(@libsql/client@0.4.3)
       kleur:
         specifier: ^4.1.5
         version: 4.1.5
@@ -9952,8 +9952,8 @@ packages:
     engines: {node: '>=10'}
     dev: true
 
-  /drizzle-orm@0.28.6(@libsql/client@0.4.3):
-    resolution: {integrity: sha512-yBe+F9htrlYER7uXgDJUQsTHFoIrI5yMm5A0bg0GiZ/kY5jNXTWoEy4KQtg35cE27sw1VbgzoMWHAgCckUUUww==}
+  /drizzle-orm@0.29.5(@libsql/client@0.4.3):
+    resolution: {integrity: sha512-jS3+uyzTz4P0Y2CICx8FmRQ1eplURPaIMWDn/yq6k4ShRFj9V7vlJk67lSf2kyYPzQ60GkkNGXcJcwrxZ6QCRw==}
     peerDependencies:
       '@aws-sdk/client-rds-data': '>=3'
       '@cloudflare/workers-types': '>=3'
@@ -9963,15 +9963,18 @@ packages:
       '@planetscale/database': '>=1'
       '@types/better-sqlite3': '*'
       '@types/pg': '*'
+      '@types/react': '>=18'
       '@types/sql.js': '*'
       '@vercel/postgres': '*'
       better-sqlite3: '>=7'
       bun-types: '*'
+      expo-sqlite: '>=13.2.0'
       knex: '*'
       kysely: '*'
       mysql2: '>=2'
       pg: '>=8'
       postgres: '>=3'
+      react: '>=18'
       sql.js: '>=1'
       sqlite3: '>=5'
     peerDependenciesMeta:
@@ -9991,6 +9994,8 @@ packages:
         optional: true
       '@types/pg':
         optional: true
+      '@types/react':
+        optional: true
       '@types/sql.js':
         optional: true
       '@vercel/postgres':
@@ -9999,6 +10004,8 @@ packages:
         optional: true
       bun-types:
         optional: true
+      expo-sqlite:
+        optional: true
       knex:
         optional: true
       kysely:
@@ -10009,6 +10016,8 @@ packages:
         optional: true
       postgres:
         optional: true
+      react:
+        optional: true
       sql.js:
         optional: true
       sqlite3: