sql-talk/src/database/schema-cache.ts

SQL Talk - 自然言語をSQLに変換するMCPサーバー（安全性保護・SSHトンネル対応） / SQL Talk - MCP Server for Natural Language to SQL conversion with safety guards and SSH tunnel support

import { readFileSync, writeFileSync, existsSync, mkdirSync } from 'fs';
import { dirname, resolve } from 'path';
import { SchemaCache, TableInfo, ColumnInfo, DatabaseEngine } from '@/types/index.js';
import { connectionManager } from './connection.js';
import { logger, PerformanceLogger } from '@/core/logger.js';
import { SchemaError } from '@/core/errors.js';
import { configManager } from '@/core/config.js';

interface PostgreSQLTableRow {
  table_schema: string;
  table_name: string;
  table_comment: string | null;
}

interface PostgreSQLColumnRow {
  table_schema: string;
  table_name: string;
  column_name: string;
  data_type: string;
  is_nullable: string;
  column_default: string | null;
  column_comment: string | null;
}

interface MySQLTableRow {
  TABLE_SCHEMA: string;
  TABLE_NAME: string;
  TABLE_COMMENT: string | null;
}

interface MySQLColumnRow {
  TABLE_SCHEMA: string;
  TABLE_NAME: string;
  COLUMN_NAME: string;
  DATA_TYPE: string;
  IS_NULLABLE: string;
  COLUMN_DEFAULT: string | null;
  COLUMN_COMMENT: string | null;
}

export class SchemaCacheManager {
  private static instance: SchemaCacheManager;
  private cache: SchemaCache | null = null;
  private cachePath: string;

  private constructor() {
    this.cachePath = resolve(process.cwd(), '.cache/schema_cache.json');
  }

  public static getInstance(): SchemaCacheManager {
    if (!SchemaCacheManager.instance) {
      SchemaCacheManager.instance = new SchemaCacheManager();
    }
    return SchemaCacheManager.instance;
  }

  public async initialize(): Promise<void> {
    const config = configManager.getConfig();
    this.cachePath = resolve(process.cwd(), config.schema_cache.path);

    // Ensure cache directory exists
    const cacheDir = dirname(this.cachePath);
    if (!existsSync(cacheDir)) {
      mkdirSync(cacheDir, { recursive: true });
    }

    // Load existing cache or create new one
    if (existsSync(this.cachePath)) {
      try {
        await this.loadFromFile();
        logger.info('Schema cache loaded from file');
      } catch (error) {
        logger.warn('Failed to load schema cache, creating new:', error);
        await this.refresh();
      }
    } else {
      await this.refresh();
    }

    // Refresh on start if configured
    if (config.schema_cache.refresh_on_start) {
      logger.info('Refreshing schema cache on startup');
      await this.refresh();
    }
  }

  public async refresh(scope: 'all' | 'tables' | 'columns' = 'all', target?: string): Promise<{
    added_tables: string[];
    removed_tables: string[];
    updated_columns: number;
  }> {
    return PerformanceLogger.measure('schema_refresh', async () => {
      const engine = connectionManager.getEngine();
      const connection = connectionManager.getReadOnlyConnection();

      logger.info(`Refreshing schema cache: scope=${scope}, target=${target}`);

      let newTables: TableInfo[] = [];
      
      if (scope === 'all' || scope === 'tables') {
        newTables = await this.fetchTables(engine, connection);
      } else {
        // Keep existing tables if only refreshing columns
        newTables = this.cache?.tables || [];
      }

      if (scope === 'all' || scope === 'columns') {
        // Update columns for all tables or specific target
        const tablesToUpdate = target 
          ? newTables.filter(t => `${t.schema}.${t.table}` === target)
          : newTables;

        for (const table of tablesToUpdate) {
          table.columns = await this.fetchColumns(engine, connection, table.schema, table.table);
        }
      }

      // Calculate differences
      const oldTables = this.cache?.tables || [];
      const oldTableNames = new Set(oldTables.map(t => `${t.schema}.${t.table}`));
      const newTableNames = new Set(newTables.map(t => `${t.schema}.${t.table}`));

      const addedTables = newTables
        .filter(t => !oldTableNames.has(`${t.schema}.${t.table}`))
        .map(t => `${t.schema}.${t.table}`);

      const removedTables = oldTables
        .filter(t => !newTableNames.has(`${t.schema}.${t.table}`))
        .map(t => `${t.schema}.${t.table}`);

      let updatedColumns = 0;
      for (const table of newTables) {
        const oldTable = oldTables.find(t => t.schema === table.schema && t.table === table.table);
        if (oldTable) {
          updatedColumns += table.columns.length - oldTable.columns.length;
        }
      }

      // Update cache
      this.cache = {
        engine,
        db: 'app', // TODO: get from connection config
        last_updated: new Date().toISOString(),
        tables: newTables
      };

      await this.saveToFile();

      logger.info(`Schema refresh completed: ${addedTables.length} added, ${removedTables.length} removed, ${updatedColumns} columns updated`);

      return {
        added_tables: addedTables,
        removed_tables: removedTables,
        updated_columns: Math.max(0, updatedColumns)
      };
    });
  }

  private async fetchTables(engine: DatabaseEngine, connection: any): Promise<TableInfo[]> {
    let tables: TableInfo[] = [];

    if (engine === 'postgres') {
      const sql = `
        SELECT 
          t.table_schema,
          t.table_name,
          obj_description(c.oid) as table_comment
        FROM information_schema.tables t
        LEFT JOIN pg_class c ON c.relname = t.table_name
        LEFT JOIN pg_namespace n ON n.oid = c.relnamespace AND n.nspname = t.table_schema
        WHERE t.table_type = 'BASE TABLE'
          AND t.table_schema NOT IN ('information_schema', 'pg_catalog')
        ORDER BY t.table_schema, t.table_name
      `;

      const rows = await connection.query(sql) as PostgreSQLTableRow[];
      tables = rows.map((row: PostgreSQLTableRow) => ({
        schema: row.table_schema,
        table: row.table_name,
        table_comment: row.table_comment,
        columns: [] // Will be populated separately
      }));
    } else if (engine === 'mysql') {
      const sql = `
        SELECT 
          TABLE_SCHEMA,
          TABLE_NAME,
          TABLE_COMMENT
        FROM information_schema.TABLES
        WHERE TABLE_TYPE = 'BASE TABLE'
          AND TABLE_SCHEMA NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys')
        ORDER BY TABLE_SCHEMA, TABLE_NAME
      `;

      const rows = await connection.query(sql) as MySQLTableRow[];
      tables = rows.map((row: MySQLTableRow) => ({
        schema: row.TABLE_SCHEMA,
        table: row.TABLE_NAME,
        table_comment: row.TABLE_COMMENT || null,
        columns: [] // Will be populated separately
      }));
    }

    return tables;
  }

  private async fetchColumns(engine: DatabaseEngine, connection: any, schema: string, table: string): Promise<ColumnInfo[]> {
    let columns: ColumnInfo[] = [];

    if (engine === 'postgres') {
      const sql = `
        SELECT 
          c.column_name,
          c.data_type,
          c.is_nullable,
          c.column_default,
          col_description(pgc.oid, c.ordinal_position) as column_comment
        FROM information_schema.columns c
        LEFT JOIN pg_class pgc ON pgc.relname = c.table_name
        LEFT JOIN pg_namespace pgn ON pgn.oid = pgc.relnamespace AND pgn.nspname = c.table_schema
        WHERE c.table_schema = $1 AND c.table_name = $2
        ORDER BY c.ordinal_position
      `;

      const rows = await connection.query(sql, [schema, table]) as PostgreSQLColumnRow[];
      columns = rows.map((row: PostgreSQLColumnRow) => ({
        name: row.column_name,
        type: row.data_type,
        nullable: row.is_nullable === 'YES',
        default: row.column_default,
        comment: row.column_comment
      }));
    } else if (engine === 'mysql') {
      const sql = `
        SELECT 
          COLUMN_NAME,
          DATA_TYPE,
          IS_NULLABLE,
          COLUMN_DEFAULT,
          COLUMN_COMMENT
        FROM information_schema.COLUMNS
        WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?
        ORDER BY ORDINAL_POSITION
      `;

      const rows = await connection.query(sql, [schema, table]) as MySQLColumnRow[];
      columns = rows.map((row: MySQLColumnRow) => ({
        name: row.COLUMN_NAME,
        type: row.DATA_TYPE,
        nullable: row.IS_NULLABLE === 'YES',
        default: row.COLUMN_DEFAULT,
        comment: row.COLUMN_COMMENT || null
      }));
    }

    return columns;
  }

  public getCache(): SchemaCache {
    if (!this.cache) {
      throw new SchemaError('Schema cache not initialized');
    }
    return this.cache;
  }

  public filterTables(filter?: {
    schema?: string;
    table_like?: string;
    missing_comment_only?: boolean;
  }): TableInfo[] {
    const cache = this.getCache();
    let tables = cache.tables;

    if (filter?.schema) {
      tables = tables.filter(t => t.schema === filter.schema);
    }

    if (filter?.table_like) {
      const pattern = new RegExp(filter.table_like.replace(/\*/g, '.*'), 'i');
      tables = tables.filter(t => pattern.test(t.table));
    }

    if (filter?.missing_comment_only) {
      tables = tables.filter(t => 
        !t.table_comment || 
        t.columns.some(c => !c.comment)
      );
    }

    return tables;
  }

  private async loadFromFile(): Promise<void> {
    try {
      const content = readFileSync(this.cachePath, 'utf-8');
      this.cache = SchemaCache.parse(JSON.parse(content));
    } catch (error) {
      throw new SchemaError(`Failed to load schema cache: ${error}`);
    }
  }

  private async saveToFile(): Promise<void> {
    try {
      const content = JSON.stringify(this.cache, null, 2);
      writeFileSync(this.cachePath, content, 'utf-8');
    } catch (error) {
      throw new SchemaError(`Failed to save schema cache: ${error}`);
    }
  }
}

export const schemaCacheManager = SchemaCacheManager.getInstance();