Add D1 batch support

.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.

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: <const Statements extends ReadonlyArray<Statement.Statement<any>>>(
+    statements: Statements
+  ) => Effect.Effect<
+    {
+      readonly [K in keyof Statements]: Effect.Success<Statements[K]>
+    },
+    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<readonly [string, unknown]> = [
+      ...(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<Connection>({
+      const makeBatch = (
+        transformRows: (<A extends object>(rows: ReadonlyArray<A>) => ReadonlyArray<A>) | undefined,
+        getClient: () => D1Client
+      ) =>
+      <const Statements extends ReadonlyArray<Statement.Statement<any>>>(
+        statements: Statements
+      ): Effect.Effect<
+        {
+          readonly [K in keyof Statements]: Effect.Success<Statements[K]>
+        },
+        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<Record<string, unknown>>) : rows
+              }) as any
+            })
+        ) as any
+      }
+
+      const connection = identity<Connection>({
         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
   })
 
 /**

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<string> = []
+
+      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