source/teaql-ts/src/sql/sqlite.ts

import Database from 'better-sqlite3';
import { startStandardNodeTelemetry } from '../telemetry/standard-node';
import {
  AbstractSQLTeaQLClient,
  assertSafeIdentifier,
  ColumnSchema,
  EntitySchema,
  LogicalColumnType,
  SqlQueryResult,
  SqlSession,
  standardAggregateFunction,
  ensureOptimisticIdFloor,
  TeaQLSqlDriver,
  canonicalRelationIndexes,
} from './core';

function soundex(value: unknown): string {
  const letters = String(value ?? '').toUpperCase().match(/[A-Z]/g) ?? [];
  if (!letters.length) return '?000';
  const groups: Record<string, string> = {};
  for (const [chars, code] of [['BFPV', '1'], ['CGJKQSXZ', '2'], ['DT', '3'], ['L', '4'], ['MN', '5'], ['R', '6']]) {
    for (const char of chars) groups[char] = code;
  }
  const first = letters[0]!;
  let result: string = first;
  let previous = groups[first] ?? '0';
  for (const letter of letters.slice(1)) {
    const current = groups[letter] ?? '0';
    if (current !== '0' && current !== previous) result += current;
    if (result.length === 4) break;
    previous = current;
  }
  return result.padEnd(4, '0');
}

export class SQLiteDriver implements TeaQLSqlDriver, SqlSession {
  readonly databaseKind = 'sqlite' as const;
  readonly topNRelationPlanPolicy = 'alwaysProbe' as const;
  private readonly database: Database.Database;
  private soundexRegistered = false;

  constructor(filename: string) {
    if (!filename) throw new Error('filename is required');
    this.database = new Database(filename);
    this.database.pragma('journal_mode = WAL');
    this.database.pragma('foreign_keys = ON');
  }

  identifier(value: string): string {
    return `"${assertSafeIdentifier(value)}"`;
  }

  placeholder(_index: number): string {
    return '?';
  }

  encode(value: any, column?: ColumnSchema): any {
    if (value === null || value === undefined) return value;
    if (column?.logicalType === 'json' && typeof value !== 'string') {
      return JSON.stringify(value);
    }
    if (column?.logicalType === 'boolean') return value ? 1 : 0;
    if (column?.logicalType === 'date') {
      return value instanceof Date ? value.toISOString().slice(0, 10) : value;
    }
    if (column?.logicalType === 'datetime') {
      if (value instanceof Date) return value.getTime();
      if (typeof value === 'string') {
        const millis = Date.parse(value);
        if (!Number.isNaN(millis)) return millis;
      }
    }
    return value;
  }

  contains(columnSql: string, placeholder: string): string {
    return `CAST(${columnSql} AS TEXT) LIKE '%' || ${placeholder} || '%'`;
  }

  aggregateFunction(name: string): string {
    return standardAggregateFunction(name);
  }

  private sqlType(type: LogicalColumnType): string {
    const types: Record<LogicalColumnType, string> = {
      boolean: 'INTEGER',
      double: 'REAL',
      decimal: 'NUMERIC',
      date: 'TEXT',
      datetime: 'INTEGER',
      json: 'TEXT',
      integer: 'INTEGER',
      text: 'TEXT',
    };
    return types[type];
  }

  async ensureSchema(schemas: Record<string, EntitySchema>): Promise<void> {
    if (!this.soundexRegistered) {
      this.database.function('soundex', { deterministic: true }, soundex);
      this.soundexRegistered = true;
    }
    for (const schema of Object.values(schemas)) {
      const table = this.identifier(schema.table);
      await this.query(
        `CREATE TABLE IF NOT EXISTS ${table} (` +
        '"id" INTEGER PRIMARY KEY, "version" INTEGER NOT NULL)',
      );
      const columns = await this.query(`PRAGMA table_info(${table})`);
      const existing = new Set(columns.rows.map(row => String(row.name)));
      for (const [field, column] of Object.entries(schema.columns)) {
        if (field === 'id' || field === 'version' ||
            existing.has(column.columnName)) continue;
        await this.query(
          `ALTER TABLE ${table} ADD COLUMN ` +
          `${this.identifier(column.columnName)} ${this.sqlType(column.logicalType)}` +
          `${column.nullable === false ? ' NOT NULL' : ''}`,
        );
      }
    }
    for (const index of canonicalRelationIndexes(schemas)) {
      await this.query(
        `CREATE INDEX IF NOT EXISTS ${this.identifier(index.name)} ON ` +
        `${this.identifier(index.table)} (` +
        `${this.identifier(index.foreignColumn)}, ${this.identifier(index.idColumn)} DESC)`,
      );
    }
    await this.query(
      'CREATE TABLE IF NOT EXISTS teaql_id_space (' +
      'type_name TEXT PRIMARY KEY, current_level INTEGER NOT NULL)',
    );
  }

  async transaction<T>(work: (session: SqlSession) => Promise<T>): Promise<T> {
    this.database.exec('BEGIN IMMEDIATE');
    try {
      const result = await work(this);
      this.database.exec('COMMIT');
      return result;
    } catch (error) {
      this.database.exec('ROLLBACK');
      throw error;
    }
  }

  async nextId(session: SqlSession, entity: string): Promise<string> {
    for (let attempt = 1; attempt <= 100; attempt += 1) {
      const current = await session.query(
        'SELECT current_level AS id FROM teaql_id_space WHERE type_name = ?',
        [entity],
      );
      if (!current.rowCount) {
        try {
          const inserted = await session.query(
            'INSERT INTO teaql_id_space(type_name, current_level) VALUES (?, 1)',
            [entity],
          );
          if (inserted.rowCount === 1) return '1';
        } catch (error) {
          const winner = await session.query(
            'SELECT current_level AS id FROM teaql_id_space WHERE type_name = ?',
            [entity],
          );
          if (!winner.rowCount) throw error;
        }
        continue;
      }
      const previous = Number(current.rows[0].id);
      const next = previous + 1;
      const updated = await session.query(
        'UPDATE teaql_id_space SET current_level = ? WHERE type_name = ? AND current_level = ?',
        [next, entity, previous],
      );
      if (updated.rowCount === 1) return String(next);
      if (updated.rowCount !== 0) {
        throw new Error(`ID space update for ${entity} changed ${updated.rowCount} rows`);
      }
    }
    throw new Error(
      `Unable to allocate ID for ${entity} after 100 optimistic-lock attempts`,
    );
  }

  ensureIdFloor(session: SqlSession, entity: string, floor: string): Promise<void> {
    return ensureOptimisticIdFloor(session, index => this.placeholder(index), entity, floor);
  }

  async query(sql: string, values: any[] = []): Promise<SqlQueryResult> {
    const statement = this.database.prepare(sql);
    if (statement.reader) {
      const rows = statement.all(...values);
      return { rows, rowCount: rows.length };
    }
    const result = statement.run(...values);
    return { rows: [], rowCount: result.changes };
  }

  async *stream(sql: string, values: any[] = []): AsyncIterable<any> {
    const statement = this.database.prepare(sql);
    if (!statement.reader) throw new Error('stream() requires a SELECT statement');
    for (const row of statement.iterate(...values)) yield row;
  }

  async close(): Promise<void> {
    this.database.close();
  }
}

export class SQLiteTeaQLClient extends AbstractSQLTeaQLClient {
  constructor(filename: string, schemas: Record<string, EntitySchema>) {
    super(new SQLiteDriver(filename), schemas);
    this.attachStandardTelemetryStartup(startStandardNodeTelemetry());
  }
}