STTracker/packages/db/src/dialect.ts

import { SQL, type ReservedSQL } from "bun";
import {
  CompiledQuery,
  PostgresAdapter,
  PostgresIntrospector,
  PostgresQueryCompiler,
  SqliteAdapter,
  SqliteIntrospector,
  SqliteQueryCompiler,
  type DatabaseConnection,
  type DatabaseIntrospector,
  type DialectAdapter,
  type Dialect as KyselyDialect,
  type Driver,
  type Kysely,
  type QueryCompiler,
  type QueryResult,
  type TransactionSettings,
} from "kysely";

export type Dialect = "sqlite" | "postgres";

export function dialectOf(url: string): Dialect {
  if (url.startsWith("sqlite://")) return "sqlite";
  if (url.startsWith("postgres://") || url.startsWith("postgresql://")) return "postgres";
  throw new Error(`Unsupported database URL scheme: ${url.split(":")[0]}`);
}

/** Only these commands report a meaningful affected-row count. */
const WRITE_COMMANDS = new Set(["INSERT", "UPDATE", "DELETE", "MERGE"]);

type Runner = Pick<SQL, "unsafe">;

class BunSqlConnection implements DatabaseConnection {
  constructor(
    readonly runner: Runner,
    private readonly dialect: Dialect,
  ) {}

  async executeQuery<R>(compiled: CompiledQuery): Promise<QueryResult<R>> {
    const result = await this.runner.unsafe(compiled.sql, compiled.parameters as unknown[]);
    const rows = [...result] as R[];
    const command: string | undefined = result.command;
    if (!command || !WRITE_COMMANDS.has(command)) return { rows };
    const numAffectedRows = BigInt(result.count ?? 0);
    const lastId = result.lastInsertRowid;
    if (this.dialect === "sqlite" && command === "INSERT" && lastId != null) {
      return { rows, numAffectedRows, insertId: BigInt(lastId) };
    }
    return { rows, numAffectedRows };
  }

  async *streamQuery<R>(compiled: CompiledQuery): AsyncIterableIterator<QueryResult<R>> {
    yield await this.executeQuery<R>(compiled);
  }
}

/**
 * Serialises access to the single SQLite connection. Bun.SQL's SQLite adapter
 * does not support `reserve()`, so a held connection is the whole database.
 */
class Mutex {
  private tail: Promise<void> = Promise.resolve();

  async lock(): Promise<() => void> {
    let unlock!: () => void;
    const next = new Promise<void>((resolve) => (unlock = resolve));
    const prev = this.tail;
    this.tail = prev.then(() => next);
    await prev;
    return unlock;
  }
}

class BunSqlDriver implements Driver {
  private sql: SQL | undefined;
  private readonly sqliteLock = new Mutex();
  private readonly unlocks = new WeakMap<DatabaseConnection, () => void>();

  constructor(
    private readonly url: string,
    private readonly dialect: Dialect,
  ) {}

  async init(): Promise<void> {
    this.sql = new SQL(this.url);
  }

  private get client(): SQL {
    if (!this.sql) throw new Error("Database driver is not initialised or already destroyed");
    return this.sql;
  }

  async acquireConnection(): Promise<DatabaseConnection> {
    if (this.dialect === "postgres") {
      return new BunSqlConnection(await this.client.reserve(), this.dialect);
    }
    const unlock = await this.sqliteLock.lock();
    const conn = new BunSqlConnection(this.client, this.dialect);
    this.unlocks.set(conn, unlock);
    return conn;
  }

  async releaseConnection(connection: DatabaseConnection): Promise<void> {
    if (this.dialect === "postgres") {
      ((connection as BunSqlConnection).runner as ReservedSQL).release();
      return;
    }
    this.unlocks.get(connection)?.();
    this.unlocks.delete(connection);
  }

  async beginTransaction(connection: DatabaseConnection, settings: TransactionSettings): Promise<void> {
    let stmt = "begin";
    if (this.dialect === "postgres") {
      if (settings.isolationLevel) stmt += ` isolation level ${settings.isolationLevel}`;
      if (settings.accessMode) stmt += ` ${settings.accessMode}`;
    }
    await connection.executeQuery(CompiledQuery.raw(stmt));
  }

  async commitTransaction(connection: DatabaseConnection): Promise<void> {
    await connection.executeQuery(CompiledQuery.raw("commit"));
  }

  async rollbackTransaction(connection: DatabaseConnection): Promise<void> {
    await connection.executeQuery(CompiledQuery.raw("rollback"));
  }

  async destroy(): Promise<void> {
    const sql = this.sql;
    this.sql = undefined;
    await sql?.close();
  }
}

/** One Kysely dialect over Bun.SQL for both `sqlite://` and `postgres://` URLs. */
export class BunSqlDialect implements KyselyDialect {
  readonly dialect: Dialect;

  constructor(private readonly url: string) {
    this.dialect = dialectOf(url);
  }

  createDriver(): Driver {
    return new BunSqlDriver(this.url, this.dialect);
  }

  createQueryCompiler(): QueryCompiler {
    return this.dialect === "sqlite" ? new SqliteQueryCompiler() : new PostgresQueryCompiler();
  }

  createAdapter(): DialectAdapter {
    return this.dialect === "sqlite" ? new SqliteAdapter() : new PostgresAdapter();
  }

  createIntrospector(db: Kysely<any>): DatabaseIntrospector {
    return this.dialect === "sqlite" ? new SqliteIntrospector(db) : new PostgresIntrospector(db);
  }
}