Ekklesia/packages/module-database/src/drizzle.ts

import type {
  ModuleDatabaseExecutor,
  ModuleDatabaseParameter,
  ModuleDatabaseQueryable,
  ModuleDatabaseTransaction,
  ModuleId,
  ModuleSchemaMetadataDefinition,
} from "@ekklesia/contracts";
import { parseModuleSchemaMetadataDefinition } from "@ekklesia/contracts";
import { type AnyPgTable, getTableConfig } from "drizzle-orm/pg-core";
import { drizzle, type PgRemoteDatabase } from "drizzle-orm/pg-proxy";
import type { DrizzleConfig } from "drizzle-orm/utils";

type MaybePromise<T> = Promise<T> | T;
type DrizzleSurface<TSchema extends Record<string, unknown>> = Pick<
  PgRemoteDatabase<TSchema>,
  "delete" | "insert" | "select" | "update"
>;

export type ModuleDrizzleTransaction<
  TSchema extends Record<string, unknown> = Record<string, never>,
> = DrizzleSurface<TSchema> & {
  transaction<T>(
    callback: (transaction: ModuleDrizzleTransaction<TSchema>) => MaybePromise<T>,
  ): Promise<T>;
};

export type ModuleDrizzleDatabase<TSchema extends Record<string, unknown> = Record<string, never>> =
  DrizzleSurface<TSchema> & {
    transaction<T>(
      callback: (transaction: ModuleDrizzleTransaction<TSchema>) => MaybePromise<T>,
    ): Promise<T>;
  };

type TransactionRunner<TSchema extends Record<string, unknown>> = <T>(
  callback: (transaction: ModuleDrizzleTransaction<TSchema>) => MaybePromise<T>,
) => Promise<T>;

function remoteDatabase<TSchema extends Record<string, unknown>>(
  executor: ModuleDatabaseQueryable,
  config: DrizzleConfig<TSchema>,
  runTransaction: TransactionRunner<TSchema>,
): ModuleDrizzleDatabase<TSchema> {
  const database = drizzle<TSchema>(async (text, parameters) => {
    const result = await executor.query({
      text,
      // Drizzle performs column-specific parameter encoding before this callback. Values outside
      // ABI v1 remain unchanged so the ABI executor, not the adapter, rejects them.
      values: parameters as readonly ModuleDatabaseParameter[],
      // Drizzle maps selected fields by position. Array rows also remain correct when a join returns
      // duplicate PostgreSQL column names.
      rowMode: "array",
    });
    return { rows: result.rows as unknown as unknown[][] };
  }, config);

  return Object.freeze({
    delete: database.delete.bind(database),
    insert: database.insert.bind(database),
    select: database.select.bind(database),
    transaction: runTransaction,
    update: database.update.bind(database),
  }) as ModuleDrizzleDatabase<TSchema>;
}

function transactionDatabase<TSchema extends Record<string, unknown>>(
  transaction: ModuleDatabaseTransaction,
  config: DrizzleConfig<TSchema>,
): ModuleDrizzleTransaction<TSchema> {
  return remoteDatabase(transaction, config, (callback) =>
    transaction.savepoint((savepoint) => callback(transactionDatabase(savepoint, config))),
  );
}

/**
 * Creates a module-owned Drizzle database backed exclusively by Module Database ABI v1 calls.
 *
 * Drizzle compiles and decodes on the module side. Core receives only SQL text, ordered parameter
 * values, and `rowMode: "array"`. Top-level transactions use the ABI transaction callback; nested
 * Drizzle transactions use ABI savepoints.
 */
export function moduleDrizzle<TSchema extends Record<string, unknown> = Record<string, never>>(
  executor: ModuleDatabaseExecutor,
  config: DrizzleConfig<TSchema> = {},
): ModuleDrizzleDatabase<TSchema> {
  return remoteDatabase(executor, config, (callback) =>
    executor.transaction((transaction) => callback(transactionDatabase(transaction, config))),
  );
}

/** Derives the plain ABI-v1 schema entrypoint from a module's own Drizzle tables. Drizzle stays on
 * this module-side adapter path; Core consumes only the returned strings and arrays. */
export function moduleDrizzleSchemaMetadata(
  moduleId: ModuleId,
  tables: Readonly<Record<string, AnyPgTable>>,
): ModuleSchemaMetadataDefinition {
  const metadata = Object.entries(tables)
    .sort(([left], [right]) => (left < right ? -1 : left > right ? 1 : 0))
    .map(([_exportName, table]) => {
      const config = getTableConfig(table);
      if (config.schema !== undefined && config.schema !== "public") {
        throw new TypeError("Module schema metadata supports only the public PostgreSQL schema.");
      }
      return {
        name: config.name,
        foreignKeys: config.foreignKeys.map((foreignKey) => ({
          table: getTableConfig(foreignKey.reference().foreignTable).name,
        })),
      };
    });
  return parseModuleSchemaMetadataDefinition({ moduleId, tables: metadata });
}