MetonaSqlark/src/engine/memory.ts

/**
 * metona-sqlark Memory Engine — 基于 Map 的内存存储引擎
 * @module engine/memory
 */

import type { IStorageEngine } from './interface';
import type { QueryPlan, TableSchema } from '../constants';
import { DatabaseError } from '../constants';
import { matchWhere, applyOrderBy, projectColumns } from '../query/where-matcher';

export class MemoryEngine implements IStorageEngine {
  readonly name = 'memory';

  private tables: Map<string, Map<string, Record<string, unknown>>> = new Map();
  private schemas: Map<string, TableSchema> = new Map();
  private indexes: Map<string, Map<string, Map<unknown, Set<string>>>> = new Map();
  private opened = false;

  // ---- 事务快照 ----
  private snapshot: {
    tables: Map<string, Map<string, Record<string, unknown>>>;
    schemas: Map<string, TableSchema>;
    indexes: Map<string, Map<string, Map<unknown, Set<string>>>>;
  } | null = null;

  // ---- 生命周期 ----
  async open(_dbName: string, _version: number): Promise<void> {
    if (this.opened) {
      // 幂等：已打开则忽略
      return;
    }
    this.opened = true;
  }
  async close(): Promise<void> {
    this.tables.clear(); this.schemas.clear(); this.indexes.clear(); this.opened = false;
  }
  isOpen(): boolean { return this.opened; }

  // ---- 表管理 ----
  async createTable(schema: TableSchema): Promise<void> {
    if (this.schemas.has(schema.name)) throw new DatabaseError(`Table "${schema.name}" already exists`, 'TABLE_EXISTS');
    this.schemas.set(schema.name, schema);
    this.tables.set(schema.name, new Map());
    const tableIndexes = new Map<string, Map<unknown, Set<string>>>();
    for (const [colName, colDef] of Object.entries(schema.columns)) {
      if (colDef.index || colDef.unique) tableIndexes.set(colName, new Map());
    }
    this.indexes.set(schema.name, tableIndexes);
  }

  async dropTable(tableName: string): Promise<void> {
    this.ensureTable(tableName);
    this.schemas.delete(tableName); this.tables.delete(tableName); this.indexes.delete(tableName);
  }

  async hasTable(tableName: string): Promise<boolean> { return this.schemas.has(tableName); }
  async getTableNames(): Promise<string[]> { return Array.from(this.schemas.keys()); }
  async getTableSchema(tableName: string): Promise<TableSchema | null> { return this.schemas.get(tableName) ?? null; }

  // ---- CRUD ----
  async insert(tableName: string, rows: Record<string, unknown>[]): Promise<string[]> {
    this.ensureTable(tableName);
    const schema = this.schemas.get(tableName)!;
    const table = this.tables.get(tableName)!;
    const pkColumn = this.getPrimaryKey(schema);
    const pks: string[] = [];

    for (const row of rows) {
      const validatedRow = this.validateRow(schema, row);
      const pkValue = String(validatedRow[pkColumn]);
      if (table.has(pkValue)) throw new DatabaseError(`Duplicate primary key "${pkValue}" in table "${tableName}"`, 'DUPLICATE_KEY');
      this.checkUniqueness(schema, validatedRow);
      table.set(pkValue, validatedRow);
      this.updateIndexes(tableName, validatedRow, pkValue);
      pks.push(pkValue);
    }
    return pks;
  }

  async find(tableName: string, query: QueryPlan): Promise<Record<string, unknown>[]> {
    this.ensureTable(tableName);
    const table = this.tables.get(tableName)!;
    let results = this.tryIndexLookup(tableName, table, query);

    if (query.where && Object.keys(query.where).length > 0) {
      results = results.filter((row) => matchWhere(row, query.where!));
    }
    if (query.orderBy && query.orderBy.length > 0) {
      results = applyOrderBy(results, query.orderBy);
    }
    const offset = query.offset ?? 0;
    const limit = query.limit ?? results.length;
    results = results.slice(offset, offset + limit);
    if (query.columns && query.columns.length > 0 && query.columns[0] !== '*') {
      results = results.map((row) => projectColumns(row, query.columns!));
    }
    return results;
  }

  async update(tableName: string, query: QueryPlan, updates: Record<string, unknown>): Promise<number> {
    this.ensureTable(tableName);
    const schema = this.schemas.get(tableName)!;
    const table = this.tables.get(tableName)!;
    let count = 0;
    for (const [pk, row] of table) {
      if (!query.where || Object.keys(query.where).length === 0 || matchWhere(row, query.where)) {
        const updated = { ...row, ...updates };
        this.validateRow(schema, updated);
        table.set(pk, updated);
        count++;
      }
    }
    return count;
  }

  async delete(tableName: string, query: QueryPlan): Promise<number> {
    this.ensureTable(tableName);
    const table = this.tables.get(tableName)!;
    const toDelete: string[] = [];
    for (const [pk, row] of table) {
      if (!query.where || Object.keys(query.where).length === 0 || matchWhere(row, query.where)) {
        toDelete.push(pk);
      }
    }
    // 级联删除：检查引用此表的其他表
    let cascadeCount = 0;
    for (const pk of toDelete) {
      const row = table.get(pk);
      if (row) cascadeCount += await this.cascadeDelete(tableName, pk, row);
    }
    for (const pk of toDelete) table.delete(pk);
    return toDelete.length + cascadeCount;
  }

  async count(tableName: string, query?: QueryPlan): Promise<number> {
    this.ensureTable(tableName);
    const table = this.tables.get(tableName)!;
    if (!query?.where || Object.keys(query.where).length === 0) return table.size;
    let result = 0;
    for (const row of table.values()) { if (matchWhere(row, query.where)) result++; }
    return result;
  }

  async clear(tableName: string): Promise<void> {
    this.ensureTable(tableName);
    this.tables.get(tableName)!.clear();
    const tableIndexes = this.indexes.get(tableName);
    if (tableIndexes) for (const colIndex of tableIndexes.values()) colIndex.clear();
  }

  // ---- 动态索引（v0.3.0） ----

  async createIndex(tableName: string, column: string, unique?: boolean): Promise<void> {
    this.ensureTable(tableName);
    const schema = this.schemas.get(tableName)!;
    const colDef = schema.columns[column];
    if (!colDef) throw new DatabaseError(`Column "${column}" does not exist in table "${tableName}"`, 'COLUMN_NOT_FOUND');
    if (colDef.index || colDef.unique) return; // 已存在
    colDef.index = true;
    if (unique) colDef.unique = true;

    const tableIndexes = this.indexes.get(tableName)!;
    if (!tableIndexes.has(column)) tableIndexes.set(column, new Map());
    const colIndex = tableIndexes.get(column)!;
    const table = this.tables.get(tableName)!;
    for (const [pk, row] of table) {
      const value = row[column];
      if (value !== undefined && value !== null) {
        if (!colIndex.has(value)) colIndex.set(value, new Set());
        colIndex.get(value)!.add(pk);
      }
    }
  }

  async dropIndex(tableName: string, column: string, _indexName?: string): Promise<void> {
    this.ensureTable(tableName);
    const schema = this.schemas.get(tableName)!;
    const colDef = schema.columns[column];
    if (!colDef) throw new DatabaseError(`Column "${column}" does not exist in table "${tableName}"`, 'COLUMN_NOT_FOUND');
    colDef.index = false;
    colDef.unique = false;
    const tableIndexes = this.indexes.get(tableName);
    if (tableIndexes) tableIndexes.delete(column);
  }

  // ---- 事务 ----

  async beginTransaction(): Promise<void> {
    if (this.snapshot) throw new DatabaseError('Transaction already in progress', 'TX_ACTIVE');
    this.snapshot = {
      tables: this.deepCloneMapMap(this.tables),
      schemas: new Map(this.schemas),
      indexes: this.deepCloneIndexes(this.indexes),
    };
  }

  async commitTransaction(): Promise<void> {
    if (!this.snapshot) throw new DatabaseError('No active transaction', 'TX_NONE');
    this.snapshot = null;
  }

  async rollbackTransaction(): Promise<void> {
    if (!this.snapshot) throw new DatabaseError('No active transaction', 'TX_NONE');
    this.tables = this.snapshot.tables;
    this.schemas = this.snapshot.schemas;
    this.indexes = this.snapshot.indexes;
    this.snapshot = null;
  }

  // ---- 事务快照辅助 ----

  private deepCloneMapMap(source: Map<string, Map<string, Record<string, unknown>>>): Map<string, Map<string, Record<string, unknown>>> {
    const clone = new Map<string, Map<string, Record<string, unknown>>>();
    for (const [k, v] of source) {
      const innerClone = new Map<string, Record<string, unknown>>();
      for (const [ik, iv] of v) innerClone.set(ik, { ...iv });
      clone.set(k, innerClone);
    }
    return clone;
  }

  private deepCloneIndexes(source: Map<string, Map<string, Map<unknown, Set<string>>>>): Map<string, Map<string, Map<unknown, Set<string>>>> {
    const clone = new Map<string, Map<string, Map<unknown, Set<string>>>>();
    for (const [tableName, tableIndexes] of source) {
      const tableClone = new Map<string, Map<unknown, Set<string>>>();
      for (const [col, colIndex] of tableIndexes) {
        const colClone = new Map<unknown, Set<string>>();
        for (const [val, pkSet] of colIndex) colClone.set(val, new Set(pkSet));
        tableClone.set(col, colClone);
      }
      clone.set(tableName, tableClone);
    }
    return clone;
  }

  // ---- 内部辅助 ----
  private ensureTable(tableName: string): void {
    if (!this.tables.has(tableName)) throw new DatabaseError(`Table "${tableName}" does not exist`, 'TABLE_NOT_FOUND');
  }

  private getPrimaryKey(schema: TableSchema): string {
    for (const [name, col] of Object.entries(schema.columns)) { if (col.primaryKey) return name; }
    return Object.keys(schema.columns)[0];
  }

  private validateRow(schema: TableSchema, row: Record<string, unknown>): Record<string, unknown> {
    const validated: Record<string, unknown> = {};
    for (const [colName, colDef] of Object.entries(schema.columns)) {
      let value = row[colName];
      if (value === undefined && colDef.default !== undefined) value = colDef.default;
      if (colDef.required && (value === undefined || value === null)) {
        throw new DatabaseError(`Column "${colName}" is required in table "${schema.name}"`, 'VALIDATION_ERROR');
      }
      if (value !== undefined && value !== null) this.checkType(colName, colDef.type, value);
      if (value !== undefined) validated[colName] = value;
    }
    return validated;
  }

  private checkType(colName: string, type: string, value: unknown): void {
    const jsType = typeof value;
    switch (type) {
      case 'string': if (jsType !== 'string') throw new DatabaseError(`Column "${colName}" expects string, got ${jsType}`, 'TYPE_ERROR'); break;
      case 'number': if (jsType !== 'number') throw new DatabaseError(`Column "${colName}" expects number, got ${jsType}`, 'TYPE_ERROR'); break;
      case 'boolean': if (jsType !== 'boolean') throw new DatabaseError(`Column "${colName}" expects boolean, got ${jsType}`, 'TYPE_ERROR'); break;
      case 'date': if (jsType !== 'string' || isNaN(Date.parse(value as string))) throw new DatabaseError(`Column "${colName}" expects valid date`, 'TYPE_ERROR'); break;
      case 'json': if (jsType !== 'object') throw new DatabaseError(`Column "${colName}" expects object/array, got ${jsType}`, 'TYPE_ERROR'); break;
    }
  }

  /** O(1) 唯一性检查：利用哈希索引 */
  private checkUniqueness(schema: TableSchema, row: Record<string, unknown>): void {
    const tableIndexes = this.indexes.get(schema.name);
    if (!tableIndexes) return;
    for (const [colName, colDef] of Object.entries(schema.columns)) {
      if (!colDef.unique || row[colName] === undefined || row[colName] === null) continue;
      const colIndex = tableIndexes.get(colName);
      if (colIndex && colIndex.has(row[colName])) {
        throw new DatabaseError(`Unique constraint violation on column "${colName}" in table "${schema.name}"`, 'UNIQUE_VIOLATION');
      }
    }
  }

  /** 索引查找 */
  private tryIndexLookup(
    tableName: string, table: Map<string, Record<string, unknown>>, query: QueryPlan,
  ): Record<string, unknown>[] {
    const tableIndexes = this.indexes.get(tableName);
    if (!tableIndexes || !query.where) return Array.from(table.values());
    for (const [col, condition] of Object.entries(query.where)) {
      if (typeof condition !== 'object' || condition === null) {
        const colIndex = tableIndexes.get(col);
        if (colIndex) {
          const pks = colIndex.get(condition);
          if (pks) {
            const result: Record<string, unknown>[] = [];
            for (const pk of pks) { const r = table.get(pk); if (r) result.push(r); }
            return result;
          }
          return [];
        }
      }
    }
    return Array.from(table.values());
  }

  /** 更新索引 */
  private updateIndexes(tableName: string, row: Record<string, unknown>, pk: string): void {
    const tableIndexes = this.indexes.get(tableName);
    if (!tableIndexes) return;
    for (const [colName, colIndex] of tableIndexes) {
      const value = row[colName];
      if (value !== undefined && value !== null) {
        if (!colIndex.has(value)) colIndex.set(value, new Set());
        colIndex.get(value)!.add(pk);
      }
    }
  }

  // ---- 外键级联 ----

  /**
   * 级联删除：查找引用 tableName.pkValue 的所有表的行并删除。
   * @returns 级联删除的行数
   */
  private async cascadeDelete(tableName: string, pkValue: string, _row: Record<string, unknown>): Promise<number> {
    let totalCascade = 0;

    for (const [refTableName, refSchema] of this.schemas) {
      if (refTableName === tableName) continue;

      for (const [colName, colDef] of Object.entries(refSchema.columns)) {
        if (!colDef.references || !colDef.onDelete) continue;

        const [refTable] = colDef.references.split('.');
        if (refTable !== tableName) continue;

        const refTableData = this.tables.get(refTableName);
        if (!refTableData) continue;

        // 查找所有引用此主键的行
        const toDelete: string[] = [];
        for (const [refPk, refRow] of refTableData) {
          if (String(refRow[colName]) === pkValue) {
            toDelete.push(refPk);
          }
        }

        // RESTRICT: 存在引用行时禁止删除
        if (colDef.onDelete === 'RESTRICT' && toDelete.length > 0) {
          throw new DatabaseError(
            `Cannot delete from "${tableName}": foreign key "${colName}" in "${refTableName}" has dependent rows`,
            'FOREIGN_KEY_VIOLATION',
          );
        }

        if (colDef.onDelete === 'CASCADE') {
          // 递归级联
          for (const refPk of toDelete) {
            const refRow = refTableData.get(refPk);
            if (refRow) {
              totalCascade += await this.cascadeDelete(refTableName, refPk, refRow);
            }
            refTableData.delete(refPk);
            totalCascade++;
          }
        } else if (colDef.onDelete === 'SET NULL') {
          for (const refPk of toDelete) {
            const refRow = refTableData.get(refPk);
            if (refRow) {
              refRow[colName] = null;
            }
          }
        }
      }
    }

    return totalCascade;
  }
}