Skip to content

Commit

Permalink
Add batch support (#10361)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
bholmesdev authored Mar 7, 2024
1 parent 52fba64 commit 988aad6
Showing 6 changed files with 146 additions and 66 deletions.
5 changes: 5 additions & 0 deletions .changeset/weak-weeks-pump.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"@astrojs/db": minor
---

Add support for batch queries with `db.batch()`. This includes an internal bump to Drizzle v0.29.
2 changes: 1 addition & 1 deletion packages/db/package.json
Original file line number Diff line number Diff line change
@@ -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",
15 changes: 12 additions & 3 deletions packages/db/src/core/cli/commands/execute/index.ts
Original file line number Diff line number Diff line change
@@ -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();
}
142 changes: 99 additions & 43 deletions packages/db/src/runtime/db-client.ts
Original file line number Diff line number Diff line change
@@ -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 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)
}`
);
}
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)
}`
);
}

// Drizzle expects each row as an array of its values
const rowValues: unknown[][] = [];
if (method === 'run') return remoteResult;

for (const row of rows) {
if (row != null && typeof row === 'object') {
rowValues.push(Object.values(row));
// 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] };
}
if (method === 'get') {
return { rows: rowValues[0] };
}

return { rows: rowValues };
});
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[][] = [];

(db as any).batch = (_drizzleQueries: Array<Promise<unknown>>) => {
throw new Error('db.batch() is not currently supported.');
};
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;
}
);
return db;
}
31 changes: 16 additions & 15 deletions packages/db/test/fixtures/basics/db/seed.ts
Original file line number Diff line number Diff line change
@@ -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' },
]),
]);
}
17 changes: 13 additions & 4 deletions pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 988aad6

Please sign in to comment.