@boundless-oss/atlas/src/modules/migration-management/migration-manager.ts

Atlas - MCP Server for comprehensive startup project management

import { promises as fs } from 'fs';
import path from 'path';
import crypto from 'crypto';
import {
  Migration,
  MigrationPlan,
  MigrationHistory,
  MigrationRecord,
  MigrationConfig,
  MigrationStatus,
  MigrationLog,
  MigrationValidation,
  MigrationTemplate,
  MigrationRisk,
  DependencyGraph,
} from './types.js';

export class MigrationManager {
  private config: MigrationConfig;
  private history: MigrationHistory | null = null;

  constructor(config: MigrationConfig) {
    this.config = config;
  }

  async createMigration(
    name: string,
    type: Migration['type'],
    template?: MigrationTemplate
  ): Promise<Migration> {
    const version = this.generateVersion();
    const fileName = `${version}_${name.toLowerCase().replace(/\s+/g, '_')}.js`;
    const filePath = path.join(this.config.migrationsPath, fileName);

    // Create migration content
    let content: string;
    if (template) {
      content = this.applyTemplate(template);
    } else {
      content = this.getDefaultTemplate(type);
    }

    // Create migration object
    const migration: Migration = {
      id: crypto.randomUUID(),
      version,
      name,
      type,
      status: 'pending',
      createdAt: new Date().toISOString(),
      checksum: this.calculateChecksum(content),
      up: {
        type: 'javascript',
        content,
        transaction: true,
      },
      down: {
        type: 'javascript',
        content: '// Rollback logic here',
        transaction: true,
      },
    };

    // Write migration file
    await fs.mkdir(this.config.migrationsPath, { recursive: true });
    await fs.writeFile(filePath, content, 'utf-8');

    return migration;
  }

  async runMigrations(options: {
    target?: string;
    dryRun?: boolean;
    force?: boolean;
  } = {}): Promise<MigrationRecord[]> {
    await this.loadHistory();
    
    const pending = await this.getPendingMigrations();
    if (pending.length === 0) {
      return [];
    }

    // Filter to target version if specified
    let toRun = pending;
    if (options.target) {
      const targetIndex = pending.findIndex(m => m.version === options.target);
      if (targetIndex >= 0) {
        toRun = pending.slice(0, targetIndex + 1);
      }
    }

    // Create execution plan
    const plan = await this.createMigrationPlan(toRun);
    
    // Validate plan
    if (!options.force) {
      const validation = await this.validateMigration(plan.migrations[0]);
      if (!validation.isValid) {
        throw new Error(`Migration validation failed: ${validation.errors.join(', ')}`);
      }
    }

    // Execute migrations
    const records: MigrationRecord[] = [];
    
    for (const migration of plan.migrations) {
      if (options.dryRun) {
        this.log('info', `[DRY RUN] Would execute migration: ${migration.version} - ${migration.name}`);
        continue;
      }

      const record = await this.executeMigration(migration);
      records.push(record);
      
      if (record.status === 'failed') {
        throw new Error(`Migration ${migration.version} failed: ${record.error}`);
      }
    }

    return records;
  }

  async rollback(options: {
    steps?: number;
    target?: string;
    force?: boolean;
  } = {}): Promise<MigrationRecord[]> {
    await this.loadHistory();
    
    if (!this.history || this.history.migrations.length === 0) {
      throw new Error('No migrations to rollback');
    }

    // Determine which migrations to rollback
    let toRollback: MigrationRecord[] = [];
    
    if (options.target) {
      // Rollback to specific version
      const targetIndex = this.history.migrations.findIndex(m => m.version === options.target);
      if (targetIndex >= 0) {
        toRollback = this.history.migrations.slice(targetIndex + 1).reverse();
      }
    } else if (options.steps) {
      // Rollback specific number of steps
      toRollback = this.history.migrations.slice(-options.steps).reverse();
    } else {
      // Rollback last migration
      toRollback = [this.history.migrations[this.history.migrations.length - 1]];
    }

    const records: MigrationRecord[] = [];
    
    for (const migrationRecord of toRollback) {
      const migration = await this.loadMigration(migrationRecord.version);
      if (!migration.down) {
        if (!options.force) {
          throw new Error(`Migration ${migration.version} cannot be rolled back (no down script)`);
        }
        continue;
      }

      const record = await this.executeRollback(migration);
      records.push(record);
      
      if (record.status === 'failed') {
        throw new Error(`Rollback of ${migration.version} failed: ${record.error}`);
      }
    }

    return records;
  }

  async getMigrationStatus(): Promise<{
    current: string;
    pending: number;
    executed: number;
    failed: number;
  }> {
    await this.loadHistory();
    
    const pending = await this.getPendingMigrations();
    const executed = this.history?.migrations.filter(m => m.status === 'completed').length || 0;
    const failed = this.history?.migrations.filter(m => m.status === 'failed').length || 0;
    
    return {
      current: this.history?.currentVersion || 'none',
      pending: pending.length,
      executed,
      failed,
    };
  }

  async validateMigration(migration: Migration): Promise<MigrationValidation> {
    const validation: MigrationValidation = {
      isValid: true,
      errors: [],
      warnings: [],
      suggestions: [],
    };

    // Check syntax
    if (this.config.validation?.syntaxCheck) {
      try {
        await this.validateSyntax(migration);
      } catch (error: any) {
        validation.errors.push(`Syntax error: ${error.message}`);
        validation.isValid = false;
      }
    }

    // Check dependencies
    if (migration.dependencies && migration.dependencies.length > 0) {
      for (const dep of migration.dependencies) {
        const exists = await this.migrationExists(dep);
        if (!exists) {
          validation.errors.push(`Missing dependency: ${dep}`);
          validation.isValid = false;
        }
      }
    }

    // Check rollback script
    if (this.config.validation?.requireRollback && !migration.down) {
      validation.errors.push('Rollback script is required but not provided');
      validation.isValid = false;
    }

    // Warnings
    if (migration.type === 'schema' && !migration.down) {
      validation.warnings.push('Schema migration without rollback script - this may cause issues');
    }

    if (!migration.description) {
      validation.warnings.push('Migration lacks description - consider adding one for clarity');
    }

    // Suggestions
    if (migration.up.transaction === false && migration.type === 'data') {
      validation.suggestions.push('Consider using transactions for data migrations');
    }

    return validation;
  }

  private async createMigrationPlan(migrations: Migration[]): Promise<MigrationPlan> {
    // Build dependency graph
    const graph = this.buildDependencyGraph(migrations);
    
    // Topological sort for execution order
    const sorted = this.topologicalSort(graph, migrations);
    
    // Assess risks
    const risks = this.assessMigrationRisks(sorted);
    
    // Estimate duration
    const estimatedDuration = sorted.reduce((total, m) => {
      return total + this.estimateMigrationDuration(m);
    }, 0);
    
    // Check if all migrations have rollback scripts
    const rollbackable = sorted.every(m => !!m.down);
    
    return {
      migrations: sorted,
      estimatedDuration,
      risks,
      rollbackable,
      dependencies: graph,
    };
  }

  private buildDependencyGraph(migrations: Migration[]): DependencyGraph {
    const nodes = migrations.map(m => m.version);
    const edges: Array<[string, string]> = [];
    
    migrations.forEach(migration => {
      if (migration.dependencies) {
        migration.dependencies.forEach(dep => {
          edges.push([dep, migration.version]);
        });
      }
    });
    
    return { nodes, edges };
  }

  private topologicalSort(graph: DependencyGraph, migrations: Migration[]): Migration[] {
    const visited = new Set<string>();
    const result: Migration[] = [];
    
    const visit = (version: string) => {
      if (visited.has(version)) return;
      visited.add(version);
      
      // Visit dependencies first
      const deps = graph.edges.filter(([from, to]) => to === version).map(([from]) => from);
      deps.forEach(dep => visit(dep));
      
      const migration = migrations.find(m => m.version === version);
      if (migration) result.push(migration);
    };
    
    graph.nodes.forEach(node => visit(node));
    return result;
  }

  private assessMigrationRisks(migrations: Migration[]): MigrationRisk[] {
    const risks: MigrationRisk[] = [];
    
    migrations.forEach(migration => {
      // Schema changes are high risk
      if (migration.type === 'schema') {
        risks.push({
          migration: migration.version,
          description: 'Schema changes may cause downtime',
          severity: 'high',
          mitigation: 'Ensure backup is available and test in staging',
        });
      }
      
      // Large data migrations
      if (migration.type === 'data' && migration.up.content.includes('UPDATE') && !migration.up.content.includes('LIMIT')) {
        risks.push({
          migration: migration.version,
          description: 'Unbounded data update may affect performance',
          severity: 'medium',
          mitigation: 'Consider batching updates',
        });
      }
      
      // No rollback script
      if (!migration.down) {
        risks.push({
          migration: migration.version,
          description: 'No rollback script available',
          severity: 'medium',
          mitigation: 'Manual intervention may be required for rollback',
        });
      }
    });
    
    return risks;
  }

  private estimateMigrationDuration(migration: Migration): number {
    // Base estimates by type (in milliseconds)
    const baseEstimates: Record<Migration['type'], number> = {
      schema: 5000,
      data: 30000,
      seed: 10000,
      index: 60000,
      procedure: 2000,
      config: 1000,
    };
    
    return baseEstimates[migration.type] || 5000;
  }

  private async executeMigration(migration: Migration): Promise<MigrationRecord> {
    const record: MigrationRecord = {
      migrationId: migration.id,
      version: migration.version,
      name: migration.name,
      status: 'running',
      executedAt: new Date().toISOString(),
      duration: 0,
      logs: [],
    };
    
    const startTime = Date.now();
    
    try {
      // Run pre-migration hooks
      if (this.config.hooks?.beforeEach) {
        await this.runHook(this.config.hooks.beforeEach, migration);
      }
      
      // Execute migration
      this.log('info', `Executing migration: ${migration.version} - ${migration.name}`);
      await this.runMigrationScript(migration.up, migration);
      
      // Run post-migration hooks
      if (this.config.hooks?.afterEach) {
        await this.runHook(this.config.hooks.afterEach, migration);
      }
      
      record.status = 'completed';
      record.duration = Date.now() - startTime;
      this.log('info', `Migration completed in ${record.duration}ms`);
      
      // Update history
      await this.updateHistory(record);
      
    } catch (error: any) {
      record.status = 'failed';
      record.error = error.message;
      record.duration = Date.now() - startTime;
      this.log('error', `Migration failed: ${error.message}`);
      
      // Run error hooks
      if (this.config.hooks?.onError) {
        await this.runHook(this.config.hooks.onError, migration);
      }
    }
    
    return record;
  }

  private async executeRollback(migration: Migration): Promise<MigrationRecord> {
    if (!migration.down) {
      throw new Error('No rollback script available');
    }
    
    const record: MigrationRecord = {
      migrationId: migration.id,
      version: migration.version,
      name: migration.name,
      status: 'running',
      executedAt: new Date().toISOString(),
      duration: 0,
      logs: [],
    };
    
    const startTime = Date.now();
    
    try {
      this.log('info', `Rolling back migration: ${migration.version} - ${migration.name}`);
      await this.runMigrationScript(migration.down, migration);
      
      record.status = 'rolled_back';
      record.duration = Date.now() - startTime;
      this.log('info', `Rollback completed in ${record.duration}ms`);
      
      // Remove from history
      await this.removeFromHistory(migration.version);
      
    } catch (error: any) {
      record.status = 'failed';
      record.error = error.message;
      record.duration = Date.now() - startTime;
      this.log('error', `Rollback failed: ${error.message}`);
    }
    
    return record;
  }

  private async runMigrationScript(script: Migration['up'], migration: Migration): Promise<void> {
    switch (script.type) {
      case 'javascript':
      case 'typescript':
        // Dynamic import and execute
        const migrationPath = path.join(this.config.migrationsPath, `${migration.version}_*.js`);
        const module = await import(migrationPath);
        if (typeof module.up === 'function') {
          await module.up();
        }
        break;
        
      case 'sql':
        // Execute SQL based on database config
        if (this.config.database) {
          // Would execute SQL here based on database type
          this.log('info', 'Executing SQL migration');
        }
        break;
        
      case 'shell':
        // Execute shell command
        const { exec } = await import('child_process');
        const { promisify } = await import('util');
        const execAsync = promisify(exec);
        await execAsync(script.content);
        break;
    }
  }

  private async getPendingMigrations(): Promise<Migration[]> {
    const allMigrations = await this.loadAllMigrations();
    const executedVersions = new Set(
      this.history?.migrations.map(m => m.version) || []
    );
    
    return allMigrations.filter(m => !executedVersions.has(m.version));
  }

  private async loadAllMigrations(): Promise<Migration[]> {
    const files = await fs.readdir(this.config.migrationsPath);
    const migrationFiles = files.filter(f => f.match(/^\d{14}_.*\.(js|ts|sql)$/));
    
    const migrations: Migration[] = [];
    
    for (const file of migrationFiles) {
      const version = file.substring(0, 14);
      const name = file.substring(15).replace(/\.(js|ts|sql)$/, '').replace(/_/g, ' ');
      const content = await fs.readFile(path.join(this.config.migrationsPath, file), 'utf-8');
      
      migrations.push({
        id: crypto.randomUUID(),
        version,
        name,
        type: this.inferMigrationType(content),
        status: 'pending',
        createdAt: new Date().toISOString(),
        checksum: this.calculateChecksum(content),
        up: {
          type: file.endsWith('.sql') ? 'sql' : 'javascript',
          content,
          transaction: true,
        },
      });
    }
    
    return migrations.sort((a, b) => a.version.localeCompare(b.version));
  }

  private async loadMigration(version: string): Promise<Migration> {
    const migrations = await this.loadAllMigrations();
    const migration = migrations.find(m => m.version === version);
    
    if (!migration) {
      throw new Error(`Migration ${version} not found`);
    }
    
    return migration;
  }

  private async loadHistory(): Promise<void> {
    // Load from storage - simplified for example
    this.history = {
      projectId: 'default',
      migrations: [],
      currentVersion: 'none',
    };
  }

  private async updateHistory(record: MigrationRecord): Promise<void> {
    if (!this.history) {
      await this.loadHistory();
    }
    
    this.history!.migrations.push(record);
    this.history!.currentVersion = record.version;
    this.history!.lastMigration = record;
    
    // Save to storage
  }

  private async removeFromHistory(version: string): Promise<void> {
    if (!this.history) return;
    
    this.history.migrations = this.history.migrations.filter(m => m.version !== version);
    
    if (this.history.migrations.length > 0) {
      const last = this.history.migrations[this.history.migrations.length - 1];
      this.history.currentVersion = last.version;
      this.history.lastMigration = last;
    } else {
      this.history.currentVersion = 'none';
      this.history.lastMigration = undefined;
    }
    
    // Save to storage
  }

  private generateVersion(): string {
    const now = new Date();
    const year = now.getFullYear();
    const month = String(now.getMonth() + 1).padStart(2, '0');
    const day = String(now.getDate()).padStart(2, '0');
    const hour = String(now.getHours()).padStart(2, '0');
    const minute = String(now.getMinutes()).padStart(2, '0');
    const second = String(now.getSeconds()).padStart(2, '0');
    
    return `${year}${month}${day}${hour}${minute}${second}`;
  }

  private calculateChecksum(content: string): string {
    return crypto.createHash('sha256').update(content).digest('hex');
  }

  private inferMigrationType(content: string): Migration['type'] {
    if (content.includes('CREATE TABLE') || content.includes('ALTER TABLE')) {
      return 'schema';
    }
    if (content.includes('INSERT INTO') && content.includes('seed')) {
      return 'seed';
    }
    if (content.includes('CREATE INDEX')) {
      return 'index';
    }
    if (content.includes('CREATE PROCEDURE') || content.includes('CREATE FUNCTION')) {
      return 'procedure';
    }
    if (content.includes('UPDATE') || content.includes('INSERT')) {
      return 'data';
    }
    return 'config';
  }

  private async validateSyntax(migration: Migration): Promise<void> {
    // Basic syntax validation based on type
    if (migration.up.type === 'sql') {
      // Check for common SQL syntax issues
      const sql = migration.up.content.toUpperCase();
      if (sql.includes('DROP TABLE') && !sql.includes('IF EXISTS')) {
        throw new Error('DROP TABLE should use IF EXISTS clause');
      }
    }
  }

  private async migrationExists(version: string): Promise<boolean> {
    const migrations = await this.loadAllMigrations();
    return migrations.some(m => m.version === version);
  }

  private async runHook(hook: string, migration: Migration): Promise<void> {
    // Execute hook command
    const { exec } = await import('child_process');
    const { promisify } = await import('util');
    const execAsync = promisify(exec);
    
    await execAsync(hook, {
      env: {
        ...process.env,
        MIGRATION_VERSION: migration.version,
        MIGRATION_NAME: migration.name,
        MIGRATION_TYPE: migration.type,
      },
    });
  }

  private getDefaultTemplate(type: Migration['type']): string {
    const templates: Record<Migration['type'], string> = {
      schema: `
// Schema migration
exports.up = async function(db) {
  // Create tables, alter columns, etc.
};

exports.down = async function(db) {
  // Revert schema changes
};
`,
      data: `
// Data migration
exports.up = async function(db) {
  // Transform or migrate data
};

exports.down = async function(db) {
  // Revert data changes
};
`,
      seed: `
// Seed data
exports.up = async function(db) {
  // Insert seed data
};

exports.down = async function(db) {
  // Remove seed data
};
`,
      index: `
// Index migration
exports.up = async function(db) {
  // Create indexes
};

exports.down = async function(db) {
  // Drop indexes
};
`,
      procedure: `
// Stored procedure migration
exports.up = async function(db) {
  // Create procedures/functions
};

exports.down = async function(db) {
  // Drop procedures/functions
};
`,
      config: `
// Configuration migration
exports.up = async function(db) {
  // Apply configuration changes
};

exports.down = async function(db) {
  // Revert configuration changes
};
`,
    };
    
    return templates[type];
  }

  private applyTemplate(template: MigrationTemplate): string {
    let content = template.template;
    
    if (template.variables) {
      Object.entries(template.variables).forEach(([key, value]) => {
        content = content.replace(new RegExp(`{{${key}}}`, 'g'), value);
      });
    }
    
    return content;
  }

  private log(level: MigrationLog['level'], message: string, details?: any): void {
    // Would normally log to file or send to logging service
    console.log(`[${level.toUpperCase()}] ${message}`, details || '');
  }
}