koc-loop/lib/database.ts

import mysql, {
  type Pool,
  type PoolConnection,
  type ResultSetHeader,
  type RowDataPacket,
} from "mysql2/promise";
import { getDatabaseUrl } from "./runtime-env";

export type DatabaseResult<T = Record<string, unknown>> = {
  results: T[];
  success: boolean;
  meta: {
    changes: number;
    last_row_id: number;
  };
};

type Executor = Pool | PoolConnection;
type MysqlBindValue =
  | string
  | number
  | bigint
  | boolean
  | Date
  | null
  | Blob
  | Buffer
  | Uint8Array
  | MysqlBindValue[]
  | { [key: string]: MysqlBindValue };

export function normalizeSqlForMysql(input: string) {
  let sql = input.trim();
  sql = sql.replace(/^INSERT\s+OR\s+IGNORE\s+INTO\b/i, "INSERT IGNORE INTO");
  sql = sql.replace(/\s+ESCAPE\s+'\\\\'/gi, "");
  sql = sql.replace(
    /datetime\(\s*'now'\s*,\s*'-([0-9]+)\s+minutes?'\s*\)/gi,
    "DATE_SUB(UTC_TIMESTAMP(), INTERVAL $1 MINUTE)",
  );
  sql = sql.replace(
    /datetime\(\s*([^,]+?)\s*,\s*'\+([0-9]+)\s+days?'\s*\)/gi,
    "DATE_ADD($1, INTERVAL $2 DAY)",
  );

  const conflict = sql.match(
    /\s+ON\s+CONFLICT\s*\(([^)]+)\)\s+DO\s+UPDATE\s+SET\s+([\s\S]+)$/i,
  );
  if (conflict) {
    const assignments = conflict[2].replace(
      /\bexcluded\.([a-zA-Z_][a-zA-Z0-9_]*)\b/g,
      "VALUES($1)",
    );
    sql = `${sql.slice(0, conflict.index)} ON DUPLICATE KEY UPDATE ${assignments}`;
  }
  return sql;
}

function normalizeBindValue(value: unknown): MysqlBindValue {
  if (value === undefined) return null;
  if (
    typeof value === "string" &&
    /^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d{1,3})?Z$/.test(value)
  ) {
    return value.replace("T", " ").replace(/Z$/, "");
  }
  return value as MysqlBindValue;
}

export class DatabaseStatement {
  private params: unknown[] = [];
  readonly database: DatabaseClient;
  readonly sql: string;

  constructor(database: DatabaseClient, sql: string) {
    this.database = database;
    this.sql = sql;
  }

  bind(...params: unknown[]) {
    const statement = new DatabaseStatement(this.database, this.sql);
    statement.params = params;
    return statement;
  }

  async all<T = Record<string, unknown>>() {
    return this.database.execute<T>(this.sql, this.params);
  }

  async first<T = Record<string, unknown>>(column?: string) {
    const result = await this.all<T>();
    const row = result.results[0];
    if (!row) return null;
    return column ? ((row as Record<string, unknown>)[column] as T) : row;
  }

  async run() {
    return this.database.execute(this.sql, this.params);
  }

  values() {
    return [...this.params];
  }
}

export class DatabaseClient {
  private readonly executor: Executor;

  constructor(executor: Executor) {
    this.executor = executor;
  }

  prepare(sql: string) {
    return new DatabaseStatement(this, sql);
  }

  async execute<T = Record<string, unknown>>(sql: string, params: unknown[] = []) {
    const [result] = await this.executor.execute(
      normalizeSqlForMysql(sql),
      params.map(normalizeBindValue),
    );
    if (Array.isArray(result)) {
      return {
        results: result as T[],
        success: true,
        meta: { changes: 0, last_row_id: 0 },
      } satisfies DatabaseResult<T>;
    }
    const header = result as ResultSetHeader;
    return {
      results: [],
      success: true,
      meta: {
        changes: header.affectedRows ?? 0,
        last_row_id: header.insertId ?? 0,
      },
    } satisfies DatabaseResult<T>;
  }

  async batch(statements: DatabaseStatement[]) {
    const pool = getPool();
    const connection = await pool.getConnection();
    try {
      await connection.beginTransaction();
      const tx = new DatabaseClient(connection);
      const results = [];
      for (const statement of statements) {
        results.push(await tx.execute(statement.sql, statement.values()));
      }
      await connection.commit();
      return results;
    } catch (error) {
      await connection.rollback();
      throw error;
    } finally {
      connection.release();
    }
  }
}

declare global {
  var __kocLoopMysqlPool: Pool | undefined;
}

export function getPool() {
  if (!globalThis.__kocLoopMysqlPool) {
    globalThis.__kocLoopMysqlPool = mysql.createPool({
      uri: getDatabaseUrl(),
      connectionLimit: 10,
      waitForConnections: true,
      queueLimit: 0,
      charset: "utf8mb4",
      timezone: "Z",
      dateStrings: true,
      decimalNumbers: true,
      enableKeepAlive: true,
    });
  }
  return globalThis.__kocLoopMysqlPool;
}

export function getDatabase() {
  return new DatabaseClient(getPool());
}

export async function checkDatabaseConnection() {
  const [rows] = await getPool().query<RowDataPacket[]>("SELECT 1 AS healthy");
  return rows[0]?.healthy === 1;
}

export async function withDatabaseLock<T>(
  name: string,
  timeoutSeconds: number,
  operation: () => Promise<T>,
) {
  const connection = await getPool().getConnection();
  try {
    const [rows] = await connection.query<RowDataPacket[]>(
      "SELECT GET_LOCK(?, ?) AS acquired",
      [name, timeoutSeconds],
    );
    if (Number(rows[0]?.acquired ?? 0) !== 1) return null;
    try {
      return await operation();
    } finally {
      await connection.query("SELECT RELEASE_LOCK(?)", [name]);
    }
  } finally {
    connection.release();
  }
}