okengine 0.23.0 → 0.23.2

package/src/elements/store/sql-session.ts

@@ -6,6 +6,7 @@
  */
 
 import type { DomainDdlMode } from "../../config/index.ts";
+import type { SignalEmitOptions, SignalTransaction } from "../../drivers/signal-types.ts";
 import type { ClassificationMap, SqlConnection, SqlRow } from "../../drivers/types.ts";
 import {
   buildRlsIdentityPreludeSql,
@@ -16,6 +17,7 @@ import {
 import { throwOke } from "../../kernel/errors.ts";
 import { isFlowFailure } from "../../kernel/hooks.ts";
 import { lazyRequire } from "../../kernel/lazy-require.ts";
+import { currentSignalTransaction, runInSignalTransaction } from "../../kernel/signal-tx.ts";
 import { isRetryableSqlError, sqlErrorToFailure } from "./sql-errors.ts";
 import { AsyncLocalStorage } from "node:async_hooks";
 import { maskRows, tableFromSql } from "./classify.ts";
@@ -51,6 +53,24 @@ export type { WhereMap } from "./sql-condition.ts";
  */
 export type InferSelectRow<T> = T extends { readonly $inferSelect: infer R } ? R : SqlRow;
 
+/**
+ * Insert row from a Drizzle table's `$inferInsert`, or {@link SqlRow}.
+ *
+ * @typeParam T - Table handle or Drizzle table
+ */
+export type InferInsertRow<T> = T extends { readonly $inferInsert: infer R } ? R : SqlRow;
+
+/**
+ * Equality map or a Drizzle SQL fragment (`queryChunks`).
+ * A bare `unknown` would swallow the row type.
+ *
+ * @typeParam TRow - Selected or inserted row
+ */
+export type SqlPredicate<TRow> =
+  | { readonly [K in keyof TRow]?: TRow[K] }
+  | { readonly queryChunks: readonly unknown[] }
+  | undefined;
+
 /** Serialize stamp frames on a shared connection so concurrent identities cannot interleave. */
 const rlsStampTails = new WeakMap<SqlConnection, Promise<unknown>>();
 
@@ -119,6 +139,29 @@ function notifySqlCdc(event: {
  */
 const cdcMutationStorage = new AsyncLocalStorage<{ readonly mutationId: string }>();
 
+/**
+ * One frame of {@link SqlStoreHandle.transaction}.
+ * Depth 0 is the driver `BEGIN`. Deeper frames are savepoints on that pin.
+ */
+interface SqlTxFrame {
+  readonly connection: SqlConnection;
+  /** 0 for the driver transaction; 1+ for `SAVEPOINT`. */
+  readonly depth: number;
+  /**
+   * Promise chain for direct child transactions on this pin.
+   * Siblings append here so their statements cannot interleave. A child
+   * waits on this chain, not on a connection-wide tail, so a nested call
+   * does not deadlock behind its own parent.
+   */
+  childTail: Promise<void>;
+}
+
+/** Pinned connection for {@link SqlStoreHandle.transaction}. */
+const sqlTxStorage = new AsyncLocalStorage<SqlTxFrame>();
+
+/** Monotonic savepoint counter, keyed by the pinned connection. */
+const savepointSeq = new WeakMap<SqlConnection, number>();
+
 /**
  * Read the ambient mutation id, or `undefined` outside a stamped request.
  */
@@ -193,6 +236,110 @@ async function withRlsStampLock<T>(connection: SqlConnection, fn: () => Promise<
   }
 }
 
+/**
+ * Next `oke_sp_<depth>_<n>` name. `n` is per pinned connection so siblings
+ * and nested savepoints never reuse a name, even after `ROLLBACK TO`.
+ *
+ * @param connection - Pinned connection
+ * @param depth - Savepoint depth (1 for the first nested call)
+ */
+function nextSavepointId(connection: SqlConnection, depth: number): string {
+  const n = (savepointSeq.get(connection) ?? 0) + 1;
+  savepointSeq.set(connection, n);
+  return `oke_sp_${String(depth)}_${String(n)}`;
+}
+
+/**
+ * Run `fn` after earlier direct children of `frame` finish.
+ *
+ * @param frame - Parent transaction frame
+ * @param fn - Child body, including savepoint control
+ */
+async function withSqlChildLock<T>(frame: SqlTxFrame, fn: () => Promise<T>): Promise<T> {
+  const prev = frame.childTail;
+  let release!: () => void;
+  const gate = new Promise<void>((resolve) => {
+    release = resolve;
+  });
+  frame.childTail = prev.then(
+    () => gate,
+    () => gate,
+  );
+  await prev.catch(() => undefined);
+  try {
+    return await fn();
+  } finally {
+    release();
+  }
+}
+
+/** One staged signal op inside a nested transaction. */
+type StagedSignalOp =
+  | { readonly kind: "write"; readonly key: string; readonly value: unknown }
+  | {
+      readonly kind: "emit";
+      readonly signal: string;
+      readonly payload: unknown;
+      readonly options?: SignalEmitOptions;
+    };
+
+/**
+ * Child staging scope. Emits buffer here and merge into `parent` on success.
+ * Discard drops them. Neither path publishes — the outermost commit does.
+ *
+ * @param parent - Ambient signal transaction
+ */
+function openChildSignalScope(parent: SignalTransaction): {
+  readonly transaction: SignalTransaction;
+  readonly merge: () => Promise<void>;
+  readonly discard: () => void;
+} {
+  const staged: StagedSignalOp[] = [];
+  let open = true;
+
+  function assertOpen(): void {
+    if (!open) throw new Error("transaction finished");
+  }
+
+  const transaction: SignalTransaction = {
+    async write(key, value) {
+      assertOpen();
+      staged.push({ kind: "write", key, value });
+    },
+    async emit(signal, payload, options) {
+      assertOpen();
+      staged.push(
+        options === undefined
+          ? { kind: "emit", signal, payload }
+          : { kind: "emit", signal, payload, options },
+      );
+    },
+    async commit() {
+      throw new Error("nested signal scope publishes on the outermost commit");
+    },
+    async rollback() {
+      open = false;
+      staged.length = 0;
+    },
+  };
+
+  return {
+    transaction,
+    async merge() {
+      if (!open) return;
+      open = false;
+      for (const op of staged) {
+        if (op.kind === "write") await parent.write(op.key, op.value);
+        else await parent.emit(op.signal, op.payload, op.options);
+      }
+    },
+    discard() {
+      open = false;
+      staged.length = 0;
+    },
+  };
+}
+
 /** Options for a SQL session. */
 export interface SqlSessionOptions {
   /** Open connection (already role-routed). */
@@ -274,7 +421,7 @@ export interface SelectFromBuilder<TRow = SqlRow> extends PromiseLike<TRow[]> {
    *
    * @param where - Condition
    */
-  where(where: unknown): SelectWhereBuilder<TRow>;
+  where(where: SqlPredicate<TRow>): SelectWhereBuilder<TRow>;
   /**
    * Order rows with Drizzle `asc()` / `desc()` terms.
    *
@@ -312,14 +459,18 @@ export interface SelectBuilder<TLocked extends SqlRow | undefined = undefined> {
   ): SelectFromBuilder<TLocked extends undefined ? InferSelectRow<TTable> : TLocked>;
 }
 
-/** Fluent insert builder. */
-export interface InsertBuilder {
+/**
+ * Fluent insert builder.
+ *
+ * @typeParam TRow - `$inferInsert` when the table is a Drizzle table
+ */
+export interface InsertBuilder<TRow = SqlRow> {
   /**
    * Provide row values.
    *
    * @param row - Row to insert
    */
-  values(row: SqlRow): InsertValuesBuilder;
+  values(row: TRow): InsertValuesBuilder;
 }
 
 /**
@@ -340,7 +491,7 @@ export interface DeleteBuilder {
    *
    * @param where - Equality map or Drizzle SQL
    */
-  where(where: unknown): Promise<number>;
+  where(where: SqlPredicate<SqlRow>): Promise<number>;
 }
 
 /** Fluent update builder. */
@@ -360,7 +511,7 @@ export interface UpdateSetBuilder {
    *
    * @param where - Equality map or Drizzle SQL
    */
-  where(where: unknown): Promise<number>;
+  where(where: SqlPredicate<SqlRow>): Promise<number>;
 }
 
 /**
@@ -393,7 +544,7 @@ export interface SqlStoreHandle {
    *
    * @param table - Target table
    */
-  insert(table: TableHandle | unknown): InsertBuilder;
+  insert<TTable>(table: TTable): InsertBuilder<InferInsertRow<TTable>>;
   /**
    * Start an update on `table`.
    *
@@ -415,6 +566,25 @@ export interface SqlStoreHandle {
    */
   delete(table: TableHandle | unknown): DeleteBuilder;
   delete(table: TableHandle | unknown, id: string): Promise<boolean>;
+  /**
+   * Run `fn` on one pinned connection. `fx.emit` inside `fn` stages on the
+   * signal outbox and publishes only after this SQL transaction commits.
+   * A call that already has a pinned connection opens `SAVEPOINT` on that
+   * connection. Nested emits stage on a child scope and merge into the
+   * parent on success; only the outermost commit publishes.
+   *
+   * @param fn - Transaction body. `tx` is this handle, bound to the pin.
+   */
+  transaction<T>(fn: (tx: SqlStoreHandle) => Promise<T>): Promise<T>;
+  /**
+   * Execute a Drizzle select, insert, update, or delete builder.
+   * Calls `toSQL()` so `drizzle-orm` stays off the edge graph.
+   *
+   * @param query - Builder with `toSQL()`
+   */
+  run<T extends SqlRow = SqlRow>(query: {
+    toSQL(): { readonly sql: string; readonly params: readonly unknown[] };
+  }): Promise<T[]>;
   /**
    * True when at least one row matches.
    *
@@ -564,6 +734,10 @@ export function createSqlStoreHandle(
     }
   }
 
+  function activeConnection(): SqlConnection {
+    return sqlTxStorage.getStore()?.connection ?? connection;
+  }
+
   function query(sql: string, params: readonly unknown[] = []): Promise<SqlRow[]> {
     return withSchemaGuard(() => withRlsStamp((conn) => conn.query(sql, params), sql));
   }
@@ -572,9 +746,29 @@ export function createSqlStoreHandle(
     return withSchemaGuard(() => withRlsStamp((conn) => conn.exec(sql, params), sql));
   }
 
+  /**
+   * Rewrite Drizzle `$1` placeholders to the connection's `?` form.
+   *
+   * @param sql - Dialect SQL
+   * @param params - Bindings in `$n` order (1-based)
+   */
+  function fromDialect(
+    sql: string,
+    params: readonly unknown[],
+  ): { sql: string; params: unknown[] } {
+    if (!sql.includes("$")) return { sql, params: [...params] };
+    const ordered: unknown[] = [];
+    const text = sql.replace(/\$(\d+)/g, (_match, index: string) => {
+      ordered.push(params[Number(index) - 1]);
+      return "?";
+    });
+    return { sql: text, params: ordered };
+  }
+
   async function withRlsStamp<T>(fn: (conn: SqlConnection) => Promise<T>, sql: string): Promise<T> {
-    if (!rls || !RLS_CONTEXT_DRIVERS.has(connection.driverId)) return fn(connection);
-    if (isRlsStampExemptSql(sql)) return fn(connection);
+    if (!rls || !RLS_CONTEXT_DRIVERS.has(connection.driverId)) return fn(activeConnection());
+    if (isRlsStampExemptSql(sql)) return fn(activeConnection());
+    if (sqlTxStorage.getStore()) return fn(activeConnection());
     const run = (): Promise<T> => applyRlsStamp(fn);
     // PGlite is one backend session — concurrent identities must not interleave.
     // Pooled postgres pins each stamp via `transaction()` instead.
@@ -910,6 +1104,53 @@ export function createSqlStoreHandle(
     };
   }
 
+  /**
+   * Nested `transaction()` on the pinned connection: `SAVEPOINT`, then
+   * `RELEASE` or `ROLLBACK TO`. Sibling calls share `parent.childTail`.
+   *
+   * @param parent - Frame already on the async store
+   * @param fn - Nested body
+   */
+  async function runNestedTransaction<T>(
+    parent: SqlTxFrame,
+    fn: (tx: SqlStoreHandle) => Promise<T>,
+  ): Promise<T> {
+    return withSqlChildLock(parent, async () => {
+      const depth = parent.depth + 1;
+      const name = quoteIdent(nextSavepointId(parent.connection, depth));
+      const conn = parent.connection;
+      await conn.exec(`SAVEPOINT ${name}`);
+      const child: SqlTxFrame = {
+        connection: conn,
+        depth,
+        childTail: Promise.resolve(),
+      };
+      const parentSignal = currentSignalTransaction();
+      const staging = parentSignal ? openChildSignalScope(parentSignal) : undefined;
+      let released = false;
+      try {
+        const result = await sqlTxStorage.run(child, () => {
+          const body = () => fn(handle as SqlStoreHandle);
+          return staging ? runInSignalTransaction(staging.transaction, body) : body();
+        });
+        await conn.exec(`RELEASE SAVEPOINT ${name}`);
+        released = true;
+        if (staging) await staging.merge();
+        return result;
+      } catch (err) {
+        if (!released) {
+          try {
+            await conn.exec(`ROLLBACK TO SAVEPOINT ${name}`);
+          } catch {
+            // The pin may already be aborted; the original error still wins.
+          }
+        }
+        staging?.discard();
+        throw err;
+      }
+    });
+  }
+
   const handle = {
     ref,
     routedRole: options.routedRole,
@@ -1332,6 +1573,31 @@ export function createSqlStoreHandle(
       });
     },
 
+    async transaction<T>(fn: (tx: SqlStoreHandle) => Promise<T>): Promise<T> {
+      const current = sqlTxStorage.getStore();
+      if (current) return runNestedTransaction(current, fn);
+      if (!connection.transaction) {
+        throw new Error("fx.store().transaction needs SqlConnection.transaction");
+      }
+      return connection.transaction(async (txConn) =>
+        sqlTxStorage.run({ connection: txConn, depth: 0, childTail: Promise.resolve() }, () =>
+          fn(handle as SqlStoreHandle),
+        ),
+      );
+    },
+
+    async run<T extends SqlRow = SqlRow>(builder: {
+      toSQL(): { readonly sql: string; readonly params: readonly unknown[] };
+    }): Promise<T[]> {
+      const compiled = builder.toSQL();
+      const dialect = fromDialect(compiled.sql, compiled.params);
+      if (isSqlDml(dialect.sql)) {
+        await exec(dialect.sql, dialect.params);
+        return [];
+      }
+      return query(dialect.sql, dialect.params) as Promise<T[]>;
+    },
+
     async ensureTable(table: TableHandle) {
       const cols = Object.values(table.columns);
       const pk = resolvePkColumn(table);

package/src/elements/store.ts

@@ -137,6 +137,7 @@ export type {
   SelectWhereBuilder,
   SelectOrderBuilder,
   InferSelectRow,
+  SqlPredicate,
   InsertBuilder,
   InsertValuesBuilder,
   SqlSessionOptions,