@crvouga/mockingbird-service-postgres 3.0.0 → 3.0.2

package/dist/wire/connection.d.ts

@@ -102,6 +102,12 @@ export declare class Connection implements Session {
     /** Rebase this transaction's successful writes onto the latest committed state. */
     private openWorkspace;
     private replay;
+    /**
+     * Hold the primary key until the transaction ends. PostgreSQL's unique index makes a
+     * second `INSERT` of the same key wait; `ON CONFLICT DO UPDATE` then sees the committed row.
+     * Without this, each session's workspace inserts the key and commit replay raises 23505.
+     */
+    private lockInsertedPrimaryKey;
     private runQuery;
     private lockRows;
     private rowLockKey;
@@ -125,9 +131,12 @@ export declare class Connection implements Session {
     private bind;
     private describe;
     /**
-     * The row shape of a statement that returns rows, for Describe on a statement: a trial run with
-     * null parameters inside a transaction that is rolled back. Anything else (or a trial that
-     * fails) is described as returning no data.
+     * The row shape of a statement that returns rows, for Describe on a statement.
+     * `INSERT`/`UPDATE`/`DELETE ... RETURNING` is described from the catalog: a trial run binds
+     * nulls and fails NOT NULL, and Describe would answer NoData. Bun then reads one DataRow cell
+     * and treats the rest of the message as a broken protocol frame.
+     * Other row-returning statements are a trial run with null parameters, rolled back. A trial
+     * that fails is described as returning no data.
      */
     private shapeOf;
     private executePortal;

package/dist/wire/index.js

@@ -250,6 +250,12 @@ function pgError(category, message2, sqlState) {
 function unsupported(what) {
   return new PostgresError("unsupported", `${what} is not supported`);
 }
+function isPostgresError(err) {
+  if (err instanceof PostgresError) return true;
+  if (!(err instanceof Error) || err.name !== "PostgresError") return false;
+  const sqlState = err.sqlState;
+  return typeof sqlState === "string" && /^[0-9A-Z]{5}$/.test(sqlState);
+}
 
 // src/runtime/prng.ts
 var Prng = class _Prng {
@@ -3005,6 +3011,7 @@ function arrayElemNeedsQuote(s) {
 function quoteArrayElem(s) {
   return `"${s.replaceAll("\\", "\\\\").replaceAll('"', '\\"')}"`;
 }
+var UTC_OUTPUT = { zoneOffsetAt: () => 0 };
 function arrayText(arr, ctx) {
   if (arr.dims.length === 0) return "{}";
   const needsLb = arr.lbs.some((l) => l !== 1);
@@ -24587,9 +24594,48 @@ var RENAME_DATABASE = new RegExp(`^alter\\s+database\\s+${IDENT}\\s+rename\\s+to
 var identifier = (quoted, plain) => quoted === void 0 ? plain.toLowerCase() : quoted.replaceAll('""', '"');
 var ROW_LOCK = /\s+for\s+(?:no\s+key\s+update|key\s+share|update|share)(?:\s+of\s+(?:(?:"[^"]+"|[A-Za-z_][\w$]*)(?:\.(?:"[^"]+"|[A-Za-z_][\w$]*))?)(?:\s*,\s*(?:(?:"[^"]+"|[A-Za-z_][\w$]*)(?:\.(?:"[^"]+"|[A-Za-z_][\w$]*))?))*)?(?:\s+(nowait|skip\s+locked))?\s*$/i;
 var LIMIT = /\s+limit\s+(\d+)(?:\s+offset\s+(\d+))?\s*$/i;
-var FROM_TABLE = /\bfrom\s+((?:"[^"]+"|[A-Za-z_][\w$]*)(?:\.(?:"[^"]+"|[A-Za-z_][\w$]*))?)/i;
-var UPDATE_TARGET = /^update\s+((?:"[^"]+"|[A-Za-z_][\w$]*)(?:\.(?:"[^"]+"|[A-Za-z_][\w$]*))?)(?:\s+(?:as\s+)?(?:"[^"]+"|[A-Za-z_][\w$]*))?\s+set\s+[\s\S]*?(?:\s+where\s+([\s\S]*?))?(?:\s+returning\s+[\s\S]*)?$/i;
-var DELETE_TARGET = /^delete\s+from\s+((?:"[^"]+"|[A-Za-z_][\w$]*)(?:\.(?:"[^"]+"|[A-Za-z_][\w$]*))?)(?:\s+(?:as\s+)?(?:"[^"]+"|[A-Za-z_][\w$]*))?(?:\s+where\s+([\s\S]*?))?(?:\s+returning\s+[\s\S]*)?$/i;
+var SQL_NAME = String.raw`(?:"[^"]+"|[A-Za-z_][\w$]*)`;
+var SQL_RELATION = String.raw`(?:"[^"]+"|[A-Za-z_][\w$]*)(?:\.(?:"[^"]+"|[A-Za-z_][\w$]*))?`;
+var FROM_TABLE = new RegExp(String.raw`\bfrom\s+(${SQL_RELATION})`, "i");
+var UPDATE_TARGET = new RegExp(
+  String.raw`^update\s+(${SQL_RELATION})(?:\s+(?:as\s+)?(${SQL_NAME}))?\s+set\s+[\s\S]*?(?:\s+where\s+([\s\S]*?))?(?:\s+returning\s+[\s\S]*)?$`,
+  "i"
+);
+var DELETE_TARGET = new RegExp(
+  String.raw`^delete\s+from\s+(${SQL_RELATION})(?:\s+(?:as\s+)?(${SQL_NAME}))?(?:\s+where\s+([\s\S]*?))?(?:\s+returning\s+[\s\S]*)?$`,
+  "i"
+);
+var lockedMutationQuery = (sql) => {
+  const match = UPDATE_TARGET.exec(sql) ?? DELETE_TARGET.exec(sql);
+  const relation = match?.[1];
+  if (!relation) return null;
+  const alias = match?.[2];
+  const where = match?.[3];
+  const from = alias ? `${relation} AS ${alias}` : relation;
+  return `SELECT * FROM ${from}${where ? ` WHERE ${where}` : ""}`;
+};
+var lockLiteral = (expr, params) => {
+  if (!expr) return void 0;
+  switch (expr.type) {
+    case "null_lit":
+      return null;
+    case "string_lit":
+      return expr.value;
+    case "bool_lit":
+      return expr.value ? "t" : "f";
+    case "number_lit":
+      return expr.raw.replace(/^\+/, "");
+    case "param": {
+      const value = params[expr.index - 1];
+      if (value === void 0) return void 0;
+      const typed = bindValueToTyped(value, expr.index - 1);
+      if (typed.v === null) return null;
+      return datumText(typed.t, typed.v, UTC_OUTPUT);
+    }
+    default:
+      return void 0;
+  }
+};
 var TYPLEN = {
   bool: 1,
   int2: 2,
@@ -24608,7 +24654,7 @@ var TEXT_TYPES2 = /* @__PURE__ */ new Set(["text", "varchar", "bpchar", "name",
 var PG_EPOCH_MS = 9466848e5;
 var pgError2 = (category, message2, code) => new PostgresError(category, message2, code);
 var errorFields = (error) => {
-  if (error instanceof PostgresError) {
+  if (isPostgresError(error)) {
     const message2 = error.message;
     const fields = { code: error.sqlState, message: message2 };
     const constraint = /constraint "([^"]+)"/.exec(message2)?.[1];
@@ -24726,6 +24772,46 @@ var decodeParam = (oid, format, bytes) => {
       throw pgError2("internal", `binary parameter format is not supported for type oid ${oid}`, "0A000");
   }
 };
+function returningShape(db, sql) {
+  let statements;
+  try {
+    statements = parse(sql);
+  } catch {
+    return null;
+  }
+  const stmt = statements.length === 1 ? statements[0] : void 0;
+  if (!stmt || stmt.type !== "insert" && stmt.type !== "update" && stmt.type !== "delete") return null;
+  const targets = stmt.returning;
+  if (!targets || targets.length === 0) return null;
+  const table = db.state.findTable(stmt.table);
+  if (!table) return null;
+  const label = stmt.alias ?? table.name;
+  const columns = [];
+  const columnTypes = [];
+  for (const target of targets) {
+    if (target.expr.type === "star") {
+      const qualifier = target.expr.table?.at(-1);
+      if (qualifier !== void 0 && qualifier !== label && qualifier !== table.name) return null;
+      for (const column of table.columns) {
+        columns.push(column.name);
+        columnTypes.push(column.type.id);
+      }
+      continue;
+    }
+    if (target.expr.type === "colref") {
+      const name = target.expr.parts.at(-1) ?? "?column?";
+      const column = table.columns.find((item) => item.name === name);
+      columns.push(target.alias ?? name);
+      columnTypes.push(column?.type.id ?? "text");
+      continue;
+    }
+    columns.push(target.alias ?? inferColumnName(target.expr));
+    columnTypes.push("text");
+  }
+  if (columns.length === 0) return null;
+  const command2 = stmt.type === "insert" ? "INSERT" : stmt.type === "update" ? "UPDATE" : "DELETE";
+  return { columns, columnTypes, rows: [], rowCount: 0, command: command2 };
+}
 var Connection = class {
   constructor(socket, cluster, options) {
     this.socket = socket;
@@ -25004,8 +25090,8 @@ var Connection = class {
       return outcome;
     } catch (error) {
       if (this.inTransaction) this.aborted = true;
-      if (error instanceof PostgresError && error.sqlState === "40P01") this.cluster.unlockAll(this, true);
-      const status = error instanceof PostgresError ? error.sqlState : "XX000";
+      if (isPostgresError(error) && error.sqlState === "40P01") this.cluster.unlockAll(this, true);
+      const status = isPostgresError(error) ? error.sqlState : "XX000";
       this.options.onLog?.({ pid: this.pid, sql: text, durationMs: performance.now() - started, status });
       throw error;
     }
@@ -25201,14 +25287,48 @@ var Connection = class {
       }
     });
   }
+  /**
+   * Hold the primary key until the transaction ends. PostgreSQL's unique index makes a
+   * second `INSERT` of the same key wait; `ON CONFLICT DO UPDATE` then sees the committed row.
+   * Without this, each session's workspace inserts the key and commit replay raises 23505.
+   */
+  lockInsertedPrimaryKey(db, text, params) {
+    let statements;
+    try {
+      statements = parse(text);
+    } catch {
+      return;
+    }
+    const stmt = statements.length === 1 ? statements[0] : void 0;
+    if (stmt?.type !== "insert" || stmt.source === "default_values" || stmt.source.body.type !== "values") return;
+    if (!stmt.columns) return;
+    const table = db.state.findTable(stmt.table);
+    const primary = table?.constraints.find((constraint) => constraint.kind === "primary_key");
+    if (!primary || primary.columns.length === 0) return;
+    const indexes = primary.columns.map((column) => stmt.columns?.indexOf(column) ?? -1);
+    if (indexes.some((index) => index < 0)) return;
+    const identities = [];
+    for (const row of stmt.source.body.rows) {
+      const identity = [];
+      for (const index of indexes) {
+        const value = lockLiteral(row[index], params);
+        if (value === void 0) return;
+        identity.push(value);
+      }
+      identities.push(identity);
+    }
+    const relation = stmt.table.join(".");
+    for (const identity of identities) {
+      const key = `row:${this.databaseName}:${relation}:${JSON.stringify(identity)}`;
+      if (!this.cluster.tryLock(this, key, true)) throw new LockWait(key);
+    }
+  }
   runQuery(db, text, params) {
     const lock = ROW_LOCK.exec(text);
     if (!lock) {
-      const mutation = UPDATE_TARGET.exec(text) ?? DELETE_TARGET.exec(text);
-      if (mutation) {
-        const lockQuery = `SELECT * FROM ${mutation[1]}${mutation[2] ? ` WHERE ${mutation[2]}` : ""}`;
-        this.lockRows(db, lockQuery, params, "wait");
-      }
+      const lockQuery = lockedMutationQuery(text);
+      if (lockQuery) this.lockRows(db, lockQuery, params, "wait");
+      this.lockInsertedPrimaryKey(db, text, params);
       return db.prepare(text).textResult(...params);
     }
     let query = text.slice(0, lock.index).trimEnd();
@@ -25572,11 +25692,16 @@ var Connection = class {
     else this.write(backend.rowDescription(this.fields(outcome.result, portal.resultFormats)));
   }
   /**
-   * The row shape of a statement that returns rows, for Describe on a statement: a trial run with
-   * null parameters inside a transaction that is rolled back. Anything else (or a trial that
-   * fails) is described as returning no data.
+   * The row shape of a statement that returns rows, for Describe on a statement.
+   * `INSERT`/`UPDATE`/`DELETE ... RETURNING` is described from the catalog: a trial run binds
+   * nulls and fails NOT NULL, and Describe would answer NoData. Bun then reads one DataRow cell
+   * and treats the rest of the message as a broken protocol frame.
+   * Other row-returning statements are a trial run with null parameters, rolled back. A trial
+   * that fails is described as returning no data.
    */
   async shapeOf(prepared) {
+    const described = returningShape(this.database, prepared.sql);
+    if (described) return described;
     if (!RETURNS_ROWS.test(prepared.sql)) return null;
     const db = this.database;
     await this.cluster.acquireTurn(this);