plugin-postgresql-connector/src/server/services/ConnectionManager.ts

NocoBase plugin for connecting to external PostgreSQL databases

import { Pool, PoolClient, PoolConfig } from 'pg';
import CryptoJS from 'crypto-js';

export interface ConnectionConfig {
  host: string;
  port: number;
  database: string;
  username: string;
  password: string;
  ssl: boolean;
  connectionOptions?: object;
}

export interface ConnectionResult {
  success: boolean;
  connectionId?: string;
  error?: string;
}

export class ConnectionManager {
  private pools: Map<string, Pool> = new Map();
  private readonly encryptionKey: string;
  private readonly maxConnections: number;
  private readonly connectionTimeout: number;
  private readonly idleTimeout: number;

  constructor() {
    this.encryptionKey = process.env.ENCRYPTION_KEY || 'default-key-change-in-production';
    this.maxConnections = parseInt(process.env.POSTGRESQL_CONNECTOR_MAX_CONNECTIONS || '10');
    this.connectionTimeout = parseInt(process.env.POSTGRESQL_CONNECTOR_CONNECTION_TIMEOUT || '5000');
    this.idleTimeout = parseInt(process.env.POSTGRESQL_CONNECTOR_IDLE_TIMEOUT || '30000');

    if (this.encryptionKey === 'default-key-change-in-production') {
      console.warn('WARNING: Using default encryption key. Please set ENCRYPTION_KEY environment variable.');
    }
  }

  /**
   * Create a new connection pool and test the connection
   */
  async createConnection(connectionConfig: ConnectionConfig): Promise<ConnectionResult> {
    const connectionId = this.generateConnectionId();
    
    try {
      // Create connection pool configuration
      const poolConfig: PoolConfig = {
        host: connectionConfig.host,
        port: connectionConfig.port,
        database: connectionConfig.database,
        user: connectionConfig.username,
        password: connectionConfig.password,
        max: this.maxConnections,
        idleTimeoutMillis: this.idleTimeout,
        connectionTimeoutMillis: this.connectionTimeout,
        ssl: connectionConfig.ssl ? { rejectUnauthorized: false } : false,
        ...connectionConfig.connectionOptions,
      };

      // Create pool
      const pool = new Pool(poolConfig);

      // Test connection
      const client = await pool.connect();
      try {
        await client.query('SELECT NOW()');
        console.log(`Connection ${connectionId} established successfully`);
      } finally {
        client.release();
      }

      // Store pool
      this.pools.set(connectionId, pool);

      return {
        success: true,
        connectionId,
      };
    } catch (error) {
      console.error(`Connection ${connectionId} failed:`, error);
      return {
        success: false,
        error: error instanceof Error ? error.message : 'Unknown connection error',
      };
    }
  }

  /**
   * Get a client from the connection pool
   */
  async getConnection(connectionId: string): Promise<PoolClient> {
    const pool = this.pools.get(connectionId);
    if (!pool) {
      throw new Error(`Connection ${connectionId} not found`);
    }

    try {
      return await pool.connect();
    } catch (error) {
      console.error(`Failed to get connection from pool ${connectionId}:`, error);
      throw new Error(`Failed to get connection: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  /**
   * Close a specific connection pool
   */
  async closeConnection(connectionId: string): Promise<void> {
    const pool = this.pools.get(connectionId);
    if (pool) {
      try {
        await pool.end();
        this.pools.delete(connectionId);
        console.log(`Connection ${connectionId} closed successfully`);
      } catch (error) {
        console.error(`Error closing connection ${connectionId}:`, error);
        throw error;
      }
    }
  }

  /**
   * Close all connection pools
   */
  async closeAllConnections(): Promise<void> {
    const closePromises = Array.from(this.pools.keys()).map(connectionId => 
      this.closeConnection(connectionId)
    );
    
    await Promise.all(closePromises);
    console.log('All connections closed');
  }

  /**
   * Test connection without creating a pool
   */
  async testConnection(connectionConfig: ConnectionConfig): Promise<ConnectionResult> {
    const poolConfig: PoolConfig = {
      host: connectionConfig.host,
      port: connectionConfig.port,
      database: connectionConfig.database,
      user: connectionConfig.username,
      password: connectionConfig.password,
      max: 1, // Only one connection for testing
      idleTimeoutMillis: 1000,
      connectionTimeoutMillis: this.connectionTimeout,
      ssl: connectionConfig.ssl ? { rejectUnauthorized: false } : false,
    };

    const pool = new Pool(poolConfig);

    try {
      const client = await pool.connect();
      try {
        await client.query('SELECT NOW()');
        return { success: true };
      } finally {
        client.release();
      }
    } catch (error) {
      return {
        success: false,
        error: error instanceof Error ? error.message : 'Connection test failed',
      };
    } finally {
      await pool.end();
    }
  }

  /**
   * Get connection pool statistics
   */
  getConnectionStats(connectionId: string): object | null {
    const pool = this.pools.get(connectionId);
    if (!pool) {
      return null;
    }

    return {
      totalCount: pool.totalCount,
      idleCount: pool.idleCount,
      waitingCount: pool.waitingCount,
    };
  }

  /**
   * Get all active connection IDs
   */
  getActiveConnections(): string[] {
    return Array.from(this.pools.keys());
  }

  /**
   * Encrypt password for storage
   */
  encryptPassword(password: string): string {
    try {
      return CryptoJS.AES.encrypt(password, this.encryptionKey).toString();
    } catch (error) {
      console.error('Password encryption failed:', error);
      throw new Error('Failed to encrypt password');
    }
  }

  /**
   * Decrypt password from storage
   */
  decryptPassword(encryptedPassword: string): string {
    try {
      const bytes = CryptoJS.AES.decrypt(encryptedPassword, this.encryptionKey);
      const decrypted = bytes.toString(CryptoJS.enc.Utf8);
      
      if (!decrypted) {
        throw new Error('Failed to decrypt password - invalid key or corrupted data');
      }
      
      return decrypted;
    } catch (error) {
      console.error('Password decryption failed:', error);
      throw new Error('Failed to decrypt password');
    }
  }

  /**
   * Generate unique connection ID
   */
  private generateConnectionId(): string {
    const timestamp = Date.now();
    const random = Math.random().toString(36).substr(2, 9);
    return `conn_${timestamp}_${random}`;
  }

  /**
   * Cleanup method for graceful shutdown
   */
  async destroy(): Promise<void> {
    console.log('ConnectionManager: Starting cleanup...');
    await this.closeAllConnections();
    console.log('ConnectionManager: Cleanup completed');
  }
}

export default ConnectionManager;