diff --git a/.changeset/d1-batch-statements.md b/.changeset/d1-batch-statements.md new file mode 100644 index 00000000000..d1a436bb913 --- /dev/null +++ b/.changeset/d1-batch-statements.md @@ -0,0 +1,5 @@ +--- +"@effect/sql-d1": patch +--- + +Add `D1Client.batch` for executing a collection of SQL statements as a single atomic D1 batch. diff --git a/packages/sql/d1/src/D1Client.ts b/packages/sql/d1/src/D1Client.ts index 50078f5bc6a..e0b7334ae62 100644 --- a/packages/sql/d1/src/D1Client.ts +++ b/packages/sql/d1/src/D1Client.ts @@ -27,6 +27,8 @@ import { SqlError, UnknownError } from "effect/unstable/sql/SqlError" import * as Statement from "effect/unstable/sql/Statement" const ATTR_DB_SYSTEM_NAME = "db.system.name" +const ATTR_DB_OPERATION_NAME = "db.operation.name" +const ATTR_DB_QUERY_TEXT = "db.query.text" const classifyError = (cause: unknown, message: string, operation: string) => new UnknownError({ cause, message, operation }) @@ -57,6 +59,30 @@ export interface D1Client extends Client.SqlClient { readonly [TypeId]: TypeId readonly config: D1ClientConfig + /** + * Executes SQL statements as a single atomic D1 batch and returns their row results in order. + * + * **When to use** + * + * Use when you have a fixed collection of statements that should run in one + * request and roll back together if any statement fails. + * + * **Gotchas** + * + * Statements should be created by this client so that query and result name + * transformations remain consistent. + * + * @since 4.0.0 + */ + readonly batch: >>( + statements: Statements + ) => Effect.Effect< + { + readonly [K in keyof Statements]: Effect.Success + }, + SqlError + > + /** Not supported in d1 */ readonly updateValues: never } @@ -104,6 +130,10 @@ export const make = ( const transformRows = options.transformResultNames ? Statement.defaultTransforms(options.transformResultNames).array : undefined + const spanAttributes: Array = [ + ...(options.spanAttributes ? Object.entries(options.spanAttributes) : []), + [ATTR_DB_SYSTEM_NAME, "sqlite"] + ] const makeConnection = Effect.gen(function*() { const db = options.db @@ -182,7 +212,68 @@ export const make = ( catch: (cause) => new SqlError({ reason: classifyError(cause, "Failed to execute statement", "execute") }) }) - return identity({ + const makeBatch = ( + transformRows: ((rows: ReadonlyArray) => ReadonlyArray) | undefined, + getClient: () => D1Client + ) => + >>( + statements: Statements + ): Effect.Effect< + { + readonly [K in keyof Statements]: Effect.Success + }, + SqlError + > => { + if (statements.length === 0) { + return Effect.succeed([]) as any + } + return Effect.useSpan( + "sql.execute", + { kind: "client" }, + (span) => + Effect.gen(function*() { + const compiled = yield* Effect.forEach(statements, (statement) => + Effect.withFiber((fiber) => { + const transformer = fiber.getRef(Statement.CurrentTransformer) + if (transformer === undefined) { + return Effect.succeed(statement.compile()) + } + return Effect.map( + transformer(statement, getClient(), fiber, span), + (statement) => statement.compile() + ) + })) + for (const [key, value] of spanAttributes) { + span.attribute(key, value) + } + span.attribute(ATTR_DB_OPERATION_NAME, "batch") + span.attribute(ATTR_DB_QUERY_TEXT, compiled.map(([sql]) => sql).join("; ")) + + const prepared = yield* Effect.forEach(compiled, ([sql]) => Cache.get(prepareCache, sql)) + const responses = yield* Effect.tryPromise({ + try: () => + db.batch(prepared.map((statement, index) => statement.bind(...compiled[index][1]))).then( + (responses) => { + for (const response of responses) { + if (response.error) { + throw response.error + } + } + return responses + } + ), + catch: (cause) => new SqlError({ reason: classifyError(cause, "Failed to execute batch", "execute") }) + }) + + return responses.map((response) => { + const rows = response.results || [] + return transformRows ? transformRows(rows as ReadonlyArray>) : rows + }) as any + }) + ) as any + } + + const connection = identity({ execute(sql, params, transformRows) { return transformRows ? Effect.map(runCached(sql, params), transformRows) @@ -206,28 +297,45 @@ export const make = ( return Stream.die("executeStream not implemented") } }) + return { connection, makeBatch } as const }) - const connection = yield* makeConnection + const { connection, makeBatch } = yield* makeConnection const acquirer = Effect.succeed(connection) const transactionAcquirer = Effect.die("transactions are not supported in D1") - return Object.assign( + let client!: D1Client + client = Object.assign( (yield* Client.make({ acquirer, compiler, transactionAcquirer, - spanAttributes: [ - ...(options.spanAttributes ? Object.entries(options.spanAttributes) : []), - [ATTR_DB_SYSTEM_NAME, "sqlite"] - ], + spanAttributes, transformRows })) as D1Client, { [TypeId]: TypeId as TypeId, - config: options + config: options, + batch: makeBatch(transformRows, () => client) } ) + + if (transformRows !== undefined) { + const withoutTransforms = client.withoutTransforms + Object.assign(client, { + withoutTransforms: () => { + let clientWithoutTransforms!: D1Client + clientWithoutTransforms = Object.assign(withoutTransforms.call(client), { + [TypeId]: TypeId as TypeId, + config: options, + batch: makeBatch(undefined, () => clientWithoutTransforms) + }) + return clientWithoutTransforms + } + }) + } + + return client }) /** diff --git a/packages/sql/d1/test/Client.test.ts b/packages/sql/d1/test/Client.test.ts index a8678b24c9f..02da632f55b 100644 --- a/packages/sql/d1/test/Client.test.ts +++ b/packages/sql/d1/test/Client.test.ts @@ -2,6 +2,7 @@ import { D1Client } from "@effect/sql-d1" import { assert, describe, it } from "@effect/vitest" import { Cause, Effect } from "effect" import * as Reactivity from "effect/unstable/reactivity/Reactivity" +import { Statement } from "effect/unstable/sql" import { D1Miniflare } from "./utils.ts" describe("Client", () => { @@ -49,6 +50,100 @@ describe("Client", () => { assert.deepStrictEqual(rows, [{ id: 1, name: "hello" }]) }).pipe(Effect.provide(D1Miniflare.layerClient))) + it.effect("should execute statements in a batch", () => + Effect.gen(function*() { + const sql = yield* D1Client.D1Client + yield* sql`CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)` + + const results: readonly [ + ReadonlyArray<{ id: number; name: string }>, + ReadonlyArray<{ id: number; name: string }>, + ReadonlyArray<{ count: number }> + ] = yield* sql.batch( + [ + sql<{ id: number; name: string }>`INSERT INTO test (name) VALUES (${"hello"}) RETURNING *`, + sql<{ id: number; name: string }>`INSERT INTO test (name) VALUES (${"world"}) RETURNING *`, + sql<{ count: number }>`SELECT COUNT(*) AS count FROM test` + ] as const + ) + + assert.deepStrictEqual(results, [ + [{ id: 1, name: "hello" }], + [{ id: 2, name: "world" }], + [{ count: 2 }] + ]) + }).pipe(Effect.provide(D1Miniflare.layerClient))) + + it.effect("should apply result transforms to batch results", () => + Effect.gen(function*() { + const sql = yield* D1Client.D1Client + yield* sql`CREATE TABLE test (first_name TEXT, "firstName" TEXT)` + yield* sql`INSERT INTO test (first_name, "firstName") VALUES ('John', 'Jane')` + + yield* Effect.gen(function*() { + const transformed = yield* D1Client.make({ + db: sql.config.db, + transformQueryNames: (name) => name === "firstName" ? "first_name" : name, + transformResultNames: (name) => name === "first_name" ? "firstName" : name + }) + const [rows] = yield* transformed.batch([ + transformed<{ firstName: string }>`SELECT ${transformed("firstName")} FROM test` + ]) + assert.deepStrictEqual(rows, [{ firstName: "John" }]) + + const withoutTransforms = transformed.withoutTransforms() + const [rawRows] = yield* withoutTransforms.batch([ + withoutTransforms<{ firstName: string }>`SELECT ${withoutTransforms("firstName")} FROM test` + ]) + assert.deepStrictEqual(rawRows, [{ firstName: "Jane" }]) + }).pipe( + Effect.scoped, + Effect.provide(Reactivity.layer) + ) + }).pipe(Effect.provide(D1Miniflare.layerClient))) + + it.effect("should roll back a failed batch", () => + Effect.gen(function*() { + const sql = yield* D1Client.D1Client + yield* sql`CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT UNIQUE)` + + const error = yield* sql.batch([ + sql`INSERT INTO test (name) VALUES (${"duplicate"})`, + sql`INSERT INTO test (name) VALUES (${"duplicate"})` + ]).pipe(Effect.flip) + + assert.strictEqual(error.reason._tag, "UnknownError") + assert.strictEqual(error.reason.operation, "execute") + const rows = yield* sql`SELECT * FROM test` + assert.deepStrictEqual(rows, []) + }).pipe(Effect.provide(D1Miniflare.layerClient))) + + it.effect("should support an empty batch", () => + Effect.gen(function*() { + const sql = yield* D1Client.D1Client + const results: readonly [] = yield* sql.batch([]) + assert.deepStrictEqual(results, []) + }).pipe(Effect.provide(D1Miniflare.layerClient))) + + it.effect("should apply statement transformers in a batch", () => + Effect.gen(function*() { + const sql = yield* D1Client.D1Client + const queries: Array = [] + + const results = yield* sql.batch([ + sql`SELECT ${1} AS value`, + sql`SELECT ${2} AS value` + ]).pipe( + Effect.provideService(Statement.CurrentTransformer, (statement, sql) => + Effect.sync(() => { + queries.push(statement.compile()[0]) + }).pipe(Effect.as(sql`SELECT ${3} AS value`))) + ) + + assert.deepStrictEqual(queries, ["SELECT ? AS value", "SELECT ? AS value"]) + assert.deepStrictEqual(results, [[{ value: 3 }], [{ value: 3 }]]) + }).pipe(Effect.provide(D1Miniflare.layerClient))) + it.effect("should defect on transactions", () => Effect.gen(function*() { const sql = yield* D1Client.D1Client