@effect/sql-d1 4.0.0-beta.98 → 4.0.0-rc.108

package/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 })
@@ -50,13 +52,38 @@ export type TypeId = "~@effect/sql-d1/D1Client"
 /**
  * Cloudflare D1 SQL client service, extending `SqlClient` with its D1 configuration and no `updateValues` support.
  *
- * @category models
+ * @category services
  * @since 4.0.0
  */
 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**
+   *
+   * Each statement uses the query and result name transformations from the
+   * client that created it. Mixing clients can produce differently shaped row
+   * results within the same batch.
+   *
+   * @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
 }
@@ -90,6 +117,78 @@ export interface D1ClientConfig {
   readonly transformQueryNames?: ((str: string) => string) | undefined
 }
 
+type TransformRows = <A extends object>(rows: ReadonlyArray<A>) => ReadonlyArray<A>
+
+type BatchResults<Statements extends ReadonlyArray<Statement.Statement<any>>> = {
+  readonly [K in keyof Statements]: Effect.Success<Statements[K]>
+}
+
+interface StatementWithTransformRows extends Statement.Statement<any> {
+  readonly transformRows: TransformRows | undefined
+}
+
+const makeBatch = (options: {
+  readonly db: D1Database
+  readonly prepareCache: Cache.Cache<string, D1PreparedStatement, SqlError>
+  readonly spanAttributes: ReadonlyArray<readonly [string, unknown]>
+  readonly getClient: () => D1Client
+}): D1Client["batch"] =>
+<const Statements extends ReadonlyArray<Statement.Statement<any>>>(
+  statements: Statements
+) => {
+  if (statements.length === 0) {
+    return Effect.succeed([] as unknown as BatchResults<Statements>)
+  }
+  return Effect.useSpan(
+    "sql.execute",
+    { kind: "client" },
+    (span) =>
+      Effect.withFiber(Effect.fnUntraced(function*(fiber) {
+        const transformer = fiber.getRef(Statement.CurrentTransformer)
+        const prepared: Array<D1PreparedStatement> = []
+        const transforms: Array<TransformRows | undefined> = []
+        const queryTexts: Array<string> = []
+
+        for (const original of statements) {
+          const statement = transformer === undefined
+            ? original
+            : yield* transformer(original, options.getClient(), fiber, span)
+          const [sql, params] = statement.compile()
+          queryTexts.push(sql)
+          transforms.push((statement as StatementWithTransformRows).transformRows)
+          prepared.push((yield* Cache.get(options.prepareCache, sql)).bind(...params))
+        }
+
+        for (const [key, value] of options.spanAttributes) {
+          span.attribute(key, value)
+        }
+        span.attribute(ATTR_DB_OPERATION_NAME, "batch")
+        span.attribute(ATTR_DB_QUERY_TEXT, queryTexts.join("; "))
+
+        // D1 batches execute on the binding directly and intentionally cannot participate in SqlClient transactions.
+        const responses = yield* Effect.tryPromise({
+          try: () =>
+            options.db.batch<Record<string, unknown>>(prepared).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") })
+        })
+
+        const results = responses.map((response, index) => {
+          const rows = response.results || []
+          const transformRows = transforms[index]
+          return transformRows ? transformRows(rows) : rows
+        })
+        return results as BatchResults<Statements>
+      }))
+  )
+}
+
 /**
  * Creates a scoped Cloudflare D1 SQL client. Prepared statements are cached, while transactions and streaming queries are not supported by this driver.
  *
@@ -104,6 +203,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 +285,7 @@ export const make = (
           catch: (cause) => new SqlError({ reason: classifyError(cause, "Failed to execute statement", "execute") })
         })
 
-      return identity<Connection>({
+      const connection = identity<Connection>({
         execute(sql, params, transformRows) {
           return transformRows
             ? Effect.map(runCached(sql, params), transformRows)
@@ -206,28 +309,60 @@ export const make = (
           return Stream.die("executeStream not implemented")
         }
       })
+      return { connection, prepareCache } as const
     })
 
-    const connection = yield* makeConnection
+    const { connection, prepareCache } = 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({
+          db: options.db,
+          prepareCache,
+          spanAttributes,
+          getClient: () => client
+        })
       }
     )
+
+    if (options.transformQueryNames !== undefined || transformRows !== undefined) {
+      const clientWithoutTransformsBase = yield* Client.make({
+        acquirer: Effect.succeed(connection),
+        compiler: compiler.withoutTransform,
+        transactionAcquirer,
+        spanAttributes,
+        transformRows: undefined
+      })
+      let clientWithoutTransforms!: D1Client
+      clientWithoutTransforms = Object.assign(clientWithoutTransformsBase as D1Client, {
+        [TypeId]: TypeId as TypeId,
+        config: options,
+        batch: makeBatch({
+          db: options.db,
+          prepareCache,
+          spanAttributes,
+          getClient: () => clientWithoutTransforms
+        }),
+        withoutTransforms: () => clientWithoutTransforms
+      })
+      Object.assign(client, {
+        withoutTransforms: () => clientWithoutTransforms
+      })
+    }
+
+    return client
   })
 
 /**