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

NocoBase plugin for connecting to external PostgreSQL databases

import { PoolClient } from 'pg';
import { ConnectionManager } from './ConnectionManager';
import { format } from 'sql-formatter';

export interface QueryResult {
  rows: any[];
  rowCount: number;
  fields?: Array<{
    name: string;
    dataTypeID: number;
    dataTypeSize: number;
    dataTypeName?: string;
  }>;
  executionTime: number;
  query: string;
  formattedQuery?: string;
}

export interface QueryParameter {
  name: string;
  value: any;
  type?: string;
}

export interface QueryOptions {
  timeout?: number;
  maxRows?: number;
  formatQuery?: boolean;
  includeMetadata?: boolean;
}

export interface SchemaInfo {
  tables: Array<{
    table_name: string;
    table_type: string;
    table_schema: string;
  }>;
  views: Array<{
    view_name: string;
    view_definition?: string;
    table_schema: string;
  }>;
  functions: Array<{
    function_name: string;
    routine_type: string;
    data_type?: string;
    routine_definition?: string;
  }>;
  procedures: Array<{
    procedure_name: string;
    routine_type: string;
    routine_definition?: string;
  }>;
}

export abstract class QueryStrategy {
  abstract execute(
    client: PoolClient,
    query: string,
    params: any[],
    options: QueryOptions
  ): Promise<QueryResult>;

  protected getDataTypeName(dataTypeID: number): string {
    const typeMap: { [key: number]: string } = {
      16: 'BOOLEAN',
      17: 'BYTEA',
      20: 'BIGINT',
      21: 'SMALLINT',
      23: 'INTEGER',
      25: 'TEXT',
      700: 'REAL',
      701: 'DOUBLE_PRECISION',
      1043: 'VARCHAR',
      1082: 'DATE',
      1083: 'TIME',
      1114: 'TIMESTAMP',
      1184: 'TIMESTAMPTZ',
      1700: 'NUMERIC',
      2950: 'UUID',
      3802: 'JSONB',
      114: 'JSON',
    };
    return typeMap[dataTypeID] || 'UNKNOWN';
  }

  protected formatExecutionTime(startTime: number): number {
    return Date.now() - startTime;
  }
}

export class SelectQueryStrategy extends QueryStrategy {
  async execute(
    client: PoolClient,
    query: string,
    params: any[] = [],
    options: QueryOptions = {}
  ): Promise<QueryResult> {
    const startTime = Date.now();
    
    try {
      // Apply row limit if specified
      let finalQuery = query;
      if (options.maxRows && !query.toLowerCase().includes('limit')) {
        finalQuery = `${query} LIMIT ${options.maxRows}`;
      }

      const result = await client.query(finalQuery, params);
      
      return {
        rows: result.rows,
        rowCount: result.rowCount || 0,
        fields: result.fields?.map(field => ({
          name: field.name,
          dataTypeID: field.dataTypeID,
          dataTypeSize: field.dataTypeSize,
          dataTypeName: this.getDataTypeName(field.dataTypeID),
        })),
        executionTime: this.formatExecutionTime(startTime),
        query: finalQuery,
        formattedQuery: options.formatQuery ? format(finalQuery) : undefined,
      };
    } catch (error) {
      throw new Error(`SELECT query failed: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }
}

export class ModifyQueryStrategy extends QueryStrategy {
  async execute(
    client: PoolClient,
    query: string,
    params: any[] = [],
    options: QueryOptions = {}
  ): Promise<QueryResult> {
    const startTime = Date.now();
    
    try {
      const result = await client.query(query, params);
      
      return {
        rows: result.rows || [],
        rowCount: result.rowCount || 0,
        fields: result.fields?.map(field => ({
          name: field.name,
          dataTypeID: field.dataTypeID,
          dataTypeSize: field.dataTypeSize,
          dataTypeName: this.getDataTypeName(field.dataTypeID),
        })),
        executionTime: this.formatExecutionTime(startTime),
        query,
        formattedQuery: options.formatQuery ? format(query) : undefined,
      };
    } catch (error) {
      throw new Error(`Modify query failed: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }
}

export class ProcedureQueryStrategy extends QueryStrategy {
  async execute(
    client: PoolClient,
    query: string,
    params: any[] = [],
    options: QueryOptions = {}
  ): Promise<QueryResult> {
    const startTime = Date.now();
    
    try {
      // For stored procedures, we might need to use CALL syntax
      const result = await client.query(query, params);
      
      return {
        rows: result.rows || [],
        rowCount: result.rowCount || 0,
        fields: result.fields?.map(field => ({
          name: field.name,
          dataTypeID: field.dataTypeID,
          dataTypeSize: field.dataTypeSize,
          dataTypeName: this.getDataTypeName(field.dataTypeID),
        })),
        executionTime: this.formatExecutionTime(startTime),
        query,
        formattedQuery: options.formatQuery ? format(query) : undefined,
      };
    } catch (error) {
      throw new Error(`Procedure execution failed: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }
}

export class QueryExecutor {
  private queryCache: Map<string, { result: QueryResult; timestamp: number }> = new Map();
  private readonly cacheTTL: number = 300000; // 5 minutes

  constructor(private connectionManager: ConnectionManager) {}

  async executeQuery(
    connectionId: string,
    query: string,
    params: any[] = [],
    options: QueryOptions = {}
  ): Promise<QueryResult> {
    const cleanQuery = this.validateAndCleanQuery(query);
    const queryType = this.detectQueryType(cleanQuery);
    
    // Check cache for SELECT queries
    if (queryType === 'SELECT' && this.shouldUseCache(cleanQuery, params)) {
      const cached = this.getCachedResult(cleanQuery, params);
      if (cached) {
        return cached;
      }
    }

    const client = await this.connectionManager.getConnection(connectionId);
    
    try {
      const startTime = Date.now();
      
      // Set query timeout if specified
      if (options.timeout) {
        await client.query(`SET statement_timeout = ${options.timeout}`);
      }

      // Apply row limit if specified and not already present
      let finalQuery = cleanQuery;
      if (options.maxRows && !cleanQuery.toLowerCase().includes('limit')) {
        finalQuery = `${cleanQuery} LIMIT ${options.maxRows}`;
      }

      const result = await client.query(finalQuery, params);
      const executionTime = Date.now() - startTime;

      const queryResult: QueryResult = {
        rows: result.rows,
        rowCount: result.rowCount || 0,
        fields: result.fields?.map(field => ({
          name: field.name,
          dataTypeID: field.dataTypeID,
          dataTypeSize: field.dataTypeSize,
          dataTypeName: this.getDataTypeName(field.dataTypeID),
        })),
        executionTime,
        query: finalQuery,
        formattedQuery: options.formatQuery ? format(finalQuery) : undefined,
      };

      // Cache SELECT query results
      if (queryType === 'SELECT') {
        this.cacheResult(cleanQuery, params, queryResult);
      }

      return queryResult;
    } finally {
      client.release();
      
      // Reset timeout
      if (options.timeout) {
        await client.query('RESET statement_timeout');
      }
    }
  }

  async executeProcedure(
    connectionId: string,
    procedureName: string,
    params: any[] = [],
    options: QueryOptions = {}
  ): Promise<QueryResult> {
    const paramPlaceholders = params.map((_, index) => `$${index + 1}`).join(', ');
    const query = `CALL ${procedureName}(${paramPlaceholders})`;
    return this.executeQuery(connectionId, query, params, options);
  }

  async executeFunction(
    connectionId: string,
    functionName: string,
    params: any[] = [],
    options: QueryOptions = {}
  ): Promise<QueryResult> {
    const paramPlaceholders = params.map((_, index) => `$${index + 1}`).join(', ');
    const query = `SELECT ${functionName}(${paramPlaceholders})`;
    return this.executeQuery(connectionId, query, params, options);
  }

  async getViewData(
    connectionId: string,
    viewName: string,
    limit: number = 100,
    options: QueryOptions = {}
  ): Promise<QueryResult> {
    const query = `SELECT * FROM ${viewName} LIMIT $1`;
    return this.executeQuery(connectionId, query, [limit], options);
  }

  async getTableData(
    connectionId: string,
    tableName: string,
    limit: number = 100,
    offset: number = 0,
    orderBy?: string,
    options: QueryOptions = {}
  ): Promise<QueryResult> {
    let query = `SELECT * FROM ${tableName}`;
    const params: any[] = [];
    let paramIndex = 1;

    if (orderBy) {
      query += ` ORDER BY ${orderBy}`;
    }

    query += ` LIMIT $${paramIndex++}`;
    params.push(limit);

    if (offset > 0) {
      query += ` OFFSET $${paramIndex}`;
      params.push(offset);
    }

    return this.executeQuery(connectionId, query, params, options);
  }

  private validateAndCleanQuery(query: string): string {
    const trimmedQuery = query.trim();
    
    if (!trimmedQuery) {
      throw new Error('Query cannot be empty');
    }

    // Basic security checks
    const dangerousPatterns = [
      /\b(DROP|TRUNCATE|ALTER)\s+/i,
      /\b(CREATE|DROP)\s+(USER|ROLE)\s+/i,
      /\b(GRANT|REVOKE)\s+/i,
      /\bSHUTDOWN\b/i,
      /\bRESTART\b/i,
    ];

    for (const pattern of dangerousPatterns) {
      if (pattern.test(trimmedQuery)) {
        throw new Error('Potentially dangerous SQL statement detected');
      }
    }

    return trimmedQuery;
  }

  private detectQueryType(query: string): string {
    const upperQuery = query.toUpperCase().trim();
    
    if (upperQuery.startsWith('SELECT')) return 'SELECT';
    if (upperQuery.startsWith('INSERT')) return 'INSERT';
    if (upperQuery.startsWith('UPDATE')) return 'UPDATE';
    if (upperQuery.startsWith('DELETE')) return 'DELETE';
    if (upperQuery.startsWith('CALL')) return 'CALL';
    if (upperQuery.startsWith('WITH')) return 'SELECT'; // CTE queries
    
    // Check for function calls
    if (upperQuery.startsWith('SELECT') && upperQuery.includes('(')) {
      return 'FUNCTION';
    }
    
    return 'SELECT'; // Default to SELECT for safety
  }

  private getDataTypeName(dataTypeID: number): string {
    const typeMap: { [key: number]: string } = {
      16: 'BOOLEAN',
      17: 'BYTEA',
      20: 'BIGINT',
      21: 'SMALLINT',
      23: 'INTEGER',
      25: 'TEXT',
      700: 'REAL',
      701: 'DOUBLE_PRECISION',
      1043: 'VARCHAR',
      1082: 'DATE',
      1083: 'TIME',
      1114: 'TIMESTAMP',
      1184: 'TIMESTAMPTZ',
      1700: 'NUMERIC',
      2950: 'UUID',
      3802: 'JSONB',
      114: 'JSON',
    };
    return typeMap[dataTypeID] || 'UNKNOWN';
  }

  private shouldUseCache(query: string, params: any[]): boolean {
    // Only cache simple SELECT queries without parameters that might change frequently
    const upperQuery = query.toUpperCase();
    
    // Don't cache queries with NOW(), CURRENT_TIMESTAMP, etc.
    const nonCacheablePatterns = [
      /\bNOW\(\)/i,
      /\bCURRENT_TIMESTAMP\b/i,
      /\bCURRENT_DATE\b/i,
      /\bRANDOM\(\)/i,
      /\bGEN_RANDOM_UUID\(\)/i,
    ];

    return !nonCacheablePatterns.some(pattern => pattern.test(upperQuery));
  }

  private getCachedResult(query: string, params: any[]): QueryResult | null {
    const cacheKey = this.generateCacheKey(query, params);
    const cached = this.queryCache.get(cacheKey);
    
    if (cached && Date.now() - cached.timestamp < this.cacheTTL) {
      return { ...cached.result }; // Return a copy
    }
    
    if (cached) {
      this.queryCache.delete(cacheKey);
    }
    
    return null;
  }

  private cacheResult(query: string, params: any[], result: QueryResult): void {
    const cacheKey = this.generateCacheKey(query, params);
    this.queryCache.set(cacheKey, {
      result: { ...result }, // Store a copy
      timestamp: Date.now(),
    });

    // Cleanup old cache entries if cache gets too large
    if (this.queryCache.size > 100) {
      this.cleanupCache();
    }
  }

  private generateCacheKey(query: string, params: any[]): string {
    return `${query}:${JSON.stringify(params)}`;
  }

  private cleanupCache(): void {
    const now = Date.now();
    for (const [key, value] of this.queryCache) {
      if (now - value.timestamp > this.cacheTTL) {
        this.queryCache.delete(key);
      }
    }
  }

  /**
   * Get query execution statistics
   */
  getStatistics(): object {
    return {
      cacheSize: this.queryCache.size,
      supportedOperations: ['SELECT', 'INSERT', 'UPDATE', 'DELETE', 'CALL', 'PROCEDURE', 'FUNCTION'],
      cacheTTL: this.cacheTTL,
    };
  }

  /**
   * Clear query cache
   */
  clearCache(): void {
    this.queryCache.clear();
  }

  /**
   * Analyze query complexity (basic implementation)
   */
  analyzeQuery(query: string): object {
    const upperQuery = query.toUpperCase();
    const complexity = {
      joins: (upperQuery.match(/\bJOIN\b/g) || []).length,
      subqueries: (upperQuery.match(/\bSELECT\b/g) || []).length - 1,
      aggregations: (upperQuery.match(/\b(COUNT|SUM|AVG|MIN|MAX|GROUP BY)\b/g) || []).length,
      orderBy: upperQuery.includes('ORDER BY'),
      limit: upperQuery.includes('LIMIT'),
      estimated_complexity: 'low',
    };

    // Simple complexity estimation
    const score = complexity.joins * 2 + complexity.subqueries * 3 + complexity.aggregations;
    if (score > 10) {
      complexity.estimated_complexity = 'high';
    } else if (score > 5) {
      complexity.estimated_complexity = 'medium';
    }

    return complexity;
  }
}

export default QueryExecutor;