@pikku/kysely 0.13.6 → 0.13.10

package/src/kysely-virtual-user-run-store.test.ts

@@ -0,0 +1,212 @@
+import { describe, test, beforeEach } from 'node:test'
+import assert from 'node:assert/strict'
+import Database from 'better-sqlite3'
+import { CamelCasePlugin, Kysely, SqliteDialect } from 'kysely'
+
+import { KyselyVirtualUserRunStore } from './kysely-virtual-user-run-store.js'
+import { SerializePlugin } from './serialize-plugin.js'
+import type { KyselyPikkuDB } from './kysely-tables.js'
+
+// SerializePlugin is here because most projects in this package run with it, and
+// it is the harder case: it deserialises JSON columns the store has already
+// stringified. The store has to survive that, and its absence.
+const newDb = () =>
+  new Kysely<KyselyPikkuDB>({
+    dialect: new SqliteDialect({ database: new Database(':memory:') }),
+    plugins: [new CamelCasePlugin(), new SerializePlugin()],
+  })
+
+const TALLY = {
+  steps: 12,
+  calls: 30,
+  mutations: 4,
+  tokensIn: 900,
+  tokensOut: 120,
+  elapsed: 4200,
+} as any
+
+const FINDING = {
+  kind: 'schema-mismatch',
+  detail: 'listDocs returned a doc with no title',
+  rpcName: 'listDocs',
+  step: 7,
+} as any
+
+describe('KyselyVirtualUserRunStore', () => {
+  let db: Kysely<KyselyPikkuDB>
+  let store: KyselyVirtualUserRunStore
+
+  beforeEach(() => {
+    db = newDb()
+    store = new KyselyVirtualUserRunStore(db)
+  })
+
+  // The runtime never needs this table, so it arrives with the feature that
+  // fills it rather than in every database — which means the first write has to
+  // create it.
+  test('creates its table on first use', async () => {
+    const runId = await store.start({
+      persona: 'susan',
+      disposition: 'realistic',
+      seed: 42,
+    })
+    assert.ok(runId)
+  })
+
+  test('a run starts as running, with no findings and no finish time', async () => {
+    const runId = await store.start({
+      persona: 'susan',
+      disposition: 'careless',
+      seed: 7,
+      goals: ['find the export button'],
+      memory: { docId: 'doc_1' },
+      startedBy: 'user_1',
+    })
+
+    const run = await store.get(runId)
+    assert.equal(run?.status, 'running')
+    assert.deepEqual(run?.findings, [])
+    assert.equal(run?.finishedAt, null)
+    assert.equal(run?.error, null)
+    assert.deepEqual(run?.goals, ['find the export button'])
+    assert.deepEqual(run?.memory, { docId: 'doc_1' })
+    assert.equal(run?.startedBy, 'user_1')
+  })
+
+  // The seed is the whole reason a finding is reproducible, and a BIGINT comes
+  // back as a string on some drivers — so it has to survive the round trip as a
+  // number.
+  test('the seed round-trips as a number', async () => {
+    const runId = await store.start({
+      persona: 'susan',
+      disposition: 'realistic',
+      seed: 2_147_483_646,
+    })
+    const run = await store.get(runId)
+    assert.strictEqual(run?.seed, 2_147_483_646)
+  })
+
+  test('completing a run stores what it found', async () => {
+    const runId = await store.start({
+      persona: 'susan',
+      disposition: 'auditor',
+      seed: 1,
+      memory: { docId: 'doc_1' },
+    })
+
+    await store.complete(runId, {
+      findings: [FINDING],
+      tally: TALLY,
+      memory: { docId: 'doc_1', orgId: 'org_9' },
+      stoppedBy: 'budget-steps',
+    })
+
+    const run = await store.get(runId)
+    assert.equal(run?.status, 'completed')
+    assert.equal(run?.findings.length, 1)
+    assert.equal(run?.findings[0]?.detail, FINDING.detail)
+    assert.deepEqual(run?.tally, TALLY)
+    assert.equal(run?.stoppedBy, 'budget-steps')
+    // Overwritten, not merged: the engine's memory already carries what it was
+    // given plus what it learned.
+    assert.deepEqual(run?.memory, { docId: 'doc_1', orgId: 'org_9' })
+    assert.ok(run?.finishedAt instanceof Date)
+  })
+
+  // A crashed run and a run that found nothing are different answers, and this
+  // record is the only place that distinction survives.
+  test('a failed run is failed, not an empty success', async () => {
+    const runId = await store.start({
+      persona: 'susan',
+      disposition: 'realistic',
+      seed: 1,
+    })
+
+    await store.fail(runId, 'SCENARIO_ACTOR_SECRET is not set')
+
+    const run = await store.get(runId)
+    assert.equal(run?.status, 'failed')
+    assert.equal(run?.error, 'SCENARIO_ACTOR_SECRET is not set')
+    assert.deepEqual(run?.findings, [])
+    assert.ok(run?.finishedAt instanceof Date)
+  })
+
+  test('an unknown run is null rather than a throw', async () => {
+    await store.init()
+    assert.equal(await store.get('nope'), null)
+  })
+
+  test('lists newest first', async () => {
+    const first = await store.start({
+      persona: 'susan',
+      disposition: 'realistic',
+      seed: 1,
+    })
+    const second = await store.start({
+      persona: 'raj',
+      disposition: 'realistic',
+      seed: 2,
+    })
+
+    const runs = await store.list()
+    assert.equal(runs.length, 2)
+    assert.deepEqual(
+      new Set(runs.map((run) => run.runId)),
+      new Set([first, second])
+    )
+  })
+
+  // Comparing what a persona finds this week against last week is the whole
+  // reason runs are kept.
+  test('narrows to one persona s history', async () => {
+    await store.start({ persona: 'susan', disposition: 'realistic', seed: 1 })
+    await store.start({ persona: 'raj', disposition: 'realistic', seed: 2 })
+
+    const runs = await store.list({ persona: 'susan' })
+    assert.equal(runs.length, 1)
+    assert.equal(runs[0]?.persona, 'susan')
+  })
+
+  test('completing an unknown run does not throw', async () => {
+    await store.init()
+    await store.complete('nope', {
+      findings: [],
+      tally: TALLY,
+      memory: {},
+      stoppedBy: 'exhausted',
+    })
+  })
+
+  // Not every project installs SerializePlugin, and a bare SQLite driver cannot
+  // bind a Date at all — so the whole round trip has to hold without it.
+  test('round-trips without SerializePlugin', async () => {
+    const bare = new Kysely<KyselyPikkuDB>({
+      dialect: new SqliteDialect({ database: new Database(':memory:') }),
+      plugins: [new CamelCasePlugin()],
+    })
+    const bareStore = new KyselyVirtualUserRunStore(bare)
+
+    const runId = await bareStore.start({
+      persona: 'susan',
+      disposition: 'realistic',
+      seed: 99,
+      goals: ['export a doc'],
+      memory: { docId: 'doc_1' },
+    })
+    await bareStore.complete(runId, {
+      findings: [FINDING],
+      tally: TALLY,
+      memory: { docId: 'doc_1' },
+      stoppedBy: 'exhausted',
+    })
+
+    const run = await bareStore.get(runId)
+    assert.equal(run?.status, 'completed')
+    assert.strictEqual(run?.seed, 99)
+    assert.deepEqual(run?.goals, ['export a doc'])
+    assert.deepEqual(run?.tally, TALLY)
+    assert.equal(run?.findings[0]?.detail, FINDING.detail)
+    assert.ok(run?.createdAt instanceof Date)
+    assert.ok(run?.finishedAt instanceof Date)
+  })
+})

package/src/kysely-virtual-user-run-store.ts

@@ -0,0 +1,149 @@
+import type { Kysely, Selectable } from 'kysely'
+import { randomUUID } from 'node:crypto'
+import type {
+  VirtualUserDisposition,
+  VirtualUserRunOutcome,
+  VirtualUserRunRecord,
+  VirtualUserRunStart,
+  VirtualUserRunStore,
+} from '@pikku/core/virtual-user'
+import type { KyselyPikkuDB, VirtualUserRunTable } from './kysely-tables.js'
+import { parseJson } from './kysely-json.js'
+import { ensurePikkuSchema } from './schema/index.js'
+import { virtualUserSchema } from './schema/virtual-user.schema.js'
+
+/**
+ * Records virtual-user runs in a `virtualUserRun` table.
+ *
+ * The store is the whole record of a run: `runVirtualUser` returns as soon as
+ * the row exists and the engine keeps going without it, so nothing else knows
+ * the run happened. That is why `fail()` exists as its own call — a run that
+ * crashed and a run that found nothing are different answers, and a row left at
+ * `running` is neither.
+ *
+ * Works with or without `SerializePlugin`, which not every project installs: the
+ * JSON columns are written as strings (which the plugin passes through) and read
+ * back through `parseJson` (which passes through what the plugin already
+ * parsed). Timestamps are written as ISO strings for the same reason — a bare
+ * SQLite driver cannot bind a `Date` at all.
+ *
+ * SECURITY: an `adversarial` run's findings are working exploits carrying live
+ * ids. This is a privileged store; the scaffolded reads are scope-gated, and
+ * exposing these rows more widely publishes your own exploits.
+ */
+export class KyselyVirtualUserRunStore implements VirtualUserRunStore {
+  private initialized = false
+
+  constructor(private db: Kysely<KyselyPikkuDB>) {}
+
+  /**
+   * Creates the table on first use, like the audit sink — the runtime does not
+   * need it, so it arrives with the feature that fills it rather than in every
+   * database.
+   */
+  public async init(): Promise<void> {
+    if (this.initialized) return
+    await ensurePikkuSchema(this.db, virtualUserSchema)
+    this.initialized = true
+  }
+
+  async start(run: VirtualUserRunStart): Promise<string> {
+    await this.init()
+    const runId = randomUUID()
+    await this.db
+      .insertInto('virtualUserRun')
+      .values({
+        runId,
+        persona: run.persona,
+        disposition: run.disposition,
+        // Text on every engine (see the schema), so the seed round-trips
+        // through drivers that hand a BIGINT back as a string.
+        seed: String(run.seed),
+        status: 'running',
+        goals: JSON.stringify(run.goals ?? []),
+        memory: JSON.stringify(run.memory ?? {}),
+        findings: '[]',
+        startedBy: run.startedBy ?? null,
+      })
+      .execute()
+    return runId
+  }
+
+  async complete(runId: string, outcome: VirtualUserRunOutcome): Promise<void> {
+    await this.init()
+    await this.db
+      .updateTable('virtualUserRun')
+      .set({
+        status: 'completed',
+        findings: JSON.stringify(outcome.findings),
+        tally: JSON.stringify(outcome.tally),
+        // Overwritten rather than merged: the engine's memory already carries
+        // what it was given, plus what it learned on the way.
+        memory: JSON.stringify(outcome.memory),
+        stoppedBy: outcome.stoppedBy,
+        finishedAt: new Date().toISOString(),
+      })
+      .where('runId', '=', runId)
+      .execute()
+  }
+
+  async fail(runId: string, error: string): Promise<void> {
+    await this.init()
+    await this.db
+      .updateTable('virtualUserRun')
+      .set({
+        status: 'failed',
+        error,
+        finishedAt: new Date().toISOString(),
+      })
+      .where('runId', '=', runId)
+      .execute()
+  }
+
+  async get(runId: string): Promise<VirtualUserRunRecord | null> {
+    await this.init()
+    const row = await this.db
+      .selectFrom('virtualUserRun')
+      .selectAll()
+      .where('runId', '=', runId)
+      .executeTakeFirst()
+    return row ? this.toRecord(row) : null
+  }
+
+  async list(options?: {
+    persona?: string
+    limit?: number
+    offset?: number
+  }): Promise<VirtualUserRunRecord[]> {
+    await this.init()
+    let query = this.db.selectFrom('virtualUserRun').selectAll()
+    if (options?.persona) {
+      query = query.where('persona', '=', options.persona)
+    }
+    const rows = await query
+      .orderBy('createdAt', 'desc')
+      .limit(options?.limit ?? 50)
+      .offset(options?.offset ?? 0)
+      .execute()
+    return rows.map((row) => this.toRecord(row))
+  }
+
+  private toRecord(row: Selectable<VirtualUserRunTable>): VirtualUserRunRecord {
+    return {
+      runId: row.runId,
+      persona: row.persona,
+      disposition: row.disposition as VirtualUserDisposition,
+      seed: Number(row.seed),
+      status: row.status,
+      goals: parseJson(row.goals) ?? [],
+      memory: parseJson(row.memory) ?? {},
+      findings: parseJson(row.findings) ?? [],
+      tally: row.tally ? (parseJson(row.tally) ?? null) : null,
+      stoppedBy: row.stoppedBy,
+      error: row.error,
+      startedBy: row.startedBy,
+      createdAt: new Date(row.createdAt),
+      finishedAt: row.finishedAt ? new Date(row.finishedAt) : null,
+    }
+  }
+}

package/src/kysely-workflow-service.ts

@@ -505,23 +505,6 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
     return { completedNodeIds, failedNodeIds, branchKeys }
   }
 
-  async getNodesWithoutSteps(
-    runId: string,
-    nodeIds: string[]
-  ): Promise<string[]> {
-    if (nodeIds.length === 0) return []
-
-    const result = await this.db
-      .selectFrom('workflowStep')
-      .select('stepName')
-      .where('workflowRunId', '=', runId)
-      .where('stepName', 'in', nodeIds)
-      .execute()
-
-    const existingStepNames = new Set(result.map((r) => r.stepName))
-    return nodeIds.filter((id) => !existingStepNames.has(id))
-  }
-
   async getStepInstances(
     runId: string
   ): Promise<

package/src/schema/audit.schema.ts

@@ -0,0 +1,55 @@
+import type { PikkuSchema } from './pikku-schema.types.js'
+
+/**
+ * The `audit` table {@link KyselyAuditService} writes to and reads back.
+ *
+ * Deliberately not in `pikkuSchemas`: the runtime never needs this table, only
+ * a project that has wired a durable audit sink does, and creating it for
+ * everyone would put an empty table in every database. `KyselyAuditService.init()`
+ * applies it, so it arrives exactly when the sink that fills it does.
+ *
+ * Every column is text on every engine, matching the platform audit-queue
+ * consumer's row shape — a locally-run project and a deployed stage write rows
+ * the same reader can read. `occurredAt` is an ISO 8601 string rather than a
+ * timestamp for that reason: string ordering is chronological ordering, and the
+ * queue consumer has no shared type to agree on.
+ */
+export const auditSchema: PikkuSchema = {
+  name: 'audit',
+  statements: [
+    (db) =>
+      db.schema
+        .createTable('audit')
+        .addColumn('auditId', 'text', (col) => col.primaryKey())
+        .addColumn('occurredAt', 'text', (col) => col.notNull())
+        .addColumn('type', 'text', (col) => col.notNull())
+        .addColumn('source', 'text', (col) => col.defaultTo('auto').notNull())
+        .addColumn('outcome', 'text')
+        .addColumn('functionId', 'text')
+        .addColumn('wireType', 'text')
+        .addColumn('traceId', 'text')
+        .addColumn('transactionId', 'text')
+        .addColumn('queryId', 'text')
+        .addColumn('userId', 'text')
+        .addColumn('orgId', 'text')
+        .addColumn('pikkuUserId', 'text')
+        .addColumn('tables', 'text')
+        .addColumn('changedCols', 'text')
+        .addColumn('event', 'text')
+        .addColumn('old', 'text')
+        .addColumn('data', 'text'),
+
+    // The trail is only ever read newest-first, and the two filters the console
+    // offers are user and type.
+    (db) =>
+      db.schema
+        .createIndex('idx_audit_occurred_at')
+        .on('audit')
+        .column('occurredAt'),
+
+    (db) =>
+      db.schema.createIndex('idx_audit_user_id').on('audit').column('userId'),
+
+    (db) => db.schema.createIndex('idx_audit_type').on('audit').column('type'),
+  ],
+}

package/src/schema/credential.schema.ts

@@ -7,8 +7,9 @@ import { rawStatement, type PikkuSchema } from './pikku-schema.types.js'
  * The uniqueness rule is an expression index — `user_id` is nullable and two
  * rows with a null user must still collide on name — which the schema builder
  * cannot express, so it stays raw. Raw SQL is not rewritten by `withSchema`,
- * which is why the table is named explicitly rather than left to the search
- * path.
+ * which is why the table goes through `ctx.table` rather than being left to the
+ * search path. The index's own name stays unqualified: an index belongs to the
+ * schema of the table it is on, and postgres rejects one that says otherwise.
  */
 export const credentialSchema: PikkuSchema = {
   name: 'credential',
@@ -29,9 +30,19 @@ export const credentialSchema: PikkuSchema = {
         ),
 
     rawStatement(
-      sql`CREATE UNIQUE INDEX credentials_name_user_id_unique ON credentials (name, COALESCE(user_id, ''))`
+      ({ table }) =>
+        sql`CREATE UNIQUE INDEX credentials_name_user_id_unique ON ${table('credentials')} (name, COALESCE(user_id, ''))`
     ),
 
+    (db) =>
+      db.schema
+        .createTable('credentialKekSalts')
+        .addColumn('keyVersion', 'integer', (col) => col.primaryKey())
+        .addColumn('salt', 'varchar(64)', (col) => col.notNull())
+        .addColumn('createdAt', 'timestamp', (col) =>
+          col.defaultTo(sql`CURRENT_TIMESTAMP`).notNull()
+        ),
+
     (db) =>
       db.schema
         .createTable('credentialsAudit')

package/src/schema/index.ts

@@ -1,8 +1,10 @@
 import { CamelCasePlugin, type Kysely } from 'kysely'
-import type {
-  PikkuSchema,
-  RequiredTypes,
-  SchemaRequirement,
+import {
+  schemaContext,
+  type PikkuSchema,
+  type RequiredTypes,
+  type SchemaContext,
+  type SchemaRequirement,
 } from './pikku-schema.types.js'
 import { aiSchema } from './ai.schema.js'
 import { channelSchema } from './channel.schema.js'
@@ -17,11 +19,24 @@ import { workflowSchema } from './workflow.schema.js'
 export type {
   PikkuSchema,
   RequiredTypes,
+  SchemaContext,
   SchemaRequirement,
   SchemaStatement,
   SchemaStatementFactory,
 } from './pikku-schema.types.js'
-export { rawStatement, requiredType } from './pikku-schema.types.js'
+export {
+  rawStatement,
+  requiredType,
+  schemaContext,
+  unqualifiedContext,
+} from './pikku-schema.types.js'
+
+/**
+ * Opt-in, so it is named here rather than folded into `pikkuSchemas` — the
+ * runtime works without an audit table, and `KyselyAuditService.init()` applies
+ * this for the projects that wire the sink.
+ */
+export { auditSchema } from './audit.schema.js'
 
 /**
  * Every table the pikku runtime needs, in dependency order.
@@ -53,6 +68,29 @@ export const pikkuSchemas: PikkuSchema[] = [
  */
 const bind = (db: Kysely<any>) => db.withPlugin(new CamelCasePlugin())
 
+/**
+ * The schema a connection is bound to, or `undefined` for one that is not.
+ *
+ * Read back out of compiled DDL rather than off the connection, because kysely
+ * keeps the `withSchema(...)` setting inside its executor and exposes no way to
+ * ask. Compiling a throwaway `create table` is the same trick `declaredTables`
+ * uses on the real statements, against the same regex — nothing is executed and
+ * the probe never leaves this function.
+ *
+ * Derived rather than passed in so that raw statements follow whatever
+ * connection they are handed. A caller that had to remember to say the schema
+ * twice would eventually say it once.
+ */
+const boundSchema = (db: Kysely<any>): string | undefined =>
+  CREATE_TABLE.exec(
+    db.schema.createTable('pikkuSchemaProbe').addColumn('id', 'text').compile()
+      .sql
+  )?.[1]
+
+/** The context for statements compiled against `db`, matching its binding. */
+const contextFor = (db: Kysely<any>): SchemaContext =>
+  schemaContext(boundSchema(db))
+
 export interface UnmetRequirement {
   schema: PikkuSchema
   requirement: SchemaRequirement
@@ -141,9 +179,10 @@ export const applyPikkuSchemas = async (
 
   await db.transaction().execute(async (trx) => {
     const bound = bind(trx)
+    const ctx = contextFor(bound)
     for (const schema of schemas) {
       for (const statement of schema.statements) {
-        const query = statement(bound, types)
+        const query = statement(bound, types, ctx)
         try {
           await query.execute()
         } catch (error) {
@@ -239,9 +278,10 @@ export const declaredTables = (
   db: Kysely<any>
 ): DeclaredTable[] => {
   const bound = bind(db)
+  const ctx = contextFor(bound)
   const tables: DeclaredTable[] = []
   for (const statement of schema.statements) {
-    const match = CREATE_TABLE.exec(statement(bound, {}).compile().sql)
+    const match = CREATE_TABLE.exec(statement(bound, {}, ctx).compile().sql)
     if (match) tables.push({ schema: match[1], name: match[2]! })
   }
   return tables
@@ -357,17 +397,26 @@ export const ensurePikkuSchema = async (
  * `types` comes from `resolveRequirements` against the database the migration
  * is being written for. Omit it and foreign keys onto another source's columns
  * compile as `text`, which is a shape, not a migration.
+ *
+ * `schema` qualifies every table the migration creates, for a project that
+ * keeps the runtime tables somewhere other than the default search path. It is
+ * applied here and not to `db` itself because the caller's connection is a
+ * throwaway the declaration was just *applied* to, in whatever schema that
+ * database actually has — qualifying that would create tables in a schema the
+ * scratch database has never heard of. Only the rendered SQL is bound.
  */
 export const compilePikkuSchemas = (
   db: Kysely<any>,
   schemas: PikkuSchema[] = pikkuSchemas,
-  types: RequiredTypes = {}
+  types: RequiredTypes = {},
+  schemaName?: string
 ): string => {
-  const bound = bind(db)
+  const bound = schemaName ? bind(db).withSchema(schemaName) : bind(db)
+  const ctx = contextFor(bound)
   return schemas
     .flatMap((schema) =>
       schema.statements.map(
-        (statement) => `${statement(bound, types).compile().sql};`
+        (statement) => `${statement(bound, types, ctx).compile().sql};`
       )
     )
     .join('\n\n')

package/src/schema/pikku-schema.types.ts

@@ -15,9 +15,51 @@ export interface SchemaStatement {
 
 export type SchemaStatementFactory = (
   db: Kysely<any>,
-  types: RequiredTypes
+  types: RequiredTypes,
+  ctx: SchemaContext
 ) => SchemaStatement
 
+/**
+ * What a statement needs to know about the schema its DDL lands in.
+ *
+ * Only raw SQL needs this. The schema builder learns the schema from the
+ * connection — `withSchema('app')` qualifies every table it names, foreign key
+ * targets included — but a raw fragment is passed through verbatim, so a
+ * statement that spells its own SQL has to spell the qualifier too.
+ */
+export interface SchemaContext {
+  /**
+   * A table as the surrounding DDL addresses it: `"app"."credentials"` on a
+   * schema-bound connection, `"credentials"` otherwise.
+   *
+   * Takes the physical name, not the camelCase one the schema builder accepts.
+   * Raw SQL does not go through `CamelCasePlugin`, so it is already written in
+   * physical names throughout and this is no exception.
+   */
+  table: (name: string) => RawBuilder<unknown>
+}
+
+/**
+ * The `SchemaContext` for DDL that names no schema.
+ *
+ * The default everywhere a schema has not been asked for, which is every engine
+ * but postgres and most postgres projects too.
+ */
+export const unqualifiedContext: SchemaContext = {
+  table: (name) => sql.raw(`"${name}"`),
+}
+
+/**
+ * A `SchemaContext` qualifying every table into `schema`.
+ *
+ * Identifiers are quoted rather than interpolated bare so a schema named after
+ * a reserved word still compiles.
+ */
+export const schemaContext = (schema: string | undefined): SchemaContext =>
+  schema
+    ? { table: (name) => sql.raw(`"${schema}"."${name}"`) }
+    : unqualifiedContext
+
 /**
  * The physical types of the columns named in `requires`, keyed `table.column`.
  *
@@ -82,13 +124,26 @@ export interface SchemaRequirement {
  * Wrap a raw SQL fragment as a statement.
  *
  * For the DDL kysely's schema builder cannot express — a unique index over an
- * expression, say. Note that raw SQL is *not* rewritten by `withSchema`, so a
- * fragment naming a table unqualified will resolve against the connection's
- * search path rather than the schema the rest of the declaration lands in.
+ * expression, say. Raw SQL is *not* rewritten by `withSchema`, so a fragment
+ * naming a table unqualified resolves against the connection's search path
+ * rather than the schema the rest of the declaration lands in. Take the
+ * function form and name tables through `ctx.table` to stay with them; the
+ * bare-fragment form is for SQL that names no table at all.
  */
-export const rawStatement =
-  (fragment: RawBuilder<unknown>): SchemaStatementFactory =>
-  (db) => ({
-    execute: () => fragment.execute(db),
-    compile: () => fragment.compile(db),
-  })
+export function rawStatement(
+  fragment: RawBuilder<unknown>
+): SchemaStatementFactory
+export function rawStatement(
+  build: (ctx: SchemaContext) => RawBuilder<unknown>
+): SchemaStatementFactory
+export function rawStatement(
+  source: RawBuilder<unknown> | ((ctx: SchemaContext) => RawBuilder<unknown>)
+): SchemaStatementFactory {
+  return (db, _types, ctx = unqualifiedContext) => {
+    const fragment = typeof source === 'function' ? source(ctx) : source
+    return {
+      execute: () => fragment.execute(db),
+      compile: () => fragment.compile(db),
+    }
+  }
+}