Ekklesia/packages/module-database/tests/drizzle.test.ts

import {
  type ModuleDatabaseArrayRow,
  ModuleDatabaseError,
  type ModuleDatabaseExecutor,
  type ModuleDatabaseField,
  type ModuleDatabaseObjectRow,
  type ModuleDatabaseQuery,
  type ModuleDatabaseQueryResult,
  type ModuleDatabaseTransaction,
} from "@ekklesia/contracts";
import { and, sql as drizzleSql, eq } from "drizzle-orm";
import {
  bigint,
  customType,
  date,
  integer,
  json,
  jsonb,
  numeric,
  pgTable,
  text,
  timestamp,
} from "drizzle-orm/pg-core";
import { describe, expect, expectTypeOf, it } from "vitest";
import {
  type ModuleDrizzleDatabase,
  moduleDrizzle,
  moduleDrizzleSchemaMetadata,
} from "../src/drizzle.js";
import { findEkklesiaDatabaseError } from "../src/index.js";

const bytea = customType<{ data: Uint8Array; driverData: Uint8Array }>({
  dataType: () => "bytea",
});

const records = pgTable("calendar_records", {
  amount: numeric("amount"),
  bigId: bigint("big_id", { mode: "bigint" }).notNull(),
  bytes: bytea("bytes"),
  count: integer("count").notNull(),
  createdAt: timestamp("created_at", { mode: "date", withTimezone: true }).notNull(),
  day: date("day", { mode: "date" }).notNull(),
  id: integer("id").primaryKey(),
  localAt: timestamp("local_at", { mode: "date" }).notNull(),
  metadata: jsonb("metadata"),
  nullable: text("nullable"),
  payload: json("payload"),
  scores: bigint("scores", { mode: "bigint" }).array(),
  tags: text("tags").array(),
  title: text("title").notNull(),
});

const categories = pgTable("calendar_categories", {
  id: integer("id").primaryKey(),
  title: text("title").notNull(),
});

type Outcome = Error | ModuleDatabaseQueryResult<ModuleDatabaseArrayRow | ModuleDatabaseObjectRow>;

interface RecordedCall {
  readonly query: ModuleDatabaseQuery;
  readonly scope: string;
}

const noFields: readonly ModuleDatabaseField[] = [];

function result(
  rows: readonly ModuleDatabaseArrayRow[],
  command = "SELECT",
): ModuleDatabaseQueryResult<ModuleDatabaseArrayRow> {
  return { command, fields: noFields, rowCount: rows.length, rows };
}

class FakeState {
  readonly calls: RecordedCall[] = [];
  readonly events: string[] = [];
  readonly outcomes: Outcome[] = [];
}

class FakeTransaction implements ModuleDatabaseTransaction {
  readonly abiVersion = 1;

  constructor(
    private readonly state: FakeState,
    private readonly scope: string,
  ) {}

  query(
    query: ModuleDatabaseQuery & { readonly rowMode: "array" },
  ): Promise<ModuleDatabaseQueryResult<ModuleDatabaseArrayRow>>;
  query(
    query: ModuleDatabaseQuery & { readonly rowMode?: "object" },
  ): Promise<ModuleDatabaseQueryResult<ModuleDatabaseObjectRow>>;
  query(
    query: ModuleDatabaseQuery,
  ): Promise<ModuleDatabaseQueryResult<ModuleDatabaseObjectRow | ModuleDatabaseArrayRow>>;
  async query(
    query: ModuleDatabaseQuery,
  ): Promise<ModuleDatabaseQueryResult<ModuleDatabaseObjectRow | ModuleDatabaseArrayRow>> {
    this.state.calls.push({ query, scope: this.scope });
    this.state.events.push(`${this.scope}:query`);
    const outcome = this.state.outcomes.shift() ?? result([]);
    if (outcome instanceof Error) {
      throw outcome;
    }
    return outcome;
  }

  async savepoint<T>(
    callback: (savepoint: ModuleDatabaseTransaction) => Promise<T> | T,
  ): Promise<T> {
    const scope = `${this.scope}.savepoint`;
    this.state.events.push(`${scope}:begin`);
    try {
      const value = await callback(new FakeTransaction(this.state, scope));
      this.state.events.push(`${scope}:release`);
      return value;
    } catch (error) {
      this.state.events.push(`${scope}:rollback`);
      throw error;
    }
  }
}

class FakeExecutor implements ModuleDatabaseExecutor {
  readonly abiVersion = 1;
  readonly state = new FakeState();

  query(
    query: ModuleDatabaseQuery & { readonly rowMode: "array" },
  ): Promise<ModuleDatabaseQueryResult<ModuleDatabaseArrayRow>>;
  query(
    query: ModuleDatabaseQuery & { readonly rowMode?: "object" },
  ): Promise<ModuleDatabaseQueryResult<ModuleDatabaseObjectRow>>;
  query(
    query: ModuleDatabaseQuery,
  ): Promise<ModuleDatabaseQueryResult<ModuleDatabaseObjectRow | ModuleDatabaseArrayRow>>;
  async query(
    query: ModuleDatabaseQuery,
  ): Promise<ModuleDatabaseQueryResult<ModuleDatabaseObjectRow | ModuleDatabaseArrayRow>> {
    return new FakeTransaction(this.state, "executor").query(query);
  }

  async transaction<T>(
    callback: (transaction: ModuleDatabaseTransaction) => Promise<T> | T,
  ): Promise<T> {
    this.state.events.push("transaction:begin");
    try {
      const value = await callback(new FakeTransaction(this.state, "transaction"));
      this.state.events.push("transaction:commit");
      return value;
    } catch (error) {
      this.state.events.push("transaction:rollback");
      throw error;
    }
  }
}

function setup(...outcomes: Outcome[]) {
  const executor = new FakeExecutor();
  executor.state.outcomes.push(...outcomes);
  return { db: moduleDrizzle(executor), executor };
}

describe("moduleDrizzle query compilation and ABI mapping", () => {
  it("derives plain Ekklesia schema metadata with foreign-key targets", () => {
    const child = pgTable("calendar_children", {
      id: integer("id").primaryKey(),
      recordId: integer("record_id").references(() => records.id),
    });

    const metadata = moduleDrizzleSchemaMetadata("calendar", { records, child });

    expect(metadata).toEqual({
      moduleId: "calendar",
      tables: [
        { name: "calendar_children", foreignKeys: [{ table: "calendar_records" }] },
        { name: "calendar_records", foreignKeys: [] },
      ],
    });
    expect(JSON.parse(JSON.stringify(metadata))).toEqual(metadata);
  });

  it("compiles a select and decodes its positional result on the module side", async () => {
    const { db, executor } = setup(result([[7, "Launch"]]));
    await expect(
      db.select({ id: records.id, title: records.title }).from(records).where(eq(records.id, 7)),
    ).resolves.toEqual([{ id: 7, title: "Launch" }]);
    expect(executor.state.calls).toEqual([
      {
        query: {
          text: 'select "id", "title" from "calendar_records" where "calendar_records"."id" = $1',
          values: [7],
          rowMode: "array",
        },
        scope: "executor",
      },
    ]);
  });

  it("compiles insert returning through the array-row ABI", async () => {
    const { db, executor } = setup(result([[8, "New"]], "INSERT"));
    await expect(
      db
        .insert(records)
        .values({
          bigId: 8n,
          count: 1,
          createdAt: new Date("2026-01-02T03:04:05.000Z"),
          day: new Date("2026-01-02T00:00:00.000Z"),
          id: 8,
          localAt: new Date("2026-01-02T03:04:05.000Z"),
          title: "New",
        })
        .returning({ id: records.id, title: records.title }),
    ).resolves.toEqual([{ id: 8, title: "New" }]);
    expect(executor.state.calls[0]?.query).toMatchObject({
      rowMode: "array",
      values: [
        8n,
        1,
        "2026-01-02T03:04:05.000Z",
        "2026-01-02T00:00:00.000Z",
        8,
        "2026-01-02T03:04:05.000Z",
        "New",
      ],
    });
    expect(executor.state.calls[0]?.query.text).toContain('insert into "calendar_records"');
    expect(executor.state.calls[0]?.query.text).toContain('returning "id", "title"');
  });

  it("compiles an update returning with a SQL expression", async () => {
    const { db, executor } = setup(result([[7, 4]], "UPDATE"));
    await expect(
      db
        .update(records)
        .set({ count: drizzleSql`${records.count} + 1` })
        .where(and(eq(records.id, 7), eq(records.title, "Launch")))
        .returning({ count: records.count, id: records.id }),
    ).resolves.toEqual([{ count: 7, id: 4 }]);
    expect(executor.state.calls[0]?.query).toEqual({
      text: 'update "calendar_records" set "count" = "calendar_records"."count" + 1 where ("calendar_records"."id" = $1 and "calendar_records"."title" = $2) returning "count", "id"',
      values: [7, "Launch"],
      rowMode: "array",
    });
  });

  it("compiles delete returning", async () => {
    const { db, executor } = setup(result([[7, "Old"]], "DELETE"));
    await expect(
      db
        .delete(records)
        .where(eq(records.id, 7))
        .returning({ id: records.id, title: records.title }),
    ).resolves.toEqual([{ id: 7, title: "Old" }]);
    expect(executor.state.calls[0]?.query).toEqual({
      text: 'delete from "calendar_records" where "calendar_records"."id" = $1 returning "id", "title"',
      values: [7],
      rowMode: "array",
    });
  });

  it("uses positions rather than duplicate database column names for joins", async () => {
    const { db, executor } = setup(result([[9, "Category", 3, "Record"]]));
    await expect(
      db
        .select({
          categoryId: categories.id,
          categoryTitle: categories.title,
          recordId: records.id,
          recordTitle: records.title,
        })
        .from(records)
        .innerJoin(categories, eq(records.id, categories.id)),
    ).resolves.toEqual([
      { categoryId: 9, categoryTitle: "Category", recordId: 3, recordTitle: "Record" },
    ]);
    expect(executor.state.calls[0]?.query.rowMode).toBe("array");
    expect(executor.state.calls[0]?.query.text).toContain(
      'select "calendar_categories"."id", "calendar_categories"."title", "calendar_records"."id", "calendar_records"."title"',
    );
  });
});

describe("moduleDrizzle result decoding", () => {
  it("lets Drizzle decode every stable low-level ABI representation", async () => {
    const bytes = new Uint8Array([0, 127, 255]);
    const { db, executor } = setup(
      result([
        [
          "1234567890.0123456789",
          "12345678901234567890",
          bytes,
          "2026-01-02 03:04:05+00",
          "2026-01-02",
          "2026-01-02 03:04:05",
          '{"kind":"jsonb"}',
          null,
          '{"kind":"json"}',
          ["3", "9007199254740993"],
          ["one", "two"],
        ],
      ]),
    );

    const rows = await db
      .select({
        amount: records.amount,
        bigId: records.bigId,
        bytes: records.bytes,
        createdAt: records.createdAt,
        day: records.day,
        localAt: records.localAt,
        metadata: records.metadata,
        nullable: records.nullable,
        payload: records.payload,
        scores: records.scores,
        tags: records.tags,
      })
      .from(records);

    expect(rows).toEqual([
      {
        amount: "1234567890.0123456789",
        bigId: 12_345_678_901_234_567_890n,
        bytes,
        createdAt: new Date("2026-01-02T03:04:05.000Z"),
        day: new Date("2026-01-02T00:00:00.000Z"),
        localAt: new Date("2026-01-02T03:04:05.000Z"),
        metadata: { kind: "jsonb" },
        nullable: null,
        payload: { kind: "json" },
        scores: [3n, 9_007_199_254_740_993n],
        tags: ["one", "two"],
      },
    ]);
    expect(rows[0]?.bytes).toBe(bytes);
    expect(executor.state.calls[0]?.query.rowMode).toBe("array");
  });
});

describe("moduleDrizzle transactions", () => {
  it("maps top-level and nested transactions to ABI callbacks and savepoints", async () => {
    const { db, executor } = setup(result([[1]]), result([[2]]), result([[3]]));
    const value = await db.transaction(async (transaction) => {
      await transaction.select({ id: records.id }).from(records);
      const nested = await transaction.transaction(async (savepoint) => {
        await savepoint.select({ id: records.id }).from(records);
        return "nested";
      });
      await transaction.select({ id: records.id }).from(records);
      return `${nested}-result`;
    });

    expect(value).toBe("nested-result");
    expect(executor.state.events).toEqual([
      "transaction:begin",
      "transaction:query",
      "transaction.savepoint:begin",
      "transaction.savepoint:query",
      "transaction.savepoint:release",
      "transaction:query",
      "transaction:commit",
    ]);
    expect(executor.state.calls.map((call) => call.scope)).toEqual([
      "transaction",
      "transaction.savepoint",
      "transaction",
    ]);
  });

  it("propagates callback errors unchanged and leaves rollback to the ABI", async () => {
    const { db, executor } = setup();
    const failure = new Error("module callback failed");
    await expect(
      db.transaction(() => {
        throw failure;
      }),
    ).rejects.toBe(failure);
    expect(executor.state.events).toEqual(["transaction:begin", "transaction:rollback"]);
  });

  it("rolls a rejected nested callback back to its ABI savepoint", async () => {
    const { db, executor } = setup();
    const failure = new Error("nested callback failed");
    await expect(
      db.transaction((transaction) => transaction.transaction(() => Promise.reject(failure))),
    ).rejects.toBe(failure);
    expect(executor.state.events).toEqual([
      "transaction:begin",
      "transaction.savepoint:begin",
      "transaction.savepoint:rollback",
      "transaction:rollback",
    ]);
  });
});

describe("moduleDrizzle error and type boundaries", () => {
  it("recovers the stable ABI code and SQLSTATE from Drizzle's wrapper", async () => {
    const databaseError = new ModuleDatabaseError("MODULE_DATABASE_QUERY_FAILED", {
      sqlState: "23505",
    });
    const { db } = setup(databaseError);
    let thrown: unknown;
    try {
      await db.select({ id: records.id }).from(records);
    } catch (error) {
      thrown = error;
    }
    expect(thrown).toBeInstanceOf(Error);
    expect(findEkklesiaDatabaseError(thrown)).toBe(databaseError);
    expect(findEkklesiaDatabaseError(thrown)).toMatchObject({
      code: "MODULE_DATABASE_QUERY_FAILED",
      sqlState: "23505",
    });
    expect(databaseError).not.toHaveProperty("query");
    expect(databaseError).not.toHaveProperty("params");
    expect(databaseError).not.toHaveProperty("detail");
    expect(databaseError).not.toHaveProperty("hint");
  });

  it("retains the module's schema inference without Core or pg database types", () => {
    const executor = new FakeExecutor();
    const schema = { records };
    const db = moduleDrizzle(executor, { schema });
    expectTypeOf(db).toMatchTypeOf<ModuleDrizzleDatabase<typeof schema>>();
    expectTypeOf(db.select().from(records)).resolves.toEqualTypeOf<
      (typeof records.$inferSelect)[]
    >();
    expect(Object.keys(db).sort()).toEqual(["delete", "insert", "select", "transaction", "update"]);
    expect(db).not.toHaveProperty("_");
    expect(db).not.toHaveProperty("$client");
    expect(db).not.toHaveProperty("execute");
    expect(db).not.toHaveProperty("query");
  });
});