@uniflowed/sql 0.35.0 → 0.36.0

package/adapters/pg.js

@@ -11,8 +11,11 @@
 // with the identity for that query alone, so the pool's own configuration is
 // left as the application set it, and generated code still sees the server's
 // text. A transaction checks a client out of the pool for its duration.
+// `:copyfrom` is a `COPY FROM STDIN`: node-postgres runs it when the query
+// object supplies `submit` and `handleCopyInResponse`, which is the protocol
+// hook it already has and does not implement for a string query.
 
-import type { Connection, QueryResult, Queryable, Run } from "../index.js";
+import type { Connection, CopyIn, QueryResult, Queryable, Run } from "../index.js";
 import { singleConnection, transactionOn } from "../index.js";
 
 type PgResult = { readonly rows: $ReadOnlyArray<mixed>, readonly rowCount: number | null, ... };
@@ -24,9 +27,32 @@ type PgQuery = {|
   readonly types  : {| readonly getTypeParser: () => (value: string) => string |},
 |};
 
+/** The connection methods a `COPY FROM STDIN` uses. `pg` passes its own. */
+type PgCopyConnection = {
+  query(text: string): mixed,
+  sendCopyFromChunk(chunk: Uint8Array): mixed,
+  endCopyFrom(): mixed,
+};
+
+/** A query object `pg` runs when `submit` is a function. */
+type PgCopyRequest = {
+  text    : string,
+  callback: ?(?mixed, ?PgResult) => mixed,
+  submit(connection: PgCopyConnection): null,
+  handleCopyInResponse(connection: PgCopyConnection): void,
+  handleCommandComplete(message: { text?: string, ... }): void,
+  handleReadyForQuery(): void,
+  handleError(error: mixed): void,
+  handleCopyData(): void,
+  handleRowDescription(): void,
+  handleDataRow(): void,
+  handleEmptyQuery(): void,
+  handlePortalSuspended(): void,
+};
+
 /** A connected `pg.Client` or a client checked out of a pool. */
 export interface PgClient {
-  query(config: PgQuery): Promise<PgResult>;
+  query(config: PgQuery | PgCopyRequest): Promise<PgResult> | void;
 }
 
 /** A client checked out of a `pg.Pool`. */
@@ -36,15 +62,99 @@ export interface PgPoolClient extends PgClient {
 
 /** A `pg.Pool`. */
 export interface PgPool {
-  query(config: PgQuery): Promise<PgResult>;
+  query(config: PgQuery | PgCopyRequest): Promise<PgResult>;
   connect(): Promise<PgPoolClient>;
 }
 
+const COPY_ROWS = /COPY (\d+)/;
+
+/**
+ * One `COPY FROM STDIN` on `runner`.
+ *
+ * The callback is left unset when `runner` is a pool: the pool sets it, and
+ * setting it here would skip the pool's own, which is what releases the
+ * client. A client sets it in the wrapper below, because `client.query` only
+ * returns a promise when the object does not already have one and nobody
+ * passed a callback.
+ */
+function copyRequest(statement: string, payload: string): PgCopyRequest {
+  let rowCount = 0;
+  let settled = false;
+  const request: PgCopyRequest = {
+    text    : statement,
+    callback: undefined,
+    submit(connection) {
+      connection.query(statement);
+      return null;
+    },
+    handleCopyInResponse(connection) {
+      // pg-protocol copies the chunk with `Buffer#copy`. A `Uint8Array` has
+      // no `copy`, and the throw leaves the backend waiting for `CopyData`.
+      connection.sendCopyFromChunk(Buffer.from(payload, "utf8"));
+      connection.endCopyFrom();
+    },
+    handleCommandComplete(message) {
+      const text = typeof message.text === "string" ? message.text : "";
+      const match = COPY_ROWS.exec(text);
+      if (match?.[1] !== undefined) {
+        rowCount = Number(match[1]);
+      }
+    },
+    handleReadyForQuery() {
+      if (settled) {
+        return;
+      }
+      settled = true;
+      request.callback?.(null, { rows: [], rowCount });
+    },
+    handleError(error) {
+      if (settled) {
+        return;
+      }
+      settled = true;
+      request.callback?.(error, undefined);
+    },
+    handleCopyData() {},
+    handleRowDescription() {},
+    handleDataRow() {},
+    handleEmptyQuery() {},
+    handlePortalSuspended() {},
+  };
+  return request;
+}
+
+function copyOn(runner: PgClient): CopyIn {
+  return (statement, payload) =>
+    new Promise((resolve, reject) => {
+      const request = copyRequest(statement, payload);
+      request.callback = (error, result) => {
+        if (error != null) {
+          reject(error instanceof Error ? error : new Error(String(error)));
+          return;
+        }
+        resolve(result?.rowCount ?? 0);
+      };
+      runner.query(request);
+    });
+}
+
+function copyOnPool(pool: PgPool): CopyIn {
+  return async (statement, payload) => {
+    const result = await pool.query(copyRequest(statement, payload));
+    return result.rowCount ?? 0;
+  };
+}
+
 const RAW = { getTypeParser: () => (value: string) => value };
 
 function runOn(client: PgClient): Run {
   return async (text, params): Promise<QueryResult> => {
     const result = await client.query({ text, values: params, rowMode: "array", types: RAW });
+    // A `COPY` request returns void. This call is a normal query, which
+    // returns a result; the union is what `query` is declared as.
+    if (result == null) {
+      throw new Error("pg query returned no result");
+    }
     const rows: Array<$ReadOnlyArray<mixed>> = [];
     for (const row of result.rows) {
       if (Array.isArray(row)) {
@@ -56,7 +166,7 @@ function runOn(client: PgClient): Run {
 }
 
 function connection(client: PgClient): Connection {
-  return { engine: "postgresql", maxParams: 65535, run: runOn(client) };
+  return { engine: "postgresql", maxParams: 65535, run: runOn(client), copy: copyOn(client) };
 }
 
 /** A [`Queryable`] over one connected `pg.Client`, which it serialises. */
@@ -70,6 +180,7 @@ export function fromPgPool(pool: PgPool): Queryable {
     engine   : "postgresql",
     maxParams: 65535,
     query    : runOn(pool),
+    copy     : copyOnPool(pool),
     transaction: async <T>(body: (tx: Queryable) => Promise<T>): Promise<T> => {
       const client = await pool.connect();
       let broken: mixed = undefined;

package/adapters/pglite.js

@@ -10,7 +10,8 @@
 // PGlite parses and serialises values itself, by type. This adapter turns both
 // off — an identity parser and serialiser for every type it knows — so rows
 // arrive as the server's text and parameters leave as the codecs wrote them,
-// exactly as with every other PostgreSQL adapter.
+// exactly as with every other PostgreSQL adapter. `:copyfrom` uses PGlite's
+// `/dev/blob` device, which is how its `COPY FROM` reads a payload.
 
 import type { Queryable, SqlParam } from "../index.js";
 import { singleConnection } from "../index.js";
@@ -24,6 +25,7 @@ export interface PGliteDatabase {
       readonly rowMode    : "array",
       readonly parsers    : { readonly [string]: (value: string) => mixed },
       readonly serializers: { readonly [string]: (value: mixed) => mixed },
+      readonly blob?      : Blob,
     |},
   ): Promise<{ readonly rows: $ReadOnlyArray<mixed>, readonly affectedRows?: number, ... }>;
   readonly parsers    : { readonly [string]: mixed };
@@ -52,6 +54,27 @@ export function fromPGlite(database: PGliteDatabase): Queryable {
   return singleConnection({
     engine   : "postgresql",
     maxParams: 65535,
+    copy: async (statement, payload) => {
+      // `COPY … FROM STDIN` is the wire form. PGlite reads the payload from
+      // the virtual file named here instead of a frontend COPY stream.
+      const text = statement.replace(" FROM STDIN", " FROM '/dev/blob'");
+      const result = await database.query(text, [], {
+        rowMode    : "array",
+        parsers    : {},
+        serializers: {},
+        blob       : new Blob([payload]),
+      });
+      if (typeof result.affectedRows === "number") {
+        return result.affectedRows;
+      }
+      let rows = 0;
+      for (let i = 0; i < payload.length; i += 1) {
+        if (payload[i] === "\n") {
+          rows += 1;
+        }
+      }
+      return rows;
+    },
     run: async (text, params) => {
       const size = Object.keys(database.parsers).length + Object.keys(database.serializers).length;
       if (size !== known) {

package/adapters/postgres.js

@@ -18,14 +18,30 @@
 // `bytea`, `x === true ? "t" : "f"` for `bool` — each of which would corrupt a
 // value the codecs already wrote as text. Declaring every parameter as
 // `unknown` (OID 705) leaves the server to infer the type, as an undeclared
-// one would, and leaves postgres.js nothing to serialise.
+// one would, and leaves postgres.js nothing to serialise. `:copyfrom` writes
+// the text-format payload to the stream `unsafe(statement, []).writable()`
+// opens for `COPY FROM STDIN`.
 
 import { SqlError } from "../index.js";
-import type { Queryable, SqlParam } from "../index.js";
+import type { CopyIn, Queryable, SqlParam } from "../index.js";
 
 type RawRows = Promise<$ReadOnlyArray<mixed> & { readonly count?: number | null, ... }>;
 
-type Unsafe = { raw(): RawRows, ... };
+/** The Node writable postgres.js resolves `COPY FROM STDIN` to. */
+type CopyWritable = {
+  write(chunk: string, callback: (error?: mixed) => mixed): mixed,
+  end(callback: (error?: mixed) => mixed): mixed,
+  once(event: "error", listener: (error: mixed) => mixed): mixed,
+  ...
+};
+
+type Unsafe = {
+  raw(): RawRows,
+  writable(): Promise<CopyWritable>,
+  /** Rejects the query. postgres.js calls it for a server error during `COPY`. */
+  reject: (error: mixed) => mixed,
+  ...
+};
 
 /** What postgres.js's `Sql` and its transaction handles have in common. */
 export interface PostgresQueries {
@@ -56,6 +72,78 @@ function text(value: mixed): mixed {
 /** `unknown`: a parameter whose type the server decides. */
 const UNKNOWN = 705;
 
+function rowCount(payload: string): number {
+  let rows = 0;
+  for (let i = 0; i < payload.length; i += 1) {
+    if (payload[i] === "\n") {
+      rows += 1;
+    }
+  }
+  return rows;
+}
+
+/** Write `payload` and wait until the server finishes the `COPY`. */
+function writeCopy(writable: CopyWritable, payload: string): Promise<void> {
+  return new Promise((resolve, reject) => {
+    let settled = false;
+    const fail = (error: mixed) => {
+      if (settled) {
+        return;
+      }
+      settled = true;
+      reject(error instanceof Error ? error : new Error(String(error)));
+    };
+    const succeed = () => {
+      if (settled) {
+        return;
+      }
+      settled = true;
+      resolve();
+    };
+    writable.once("error", fail);
+    writable.write(payload, (error) => {
+      if (error != null) {
+        fail(error);
+        return;
+      }
+      writable.end((endError) => {
+        if (endError != null) {
+          fail(endError);
+        } else {
+          succeed();
+        }
+      });
+    });
+  });
+}
+
+function copyOn(sql: PostgresQueries): CopyIn {
+  return async (statement, payload) => {
+    const query = sql.unsafe(statement, []);
+    // CopyInResponse resolves the query with the writable, and sending
+    // CopyDone clears that stream before the server answers. A failure after
+    // the data is in (a foreign key) calls `query.reject` and never reaches
+    // the writable, so its end callback would wait forever.
+    let fail: (error: mixed) => void = () => {};
+    const failed = new Promise((resolve, reject) => {
+      void resolve;
+      fail = reject;
+    });
+    failed.catch(() => {});
+    const rejectQuery = query.reject.bind(query);
+    query.reject = (error) => {
+      fail(error);
+      return rejectQuery(error);
+    };
+    const writable = await query.writable();
+    await Promise.race([writeCopy(writable, payload), failed]);
+    // CommandComplete carries `COPY n`, and postgres.js keeps that count on
+    // its internal result. A `COPY` without `WHERE` accepts every row or
+    // fails, so the number of lines written is the number the server copied.
+    return rowCount(payload);
+  };
+}
+
 function run(sql: PostgresQueries): Queryable["query"] {
   return async (statement, params) => {
     const typed = params.map((param) => (param === null ? null : sql.typed(param, UNKNOWN)));
@@ -87,6 +175,8 @@ async function scoped<T>(tx: PostgresTransaction, body: (tx: Queryable) => Promi
     maxParams: 65535,
     query: (statement, params, mode) =>
       open ? query(statement, params, mode) : Promise.reject(new SqlError({ kind: "closed" })),
+    copy: (statement, payload) =>
+      open ? copyOn(tx)(statement, payload) : Promise.reject(new SqlError({ kind: "closed" })),
     transaction: <U>(inner: (tx: Queryable) => Promise<U>): Promise<U> =>
       open
         ? tx.savepoint((savepoint) => scoped(savepoint, inner))
@@ -105,6 +195,7 @@ export function fromPostgres(sql: PostgresSql): Queryable {
     engine   : "postgresql",
     maxParams: 65535,
     query    : run(sql),
+    copy     : copyOn(sql),
     transaction: <T>(body: (tx: Queryable) => Promise<T>): Promise<T> =>
       sql.begin((tx) => scoped(tx, body)),
   };

package/index.js

@@ -16,8 +16,9 @@
 //     type. So an `int8` is a `bigint` under `pg` and under `postgres` alike,
 //     rather than whatever each driver's own parser decided.
 //   * **The parts of sqlc's semantics a driver does not have**: `sqlc.slice`
-//     expansion, `:copyfrom` as chunked multi-row inserts, `:batch*` in one
-//     transaction, and nested transactions as savepoints.
+//     expansion, `:copyfrom` (PostgreSQL `COPY` when the adapter can send it,
+//     otherwise chunked multi-row inserts), `:batch*` in one transaction, and
+//     nested transactions as savepoints.
 //   * **A synchronous path** for drivers that run a statement on the calling
 //     thread (`node:sqlite`, `better-sqlite3`, `bun:sqlite`). Generated code
 //     with `sync: true` calls it, so a transaction can stay inside `BEGIN`
@@ -65,12 +66,24 @@ export type QueryMode = "rows" | "exec";
  * `transaction` is absent where the platform has no interactive transactions
  * (Cloudflare D1). Inside a transaction it opens a savepoint.
  */
+/**
+ * Send one `COPY … FROM STDIN` payload and resolve to the number of rows the
+ * server accepted.
+ *
+ * Present on a PostgreSQL adapter that can speak `COPY`. Absent everywhere
+ * else, where `:copyfrom` stays a chunked `INSERT`. The payload is PostgreSQL
+ * text format: tab-separated fields, one row per line.
+ */
+export type CopyIn = (statement: string, payload: string) => Promise<number>;
+
 export type Queryable = {
   readonly engine: Engine,
   /** The most bound parameters one statement may carry; `:copyfrom` chunks under it. */
   readonly maxParams: number,
   query(text: string, params: $ReadOnlyArray<SqlParam>, mode: QueryMode): Promise<QueryResult>,
   readonly transaction?: <T>(body: (tx: Queryable) => Promise<T>) => Promise<T>,
+  /** PostgreSQL `COPY`, when this adapter can send one. */
+  readonly copy?: CopyIn,
   ...
 };
 
@@ -494,6 +507,96 @@ export type CopyPlan = {|
 /** The most rows one statement carries when a row has no parameters at all. */
 const COPY_ROWS_WITHOUT_PARAMS = 1000;
 
+/**
+ * `COPY <table> (<columns>) FROM STDIN` for an `INSERT … VALUES` plan, or
+ * `null` when the plan is not that shape.
+ *
+ * sqlc only accepts `INSERT INTO t (…) VALUES (…)` for `:copyfrom`, and the
+ * generator puts the table and the column list in `head`. Anything else keeps
+ * the chunked `INSERT`, which is the path every engine can run.
+ */
+function copyStatement(plan: CopyPlan): string | null {
+  const prefix = "INSERT INTO ";
+  const suffix = " VALUES ";
+  if (!plan.head.startsWith(prefix) || !plan.head.endsWith(suffix)) {
+    return null;
+  }
+  const target = plan.head.slice(prefix.length, plan.head.length - suffix.length);
+  return `COPY ${target} FROM STDIN WITH (FORMAT text)`;
+}
+
+/** One field of PostgreSQL's text `COPY` format. */
+function copyField(value: SqlParam): string {
+  if (value === null) {
+    return "\\N";
+  }
+  return escapeCopyText(copyText(value));
+}
+
+function copyText(value: string | number | bigint | Uint8Array): string {
+  if (typeof value === "string") {
+    return value;
+  }
+  if (typeof value === "number" || typeof value === "bigint") {
+    return String(value);
+  }
+  let hex = "\\x";
+  for (let i = 0; i < value.length; i += 1) {
+    hex += value[i].toString(16).padStart(2, "0");
+  }
+  return hex;
+}
+
+/** `\`, newline, carriage return and tab, the four escapes text `COPY` requires. */
+function escapeCopyText(text: string): string {
+  let out = "";
+  for (let i = 0; i < text.length; i += 1) {
+    const character = text[i];
+    if (character === "\\") {
+      out += "\\\\";
+    } else if (character === "\n") {
+      out += "\\n";
+    } else if (character === "\r") {
+      out += "\\r";
+    } else if (character === "\t") {
+      out += "\\t";
+    } else {
+      out += character;
+    }
+  }
+  return out;
+}
+
+/**
+ * Every row of a `:copyfrom` as one text-format payload.
+ *
+ * A short row throws before anything is sent, the same way the `INSERT` path
+ * does. A null field is `\N`. The payload ends each row with a newline, which
+ * is the row terminator `COPY` counts.
+ */
+function copyPayload(
+  plan: CopyPlan,
+  rows: $ReadOnlyArray<$ReadOnlyArray<SqlParam>>,
+  name: string,
+): string {
+  const width = plan.refs.length;
+  let payload = "";
+  for (const row of rows) {
+    for (let i = 0; i < width; i += 1) {
+      const index = plan.refs[i];
+      if (index >= row.length) {
+        throw new SqlError({ kind: "row", index, length: row.length }, name);
+      }
+      if (i > 0) {
+        payload += "\t";
+      }
+      payload += copyField(row[index]);
+    }
+    payload += "\n";
+  }
+  return payload;
+}
+
 function copyChunk(
   engine: Engine,
   plan  : CopyPlan,
@@ -528,12 +631,15 @@ function copyWidth(db: { readonly maxParams: number, ... }, plan: CopyPlan, name
 }
 
 /**
- * `:copyfrom` — insert every row, as few statements as `maxParams` allows.
+ * `:copyfrom` — insert every row.
  *
- * Inside one transaction when the `Queryable` can open one, so a failure in
- * the third chunk does not leave the first two behind: pgx's `CopyFrom` is one
- * `COPY`, and a caller of the generated function should not have to know it is
- * now several statements. Resolves to the number of rows inserted.
+ * A PostgreSQL adapter that implements `copy` receives one `COPY … FROM STDIN`
+ * in text format, which is one statement and so one success or one failure.
+ * Every other adapter, and a plan that is not an `INSERT … VALUES`, stays on
+ * chunked multi-row `INSERT`s. Those run inside one transaction when the
+ * `Queryable` can open one and the rows do not fit in a single statement, so
+ * a failure in the third chunk does not leave the first two behind. Resolves
+ * to the number of rows inserted.
  */
 export async function copyFrom(
   db  : Queryable,
@@ -544,6 +650,13 @@ export async function copyFrom(
   if (rows.length === 0) {
     return 0;
   }
+  const copy = db.copy;
+  if (db.engine === "postgresql" && typeof copy === "function") {
+    const statement = copyStatement(plan);
+    if (statement !== null) {
+      return copy(statement, copyPayload(plan, rows, name));
+    }
+  }
   const perStatement = copyWidth(db, plan, name);
   const insert = async (q: Queryable): Promise<number> => {
     let total = 0;
@@ -667,6 +780,8 @@ export type Connection = {|
   readonly runSync?: RunSync,
   /** The statement that opens a transaction; `BEGIN` unless the adapter knows better. */
   readonly begin?: string,
+  /** PostgreSQL `COPY` on this connection, when the driver can send one. */
+  readonly copy?: CopyIn,
 |};
 
 /**
@@ -701,10 +816,17 @@ async function scoped<T>(
     }
     return connection.run(text, params, mode);
   };
+  const copy = connection.copy;
+  const guardedCopy: CopyIn | void =
+    copy === undefined
+      ? undefined
+      : (statement, payload) =>
+          open ? copy(statement, payload) : Promise.reject(new SqlError({ kind: "closed" }));
   const tx: Queryable = {
     engine   : connection.engine,
     maxParams: connection.maxParams,
     query    : guarded,
+    ...(guardedCopy === undefined ? null : { copy: guardedCopy }),
     transaction: async <U>(inner: (tx: Queryable) => Promise<U>): Promise<U> => {
       const savepoint = `uf_sp_${depth + 1}`;
       await guarded(`SAVEPOINT ${savepoint}`, [], "exec");
@@ -910,6 +1032,19 @@ function connect(connection: Connection): Handle {
     }
     return run(text, params, mode);
   };
+  const rawCopy = connection.copy;
+  const copy =
+    rawCopy === undefined
+      ? undefined
+      : async (statement: string, payload: string): Promise<number> => {
+          if (syncOpen) {
+            unsupported("a COPY while a synchronous transaction is open");
+          }
+          while (gate !== null) {
+            await gate;
+          }
+          return rawCopy(statement, payload);
+        };
   const transactionSync = <T>(body: (tx: SyncQueryable) => T): T => {
     if (runSync == null) {
       unsupported("a synchronous transaction on an asynchronous connection");
@@ -936,6 +1071,7 @@ function connect(connection: Connection): Handle {
     transaction,
     querySync,
     transactionSync,
+    ...(copy === undefined ? null : { copy }),
   };
 }
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@uniflowed/sql",
-  "version": "0.35.0",
+  "version": "0.36.0",
   "description": "The runtime for sqlc's Flow target: exact codecs and driver adapters, part of the Unified Toolchain for Flow.",
   "type": "module",
   "license": "MIT",