@boundless-oss/atlas/src/storage/sqlite-manager.ts

Atlas - MCP Server for comprehensive startup project management

import sqlite3 from 'sqlite3';
import { promises as fs } from 'fs';
import path from 'path';
import { fileURLToPath } from 'url';

const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

export interface QueryResult<T = any> {
  success: boolean;
  data?: T;
  error?: string;
  rowsAffected?: number;
}

export interface TransactionContext {
  query<T = any>(sql: string, params?: any[]): Promise<T[]>;
  run(sql: string, params?: any[]): Promise<{ changes: number; lastID: number }>;
  get<T = any>(sql: string, params?: any[]): Promise<T | undefined>;
}

/**
 * SQLite Database Manager for Atlas
 * Provides a clean interface for database operations with proper error handling
 */
export class SQLiteManager {
  private db: sqlite3.Database | null = null;
  private isInitialized = false;
  private isInitializing = false;
  private dbPath: string;

  constructor(dbPath?: string) {
    this.dbPath = dbPath || process.env.ATLAS_DB_PATH || '.atlas/atlas.db';
  }

  /**
   * Initialize the database connection and schema
   */
  async initialize(): Promise<void> {
    // Prevent multiple initialization attempts
    if (this.isInitialized) {
      return;
    }
    
    if (this.isInitializing) {
      // Wait for the current initialization to complete
      while (this.isInitializing) {
        await new Promise(resolve => setTimeout(resolve, 100));
      }
      return;
    }
    
    this.isInitializing = true;
    
    try {
      // Ensure the directory exists
      const dbDir = path.dirname(this.dbPath);
      await fs.mkdir(dbDir, { recursive: true });

      // Connect to the database
      this.db = await this.openDatabase();
      
      // Enable foreign keys and other pragmas
      await this.runPragmas();
      
      // Load and execute schema
      await this.initializeSchema();
      
      this.isInitialized = true;
      this.isInitializing = false;
      console.error(`📊 SQLite database initialized at: ${this.dbPath}`);
    } catch (error) {
      this.isInitializing = false;
      console.error('❌ Failed to initialize SQLite database:', error);
      throw error;
    }
  }

  /**
   * Open database connection
   */
  private openDatabase(): Promise<sqlite3.Database> {
    return new Promise((resolve, reject) => {
      const db = new sqlite3.Database(this.dbPath, (err) => {
        if (err) {
          reject(err);
        } else {
          resolve(db);
        }
      });
    });
  }

  /**
   * Set database pragmas
   */
  private async runPragmas(): Promise<void> {
    if (!this.db) throw new Error('Database not connected');

    const pragmas = [
      'PRAGMA foreign_keys = ON',
      'PRAGMA journal_mode = WAL',
      'PRAGMA synchronous = NORMAL',
      'PRAGMA cache_size = 1000'
    ];

    for (const pragma of pragmas) {
      await this.runQuery(pragma);
    }
  }

  /**
   * Load and execute the database schema
   */
  private async initializeSchema(): Promise<void> {
    if (!this.db) {
      throw new Error('Database not connected');
    }

    try {
      const schemaPath = path.join(__dirname, 'schema.sql');
      const schema = await fs.readFile(schemaPath, 'utf-8');
      
      // Execute schema statements
      await this.execMultiple(schema);
      console.error('✅ Database schema initialized');
      
      // Check and perform migration if needed
      await this.checkAndPerformMigration();
    } catch (error) {
      console.error('❌ Failed to initialize database schema:', error);
      throw error;
    }
  }

  /**
   * Check if migration is needed and perform one-time migration from JSON files
   */
  private async checkAndPerformMigration(): Promise<void> {
    try {
      // Check migration status
      const migrationCheck = await this.getInternal<{ value: string }>(
        'SELECT value FROM atlas_metadata WHERE key = ?',
        ['migration_status']
      );

      if (migrationCheck.success && migrationCheck.data && migrationCheck.data.value === 'completed') {
        console.error('📋 Migration already completed, skipping');
        return;
      }

      // Import and run migration
      const { DataMigration } = await import('./migration.js');
      const migration = new DataMigration(this, '.atlas', true); // Use internal methods during initialization
      
      const migrationStatus = await migration.checkMigrationStatus();
      
      if (migrationStatus.needsMigration) {
        console.error('🔄 Performing one-time migration from JSON to SQLite...');
        const result = await migration.migrate();
        
        if (result.success) {
          // Mark migration as completed
          await this.runInternal(
            'INSERT OR REPLACE INTO atlas_metadata (key, value, updated_at) VALUES (?, ?, ?)',
            ['migration_status', 'completed', Date.now()]
          );
          console.error('✅ Migration completed successfully');
        } else {
          console.error('❌ Migration failed:', result.error);
          throw new Error(`Migration failed: ${result.error}`);
        }
      } else if (migrationStatus.hasLegacyData) {
        // Has legacy data but migration not required - force migration anyway
        console.error('🔄 Forcing migration due to legacy data presence...');
        const result = await migration.migrate();
        
        if (result.success) {
          await this.runInternal(
            'INSERT OR REPLACE INTO atlas_metadata (key, value, updated_at) VALUES (?, ?, ?)',
            ['migration_status', 'completed', Date.now()]
          );
          console.error('✅ Migration completed successfully');
        } else {
          console.error('❌ Migration failed:', result.error);
          throw new Error(`Migration failed: ${result.error}`);
        }
      } else {
        // No migration needed, mark as completed anyway
        await this.runInternal(
          'INSERT OR REPLACE INTO atlas_metadata (key, value, updated_at) VALUES (?, ?, ?)',
          ['migration_status', 'completed', Date.now()]
        );
        console.error('📋 No migration needed, fresh installation');
      }
    } catch (error) {
      console.error('❌ Migration check/execution failed:', error);
      throw error;
    }
    
    // Check and perform schema migrations
    await this.checkAndPerformSchemaMigration();
  }
  
  /**
   * Check if schema migration is needed and perform schema updates
   */
  private async checkAndPerformSchemaMigration(): Promise<void> {
    try {
      // Check if memories table has new columns
      const tableInfo = await this.query<{ name: string }>(
        "PRAGMA table_info(memories)"
      );
      
      if (tableInfo.success && tableInfo.data) {
        const columns = tableInfo.data.map(col => col.name);
        const requiredColumns = ['title', 'tags', 'importance', 'category', 'source', 'created_by'];
        const missingColumns = requiredColumns.filter(col => !columns.includes(col));
        
        if (missingColumns.length > 0) {
          console.error('🔄 Updating memories table schema...');
          
          // Add missing columns
          for (const column of missingColumns) {
            let defaultValue = '';
            let columnType = 'TEXT';
            
            switch (column) {
              case 'tags':
                defaultValue = " DEFAULT '[]'";
                break;
              case 'importance':
                defaultValue = " DEFAULT 'medium'";
                break;
              default:
                defaultValue = '';
            }
            
            await this.runInternal(
              `ALTER TABLE memories ADD COLUMN ${column} ${columnType}${defaultValue}`
            );
          }
          
          console.error('✅ Memories table schema updated');
        }
      }

      // Check if agile_epics table needs schema updates
      const epicTableInfo = await this.query<{ name: string }>(
        "PRAGMA table_info(agile_epics)"
      );
      
      if (epicTableInfo.success && epicTableInfo.data) {
        const epicColumns = epicTableInfo.data.map(col => col.name);
        const requiredEpicColumns = [
          { name: 'goals', type: 'TEXT', defaultValue: "'[]'" },
          { name: 'owner', type: 'TEXT', defaultValue: null },
          { name: 'target_date', type: 'INTEGER', defaultValue: null },
          { name: 'repositories', type: 'TEXT', defaultValue: "'[]'" }
        ];
        
        let updated = false;
        for (const column of requiredEpicColumns) {
          if (!epicColumns.includes(column.name)) {
            console.error(`🔄 Adding ${column.name} column to agile_epics table...`);
            
            const defaultClause = column.defaultValue ? ` DEFAULT ${column.defaultValue}` : '';
            await this.runInternal(
              `ALTER TABLE agile_epics ADD COLUMN ${column.name} ${column.type}${defaultClause}`
            );
            updated = true;
          }
        }
        
        if (updated) {
          console.error('✅ Agile epics table schema updated');
        }
      }
    } catch (error) {
      console.error('❌ Schema migration failed:', error);
      throw error;
    }
  }

  /**
   * Execute multiple SQL statements
   */
  private execMultiple(sql: string): Promise<void> {
    return new Promise((resolve, reject) => {
      if (!this.db) {
        reject(new Error('Database not connected'));
        return;
      }

      this.db.exec(sql, (err) => {
        if (err) {
          reject(err);
        } else {
          resolve();
        }
      });
    });
  }

  /**
   * Internal get method for use during initialization
   */
  private async getInternal<T = any>(sql: string, params: any[] = []): Promise<QueryResult<T>> {
    try {
      const row = await new Promise<T | undefined>((resolve, reject) => {
        this.db!.get(sql, params, (err, row) => {
          if (err) {
            reject(err);
          } else {
            resolve(row as T | undefined);
          }
        });
      });
      
      return {
        success: true,
        data: row
      };
    } catch (error) {
      console.error('❌ Get error:', error);
      return {
        success: false,
        error: error instanceof Error ? error.message : 'Unknown error'
      };
    }
  }

  /**
   * Internal run method for use during initialization
   */
  private async runInternal(sql: string, params: any[] = []): Promise<QueryResult<{ changes: number; lastInsertRowid: number }>> {
    try {
      const info = await this.runQuery(sql, params);
      
      return {
        success: true,
        data: {
          changes: info.changes,
          lastInsertRowid: info.lastID
        },
        rowsAffected: info.changes
      };
    } catch (error) {
      console.error('❌ Run error:', error);
      return {
        success: false,
        error: error instanceof Error ? error.message : 'Unknown error'
      };
    }
  }

  /**
   * Execute a single SQL statement
   */
  private runQuery(sql: string, params: any[] = []): Promise<{ changes: number; lastID: number }> {
    return new Promise((resolve, reject) => {
      if (!this.db) {
        reject(new Error('Database not connected'));
        return;
      }

      this.db.run(sql, params, function(err) {
        if (err) {
          reject(err);
        } else {
          resolve({ changes: this.changes, lastID: this.lastID });
        }
      });
    });
  }

  /**
   * Execute a SELECT query and return all rows
   */
  async query<T = any>(sql: string, params: any[] = []): Promise<QueryResult<T[]>> {
    if (!this.ensureInitialized()) {
      return { success: false, error: 'Database not initialized' };
    }

    try {
      const rows = await new Promise<T[]>((resolve, reject) => {
        this.db!.all(sql, params, (err, rows) => {
          if (err) {
            reject(err);
          } else {
            resolve(rows as T[]);
          }
        });
      });
      
      return {
        success: true,
        data: rows
      };
    } catch (error) {
      console.error('❌ Query error:', error);
      return {
        success: false,
        error: error instanceof Error ? error.message : 'Unknown error'
      };
    }
  }

  /**
   * Execute a query and return the first row
   */
  async get<T = any>(sql: string, params: any[] = []): Promise<QueryResult<T>> {
    if (!this.ensureInitialized()) {
      return { success: false, error: 'Database not initialized' };
    }

    try {
      const row = await new Promise<T | undefined>((resolve, reject) => {
        this.db!.get(sql, params, (err, row) => {
          if (err) {
            reject(err);
          } else {
            resolve(row as T | undefined);
          }
        });
      });
      
      return {
        success: true,
        data: row
      };
    } catch (error) {
      console.error('❌ Get error:', error);
      return {
        success: false,
        error: error instanceof Error ? error.message : 'Unknown error'
      };
    }
  }

  /**
   * Execute a SELECT query and return all rows (alias for query)
   */
  async all<T = any>(sql: string, params: any[] = []): Promise<QueryResult<T[]>> {
    return this.query<T>(sql, params);
  }

  /**
   * Execute an INSERT, UPDATE, or DELETE query
   */
  async run(sql: string, params: any[] = []): Promise<QueryResult<{ changes: number; lastInsertRowid: number }>> {
    if (!this.ensureInitialized()) {
      return { success: false, error: 'Database not initialized' };
    }

    try {
      const info = await this.runQuery(sql, params);
      
      return {
        success: true,
        data: {
          changes: info.changes,
          lastInsertRowid: info.lastID
        },
        rowsAffected: info.changes
      };
    } catch (error) {
      console.error('❌ Run error:', error);
      return {
        success: false,
        error: error instanceof Error ? error.message : 'Unknown error'
      };
    }
  }

  /**
   * Execute multiple operations in a transaction
   */
  async transaction<T>(fn: (ctx: TransactionContext) => Promise<T>): Promise<QueryResult<T>> {
    if (!this.ensureInitialized()) {
      return { success: false, error: 'Database not initialized' };
    }

    try {
      await this.runQuery('BEGIN TRANSACTION');

      const context: TransactionContext = {
        query: <U>(sql: string, params: any[] = []): Promise<U[]> => {
          return new Promise((resolve, reject) => {
            this.db!.all(sql, params, (err, rows) => {
              if (err) reject(err);
              else resolve(rows as U[]);
            });
          });
        },
        run: (sql: string, params: any[] = []): Promise<{ changes: number; lastID: number }> => {
          return this.runQuery(sql, params);
        },
        get: <U>(sql: string, params: any[] = []): Promise<U | undefined> => {
          return new Promise((resolve, reject) => {
            this.db!.get(sql, params, (err, row) => {
              if (err) reject(err);
              else resolve(row as U | undefined);
            });
          });
        }
      };

      const result = await fn(context);
      await this.runQuery('COMMIT');

      return {
        success: true,
        data: result
      };
    } catch (error) {
      console.error('❌ Transaction error:', error);
      try {
        await this.runQuery('ROLLBACK');
      } catch (rollbackError) {
        console.error('❌ Rollback error:', rollbackError);
      }
      
      return {
        success: false,
        error: error instanceof Error ? error.message : 'Unknown error'
      };
    }
  }

  /**
   * Check if a table exists
   */
  async tableExists(tableName: string): Promise<boolean> {
    const result = await this.get<{ count: number }>(
      "SELECT COUNT(*) as count FROM sqlite_master WHERE type='table' AND name=?",
      [tableName]
    );
    
    return result.success && result.data?.count === 1;
  }

  /**
   * Get database statistics
   */
  async getStats(): Promise<QueryResult<{
    tables: Array<{ name: string; rowCount: number }>;
    dbSize: number;
    pageCount: number;
    pageSize: number;
  }>> {
    if (!this.ensureInitialized()) {
      return { success: false, error: 'Database not initialized' };
    }

    try {
      // Get table names
      const tablesResult = await this.query<{ name: string }>(
        "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
      );

      if (!tablesResult.success || !tablesResult.data) {
        return { success: false, error: 'Failed to get table list' };
      }

      // Get row counts for each table
      const tables = [];
      for (const table of tablesResult.data) {
        const countResult = await this.get<{ count: number }>(
          `SELECT COUNT(*) as count FROM "${table.name}"`
        );
        tables.push({
          name: table.name,
          rowCount: countResult.data?.count || 0
        });
      }

      // Get database size info
      const pageCountResult = await this.get<{ page_count: number }>('PRAGMA page_count');
      const pageSizeResult = await this.get<{ page_size: number }>('PRAGMA page_size');

      const pageCount = pageCountResult.data?.page_count || 0;
      const pageSize = pageSizeResult.data?.page_size || 0;

      return {
        success: true,
        data: {
          tables,
          dbSize: pageCount * pageSize,
          pageCount,
          pageSize
        }
      };
    } catch (error) {
      return {
        success: false,
        error: error instanceof Error ? error.message : 'Unknown error'
      };
    }
  }

  /**
   * Close the database connection
   */
  async close(): Promise<void> {
    if (this.db) {
      await new Promise<void>((resolve, reject) => {
        this.db!.close((err) => {
          if (err) {
            reject(err);
          } else {
            resolve();
          }
        });
      });
      
      this.db = null;
      this.isInitialized = false;
      console.log('📊 SQLite database connection closed');
    }
  }

  /**
   * Check if the database is initialized
   */
  private ensureInitialized(): boolean {
    if (!this.isInitialized || !this.db) {
      console.error('❌ Database not initialized. Call initialize() first.');
      return false;
    }
    return true;
  }

  /**
   * Get the database path
   */
  getDbPath(): string {
    return this.dbPath;
  }

  /**
   * Internal get method for use during initialization (exposed for DataMigration)
   */
  getForMigration<T = any>(sql: string, params: any[] = []): Promise<QueryResult<T>> {
    return this.getInternal<T>(sql, params);
  }

  /**
   * Internal run method for use during initialization (exposed for DataMigration)
   */
  runForMigration(sql: string, params: any[] = []): Promise<QueryResult<{ changes: number; lastInsertRowid: number }>> {
    return this.runInternal(sql, params);
  }

  /**
   * Get database connection info
   */
  getConnectionInfo(): { isInitialized: boolean; dbPath: string } {
    return {
      isInitialized: this.isInitialized,
      dbPath: this.dbPath
    };
  }

  /**
   * Check if database is ready for operations
   */
  isReady(): boolean {
    return this.isInitialized && this.db !== null;
  }

  /**
   * Wait for database to be ready with timeout
   */
  async waitForReady(timeoutMs: number = 10000): Promise<boolean> {
    const startTime = Date.now();
    while (!this.isReady() && Date.now() - startTime < timeoutMs) {
      await new Promise(resolve => setTimeout(resolve, 100));
    }
    return this.isReady();
  }
}

// Singleton instance for convenience
let globalInstance: SQLiteManager | null = null;

export function getSQLiteManager(): SQLiteManager {
  if (!globalInstance) {
    globalInstance = new SQLiteManager();
  }
  return globalInstance;
}

export function createSQLiteManager(dbPath?: string): SQLiteManager {
  return new SQLiteManager(dbPath);
}

/**
 * Utility function for dashboard APIs to ensure database is ready
 */
export async function ensureDatabaseReady(retries: number = 3, delayMs: number = 1000): Promise<SQLiteManager> {
  const db = getSQLiteManager();
  
  // If database is not initialized, initialize it
  if (!db.isReady()) {
    try {
      await db.initialize();
    } catch (error) {
      // If initialization fails, it might be because another process is initializing
      // Continue with retry logic below
    }
  }
  
  for (let attempt = 1; attempt <= retries; attempt++) {
    if (db.isReady()) {
      return db;
    }
    
    console.warn(`⚠️ Database not ready, attempt ${attempt}/${retries}, waiting ${delayMs}ms...`);
    
    if (attempt < retries) {
      await new Promise(resolve => setTimeout(resolve, delayMs));
    }
  }
  
  throw new Error('Database not ready after maximum retries');
}