feat: Server Beam — sql.on across two Destinations

apps/web/src/backend/api/code-cell-bridge.ts

@@ -6,12 +6,16 @@
  * connection, the write policy, and the row cap; the worker only sees rows.
  */
 
-/** Worker → parent: run this statement on the cell's connection. */
+/** Worker → parent: run this statement on a connection (optional Server Beam alias). */
 export interface CellQueryRequest {
   type: 'cell-query';
   id: number;
   text: string;
   params: unknown[];
+  /** Server Beam endpoint alias (`sql.on('source')`). Omit = default connection. */
+  alias?: string;
+  /** True when the call used `sql.on(...)` (counts toward the per-Execute cap). */
+  viaOn?: boolean;
 }
 
 /** Parent → worker: the outcome of one `cell-query`. */
@@ -28,13 +32,20 @@ export interface CellDoneMessage {
 export type WorkerToParent = CellQueryRequest | CellDoneMessage;
 
 export function isCellQueryRequest(msg: unknown): msg is CellQueryRequest {
-  return (
-    typeof msg === 'object' &&
-    msg !== null &&
-    (msg as { type?: unknown }).type === 'cell-query' &&
-    typeof (msg as { id?: unknown }).id === 'number' &&
-    typeof (msg as { text?: unknown }).text === 'string'
-  );
+  if (
+    typeof msg !== 'object' ||
+    msg === null ||
+    (msg as { type?: unknown }).type !== 'cell-query' ||
+    typeof (msg as { id?: unknown }).id !== 'number' ||
+    typeof (msg as { text?: unknown }).text !== 'string'
+  ) {
+    return false;
+  }
+  const alias = (msg as { alias?: unknown }).alias;
+  if (alias !== undefined && typeof alias !== 'string') return false;
+  const viaOn = (msg as { viaOn?: unknown }).viaOn;
+  if (viaOn !== undefined && typeof viaOn !== 'boolean') return false;
+  return true;
 }
 
 export function isCellDoneMessage(msg: unknown): msg is CellDoneMessage {

apps/web/src/backend/api/code-cell-execute.test.ts

@@ -5,6 +5,7 @@ import {
   runCodeCellOnServer,
   validateCodeCellRequest,
 } from './code-cell-execute';
+import { MAX_SQL } from '../../shared/server-beam';
 
 /** Value planted in APP_ENCRYPTION_KEY to prove an escaped cell cannot read it. */
 const SENTINEL_SECRET = 'sentinel-must-not-leak';
@@ -240,4 +241,50 @@ describe('code cell SQL bridge', () => {
     );
     expect(result.ok).toBe(true);
   }, 45_000);
+
+  it('routes sql.on(alias) to the matching Server Beam endpoint', async () => {
+    const calls: { text: string; alias?: string }[] = [];
+    const runQuery = async (text: string, _params: unknown[], alias?: string) => {
+      calls.push({ text, alias });
+      return [{ hop: alias ?? 'none' }];
+    };
+    const result = await runCell(
+      `const a = await sql.on('source')\`SELECT \${1} AS n\`;` +
+        `const b = await sql.on('target')\`SELECT \${2} AS n\`;` +
+        `return [{ a: a[0].hop, b: b[0].hop }];`,
+      {
+        dialect: 'sqlite',
+        allowWrites: false,
+        runQuery,
+        beamDialects: { source: 'sqlite', target: 'postgres' },
+        defaultBeamAlias: 'source',
+        enforceBeamSqlOnCap: true,
+      }
+    );
+    if (!result.ok) throw new Error(result.error);
+    expect(calls.map((c) => c.alias)).toEqual(['source', 'target']);
+    expect(calls[0]!.text).toBe('SELECT ? AS n');
+    expect(calls[1]!.text).toBe('SELECT $1 AS n');
+    expect(result.rows).toEqual([['source', 'target']]);
+  }, 30_000);
+
+  it('rejects sql.on() calls beyond the Server Beam cap', async () => {
+    const runQuery = async () => [{ n: 1 }];
+    const body =
+      'const out = [];\n' +
+      `for (let i = 0; i < ${MAX_SQL + 1}; i++) {\n` +
+      "  out.push(await sql.on('source')`SELECT ${'i'} AS n`);\n" +
+      '}\n' +
+      'return out;';
+    const result = await runCell(body, {
+      dialect: 'sqlite',
+      allowWrites: false,
+      runQuery,
+      beamDialects: { source: 'sqlite' },
+      defaultBeamAlias: 'source',
+      enforceBeamSqlOnCap: true,
+    });
+    expect(result.ok).toBe(false);
+    if (!result.ok) expect(result.error).toMatch(new RegExp(`at most ${MAX_SQL} sql\\.on`, 'i'));
+  }, 60_000);
 });

apps/web/src/backend/api/code-cell-execute.ts

@@ -18,9 +18,11 @@ import {
 /** Runs one bridged statement for a cell and returns rows as objects. */
 export type CellQueryRunner = (
   text: string,
-  params: unknown[]
+  params: unknown[],
+  alias?: string
 ) => Promise<Record<string, unknown>[]>;
 import { clampMaxRows } from './sql-execute';
+import { MAX_SQL } from '../../shared/server-beam';
 
 export const MAX_CODE_CELL_LENGTH = 100_000;
 export const DEFAULT_CODE_CELL_TIMEOUT_MS = 10_000;
@@ -129,6 +131,12 @@ function runInWorkerThread(args: {
   allowWrites?: boolean;
   /** Runs one bridged `sql` statement. Absent = the cell has no connection. */
   runQuery?: CellQueryRunner;
+  /** Server Beam: alias → dialect for the worker renderer. */
+  beamDialects?: Record<string, string>;
+  /** Server Beam: default alias for plain `sql`…``. */
+  defaultBeamAlias?: string;
+  /** When true, enforce max `sql.on()` calls per Execute. */
+  enforceBeamSqlOnCap?: boolean;
 }): Promise<CodeCellResult> {
   return new Promise((resolve) => {
     let settled = false;
@@ -195,6 +203,8 @@ function runInWorkerThread(args: {
           maxRows: args.maxRows,
           dialect: args.dialect,
           allowWrites: args.allowWrites,
+          beamDialects: args.beamDialects,
+          defaultBeamAlias: args.defaultBeamAlias,
         },
         execArgv,
       });
@@ -205,6 +215,8 @@ function runInWorkerThread(args: {
 
     startTimer();
 
+    let sqlOnCount = 0;
+
     const answerQuery = async (req: CellQueryRequest) => {
       pauseClock();
       const reply = (res: CellQueryResponse) => {
@@ -216,7 +228,15 @@ function runInWorkerThread(args: {
       };
       try {
         if (!args.runQuery) throw new Error('This cell has no connection — select a credential first');
-        const rows = await args.runQuery(req.text, req.params);
+        if (args.enforceBeamSqlOnCap && req.viaOn) {
+          sqlOnCount += 1;
+          if (sqlOnCount > MAX_SQL) {
+            throw new Error(
+              `Server Beam allows at most ${MAX_SQL} sql.on() calls per editor Execute`
+            );
+          }
+        }
+        const rows = await args.runQuery(req.text, req.params, req.alias);
         reply({ type: 'cell-query-result', id: req.id, ok: true, rows, rowCount: rows.length });
       } catch (error: unknown) {
         reply({ type: 'cell-query-result', id: req.id, ok: false, error: errorMessage(error) });
@@ -252,7 +272,14 @@ function runInWorkerThread(args: {
  */
 export async function runCodeCellOnServer(
   validated: ValidatedCodeCell,
-  options?: { dialect?: string; allowWrites?: boolean; runQuery?: CellQueryRunner }
+  options?: {
+    dialect?: string;
+    allowWrites?: boolean;
+    runQuery?: CellQueryRunner;
+    beamDialects?: Record<string, string>;
+    defaultBeamAlias?: string;
+    enforceBeamSqlOnCap?: boolean;
+  }
 ): Promise<CodeCellResult & { durationMs: number }> {
   const started = Date.now();
   let body = validated.body;
@@ -277,6 +304,9 @@ export async function runCodeCellOnServer(
     dialect: options?.dialect,
     allowWrites: options?.allowWrites,
     runQuery: options?.runQuery,
+    beamDialects: options?.beamDialects,
+    defaultBeamAlias: options?.defaultBeamAlias,
+    enforceBeamSqlOnCap: options?.enforceBeamSqlOnCap,
   });
   return { ...result, durationMs: Date.now() - started };
 }

apps/web/src/backend/api/code-cell-query.ts

@@ -73,3 +73,24 @@ export function makeCellQueryRunner(
     return rows.length > MAX_CELL_QUERY_ROWS ? rows.slice(0, MAX_CELL_QUERY_ROWS) : rows;
   };
 }
+
+/**
+ * Server Beam router: pick a per-alias runner. Unknown aliases fail closed.
+ */
+export function makeBeamCellQueryRunner(
+  byAlias: Map<string, CellQueryRunner>,
+  defaultAlias?: string
+): CellQueryRunner {
+  return async (text, params, alias) => {
+    const key = alias ?? defaultAlias;
+    if (!key) {
+      throw new Error('Server Beam query missing alias — use sql.on("source") or sql.on("target")');
+    }
+    const runner = byAlias.get(key);
+    if (!runner) {
+      const known = [...byAlias.keys()].join(', ') || '(none)';
+      throw new Error(`Unknown Server Beam alias "${key}". Known: ${known}`);
+    }
+    return runner(text, params, key);
+  };
+}

apps/web/src/backend/api/code-cell-thread.ts

@@ -23,8 +23,12 @@ type Payload = {
   last: CodeCellLast;
   vars: CodeCellVars;
   maxRows: number;
-  /** Dialect of the cell's connection — decides placeholder + quoting style. */
+  /** Dialect of the default (non-beam) connection. */
   dialect?: string;
+  /** Alias → dialect for Server Beam (`sql.on`). */
+  beamDialects?: Record<string, string>;
+  /** Default alias when plain `sql`…`` is used under Server Beam. */
+  defaultBeamAlias?: string;
   /** When false, the parent rejects write/DDL statements from `sql`. */
   allowWrites?: boolean;
 };
@@ -45,38 +49,103 @@ parentPort?.on('message', (msg: CellQueryResponse) => {
   else waiter.reject(new Error(msg.error));
 });
 
+type SqlBinding = ReturnType<typeof makeSqlBinding>;
+
 /**
  * `sql` inside a cell. Renders the tagged template to `{ text, params }` for
  * the connection's dialect, then asks the parent to run it — the worker never
  * touches a driver itself.
+ *
+ * Server Beam: `sql.on('source')`…`` / `sql.on('target')`…`` select an endpoint.
  */
-function makeSqlBinding(dialect: string) {
-  const run = (query: SqlQuery): Promise<Record<string, unknown>[]> => {
+function makeSqlBinding(opts: {
+  dialect?: string;
+  beamDialects?: Record<string, string>;
+  defaultBeamAlias?: string;
+}): SqlBinding {
+  const beamDialects = opts.beamDialects ?? {};
+  const hasBeam = Object.keys(beamDialects).length > 0;
+
+  const run = (
+    query: SqlQuery,
+    alias: string | undefined,
+    viaOn: boolean
+  ): Promise<Record<string, unknown>[]> => {
     if (!isSqlQuery(query)) {
       return Promise.reject(
         new Error('sql`…` must be used as a tagged template: sql`SELECT 1`, not sql("SELECT 1")')
       );
     }
     if (!parentPort) return Promise.reject(new Error('No SQL bridge available in this context'));
+
+    let dialect = opts.dialect ?? 'postgres';
+    let resolvedAlias = alias;
+    if (hasBeam) {
+      const key = alias ?? opts.defaultBeamAlias;
+      if (!key || !beamDialects[key]) {
+        const known = Object.keys(beamDialects).join(', ') || '(none)';
+        return Promise.reject(
+          new Error(
+            alias
+              ? `Unknown Server Beam alias "${alias}". Known: ${known}`
+              : `Server Beam needs sql.on('alias') — known aliases: ${known}`
+          )
+        );
+      }
+      resolvedAlias = key;
+      dialect = beamDialects[key]!;
+    } else if (alias) {
+      return Promise.reject(
+        new Error(
+          'sql.on() needs Server Beam endpoints — check two Destinations (source, then target) and re-run'
+        )
+      );
+    }
+
     const { text, params } = renderSqlQuery(query, dialect);
     const id = nextQueryId++;
     return new Promise<Record<string, unknown>[]>((resolve, reject) => {
       pending.set(id, { resolve, reject });
-      parentPort!.postMessage({ type: 'cell-query', id, text, params });
+      parentPort!.postMessage({
+        type: 'cell-query',
+        id,
+        text,
+        params,
+        alias: resolvedAlias,
+        viaOn,
+      });
     });
   };
 
-  // Callable as a tag, and carrying the fragment helpers (sql.values, sql.id, …).
+  const attachHelpers = (
+    tag: (strings: TemplateStringsArray, ...values: unknown[]) => Promise<Record<string, unknown>[]>
+  ) =>
+    Object.assign(tag, {
+      raw: sqlTag.raw,
+      id: sqlTag.id,
+      values: sqlTag.values,
+      list: sqlTag.list,
+      run: (query: SqlQuery) => run(query, undefined, false),
+      on(alias: string) {
+        if (typeof alias !== 'string' || !alias.trim()) {
+          throw new Error(`sql.on(alias) needs a non-empty alias string`);
+        }
+        const a = alias.trim();
+        const bound = (strings: TemplateStringsArray, ...values: unknown[]) =>
+          run(sqlTag(strings, ...values), a, true);
+        return Object.assign(bound, {
+          raw: sqlTag.raw,
+          id: sqlTag.id,
+          values: sqlTag.values,
+          list: sqlTag.list,
+          run: (query: SqlQuery) => run(query, a, true),
+        });
+      },
+    });
+
   const tag = (strings: TemplateStringsArray, ...values: unknown[]) =>
-    run(sqlTag(strings, ...values));
-  return Object.assign(tag, {
-    raw: sqlTag.raw,
-    id: sqlTag.id,
-    values: sqlTag.values,
-    list: sqlTag.list,
-    /** Escape hatch for callers holding an already-built query. */
-    run,
-  });
+    run(sqlTag(strings, ...values), undefined, false);
+  return attachHelpers(tag);
 }
 
 async function main() {
@@ -86,7 +155,11 @@ async function main() {
     last: data.last,
     vars: data.vars,
     maxRows: data.maxRows,
-    sql: makeSqlBinding(data.dialect ?? 'postgres'),
+    sql: makeSqlBinding({
+      dialect: data.dialect,
+      beamDialects: data.beamDialects,
+      defaultBeamAlias: data.defaultBeamAlias,
+    }),
   });
   parentPort?.postMessage({ type: 'cell-done', result });
 }