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

NocoBase plugin for connecting to external PostgreSQL databases

import { PoolClient } from 'pg';
import { ConnectionManager } from './ConnectionManager';

export interface TableInfo {
  table_name: string;
  table_schema: string;
  table_type: string;
  table_comment?: string;
  column_count?: number;
  row_count?: number;
}

export interface ColumnInfo {
  column_name: string;
  data_type: string;
  is_nullable: string;
  column_default?: string;
  character_maximum_length?: number;
  numeric_precision?: number;
  numeric_scale?: number;
  column_comment?: string;
  ordinal_position: number;
  is_primary_key: boolean;
  is_foreign_key: boolean;
  foreign_key_table?: string;
  foreign_key_column?: string;
}

export interface ViewInfo {
  view_name: string;
  view_schema: string;
  view_definition: string;
  is_updatable: string;
  check_option?: string;
}

export interface FunctionInfo {
  function_name: string;
  function_schema: string;
  routine_type: string;
  data_type?: string;
  routine_definition?: string;
  external_language?: string;
  parameter_count?: number;
  parameters?: Array<{
    parameter_name: string;
    data_type: string;
    parameter_mode: string;
  }>;
}

export interface IndexInfo {
  index_name: string;
  table_name: string;
  column_names: string[];
  is_unique: boolean;
  is_primary: boolean;
  index_type: string;
}

export interface SchemaStatistics {
  total_tables: number;
  total_views: number;
  total_functions: number;
  total_procedures: number;
  total_indexes: number;
  database_size: string;
  schemas: string[];
}

export interface DatabaseInfo {
  database_name: string;
  database_size: string;
  encoding: string;
  collation: string;
  connection_count: number;
  version: string;
  uptime?: string;
}

export class SchemaService {
  constructor(private connectionManager: ConnectionManager) {}

  async getDatabaseInfo(connectionId: string): Promise<DatabaseInfo> {
    const client = await this.connectionManager.getConnection(connectionId);
    
    try {
      // Get basic database info
      const dbInfoQuery = `
        SELECT 
          current_database() as database_name,
          pg_encoding_to_char(encoding) as encoding,
          datcollate as collation
        FROM pg_database 
        WHERE datname = current_database()
      `;
      
      const dbInfoResult = await client.query(dbInfoQuery);
      const dbInfo = dbInfoResult.rows[0];

      // Get database size
      const sizeQuery = `SELECT pg_size_pretty(pg_database_size(current_database())) as database_size`;
      const sizeResult = await client.query(sizeQuery);

      // Get PostgreSQL version
      const versionQuery = `SELECT version() as version`;
      const versionResult = await client.query(versionQuery);

      // Get connection count
      const connectionsQuery = `
        SELECT count(*) as connection_count 
        FROM pg_stat_activity 
        WHERE datname = current_database()
      `;
      const connectionsResult = await client.query(connectionsQuery);

      // Get uptime (PostgreSQL start time)
      const uptimeQuery = `
        SELECT date_trunc('second', now() - pg_postmaster_start_time()) as uptime
      `;
      const uptimeResult = await client.query(uptimeQuery);

      return {
        database_name: dbInfo.database_name,
        database_size: sizeResult.rows[0].database_size,
        encoding: dbInfo.encoding,
        collation: dbInfo.collation,
        connection_count: parseInt(connectionsResult.rows[0].connection_count),
        version: versionResult.rows[0].version,
        uptime: uptimeResult.rows[0].uptime,
      };
    } finally {
      client.release();
    }
  }

  async getSchemaStatistics(connectionId: string, schema: string = 'public'): Promise<SchemaStatistics> {
    const client = await this.connectionManager.getConnection(connectionId);
    
    try {
      // Get table count
      const tablesQuery = `
        SELECT COUNT(*) as count 
        FROM information_schema.tables 
        WHERE table_schema = $1 AND table_type = 'BASE TABLE'
      `;
      const tablesResult = await client.query(tablesQuery, [schema]);

      // Get view count
      const viewsQuery = `
        SELECT COUNT(*) as count 
        FROM information_schema.views 
        WHERE table_schema = $1
      `;
      const viewsResult = await client.query(viewsQuery, [schema]);

      // Get function count
      const functionsQuery = `
        SELECT COUNT(*) as count 
        FROM information_schema.routines 
        WHERE routine_schema = $1 AND routine_type = 'FUNCTION'
      `;
      const functionsResult = await client.query(functionsQuery, [schema]);

      // Get procedure count
      const proceduresQuery = `
        SELECT COUNT(*) as count 
        FROM information_schema.routines 
        WHERE routine_schema = $1 AND routine_type = 'PROCEDURE'
      `;
      const proceduresResult = await client.query(proceduresQuery, [schema]);

      // Get index count
      const indexesQuery = `
        SELECT COUNT(*) as count 
        FROM pg_indexes 
        WHERE schemaname = $1
      `;
      const indexesResult = await client.query(indexesQuery, [schema]);

      // Get database size
      const sizeQuery = `SELECT pg_size_pretty(pg_database_size(current_database())) as database_size`;
      const sizeResult = await client.query(sizeQuery);

      // Get all schemas
      const schemasQuery = `
        SELECT schema_name 
        FROM information_schema.schemata 
        WHERE schema_name NOT IN ('information_schema', 'pg_catalog', 'pg_toast')
        ORDER BY schema_name
      `;
      const schemasResult = await client.query(schemasQuery);

      return {
        total_tables: parseInt(tablesResult.rows[0].count),
        total_views: parseInt(viewsResult.rows[0].count),
        total_functions: parseInt(functionsResult.rows[0].count),
        total_procedures: parseInt(proceduresResult.rows[0].count),
        total_indexes: parseInt(indexesResult.rows[0].count),
        database_size: sizeResult.rows[0].database_size,
        schemas: schemasResult.rows.map(row => row.schema_name),
      };
    } finally {
      client.release();
    }
  }

  async getTables(connectionId: string, schema: string = 'public'): Promise<TableInfo[]> {
    const client = await this.connectionManager.getConnection(connectionId);
    
    try {
      const query = `
        SELECT 
          t.table_name,
          t.table_schema,
          t.table_type,
          obj_description(c.oid) as table_comment,
          (
            SELECT COUNT(*) 
            FROM information_schema.columns 
            WHERE table_schema = t.table_schema 
            AND table_name = t.table_name
          ) as column_count,
          (
            SELECT reltuples::bigint 
            FROM pg_class c2 
            JOIN pg_namespace n ON c2.relnamespace = n.oid 
            WHERE n.nspname = t.table_schema 
            AND c2.relname = t.table_name
          ) as row_count
        FROM information_schema.tables t
        LEFT JOIN pg_class c ON c.relname = t.table_name
        LEFT JOIN pg_namespace n ON c.relnamespace = n.oid AND n.nspname = t.table_schema
        WHERE t.table_schema = $1 
        AND t.table_type = 'BASE TABLE'
        ORDER BY t.table_name
      `;
      
      const result = await client.query(query, [schema]);
      return result.rows;
    } finally {
      client.release();
    }
  }

  async getTableColumns(connectionId: string, tableName: string, schema: string = 'public'): Promise<ColumnInfo[]> {
    const client = await this.connectionManager.getConnection(connectionId);
    
    try {
      const query = `
        SELECT 
          c.column_name,
          c.data_type,
          c.is_nullable,
          c.column_default,
          c.character_maximum_length,
          c.numeric_precision,
          c.numeric_scale,
          c.ordinal_position,
          col_description(pgc.oid, c.ordinal_position) as column_comment,
          CASE 
            WHEN pk.column_name IS NOT NULL THEN true 
            ELSE false 
          END as is_primary_key,
          CASE 
            WHEN fk.column_name IS NOT NULL THEN true 
            ELSE false 
          END as is_foreign_key,
          fk.foreign_table_name as foreign_key_table,
          fk.foreign_column_name as foreign_key_column
        FROM information_schema.columns c
        LEFT JOIN pg_class pgc ON pgc.relname = c.table_name
        LEFT JOIN pg_namespace pgn ON pgc.relnamespace = pgn.oid AND pgn.nspname = c.table_schema
        LEFT JOIN (
          SELECT 
            kcu.column_name,
            kcu.table_name,
            kcu.table_schema
          FROM information_schema.table_constraints tc
          JOIN information_schema.key_column_usage kcu ON tc.constraint_name = kcu.constraint_name
          WHERE tc.constraint_type = 'PRIMARY KEY'
        ) pk ON pk.column_name = c.column_name 
             AND pk.table_name = c.table_name 
             AND pk.table_schema = c.table_schema
        LEFT JOIN (
          SELECT 
            kcu.column_name,
            kcu.table_name,
            kcu.table_schema,
            ccu.table_name as foreign_table_name,
            ccu.column_name as foreign_column_name
          FROM information_schema.table_constraints tc
          JOIN information_schema.key_column_usage kcu ON tc.constraint_name = kcu.constraint_name
          JOIN information_schema.constraint_column_usage ccu ON tc.constraint_name = ccu.constraint_name
          WHERE tc.constraint_type = 'FOREIGN KEY'
        ) fk ON fk.column_name = c.column_name 
             AND fk.table_name = c.table_name 
             AND fk.table_schema = c.table_schema
        WHERE c.table_name = $1 AND c.table_schema = $2
        ORDER BY c.ordinal_position
      `;
      
      const result = await client.query(query, [tableName, schema]);
      return result.rows;
    } finally {
      client.release();
    }
  }

  async getViews(connectionId: string, schema: string = 'public'): Promise<ViewInfo[]> {
    const client = await this.connectionManager.getConnection(connectionId);
    
    try {
      const query = `
        SELECT 
          table_name as view_name,
          table_schema as view_schema,
          view_definition,
          is_updatable,
          check_option
        FROM information_schema.views 
        WHERE table_schema = $1
        ORDER BY table_name
      `;
      
      const result = await client.query(query, [schema]);
      return result.rows;
    } finally {
      client.release();
    }
  }

  async getFunctions(connectionId: string, schema: string = 'public'): Promise<FunctionInfo[]> {
    const client = await this.connectionManager.getConnection(connectionId);
    
    try {
      const query = `
        SELECT 
          r.routine_name as function_name,
          r.routine_schema as function_schema,
          r.routine_type,
          r.data_type,
          r.routine_definition,
          r.external_language,
          (
            SELECT COUNT(*) 
            FROM information_schema.parameters p 
            WHERE p.specific_name = r.specific_name
          ) as parameter_count
        FROM information_schema.routines r
        WHERE r.routine_schema = $1
        ORDER BY r.routine_name, r.routine_type
      `;
      
      const result = await client.query(query, [schema]);
      
      // Get parameters for each function
      for (const func of result.rows) {
        const paramsQuery = `
          SELECT 
            parameter_name,
            data_type,
            parameter_mode
          FROM information_schema.parameters 
          WHERE specific_name = (
            SELECT specific_name 
            FROM information_schema.routines 
            WHERE routine_name = $1 AND routine_schema = $2 AND routine_type = $3
          )
          ORDER BY ordinal_position
        `;
        
        const paramsResult = await client.query(paramsQuery, [
          func.function_name, 
          func.function_schema, 
          func.routine_type
        ]);
        
        func.parameters = paramsResult.rows;
      }
      
      return result.rows;
    } finally {
      client.release();
    }
  }

  async getIndexes(connectionId: string, tableName?: string, schema: string = 'public'): Promise<IndexInfo[]> {
    const client = await this.connectionManager.getConnection(connectionId);
    
    try {
      let query = `
        SELECT 
          i.indexname as index_name,
          i.tablename as table_name,
          array_agg(a.attname ORDER BY a.attnum) as column_names,
          idx.indisunique as is_unique,
          idx.indisprimary as is_primary,
          am.amname as index_type
        FROM pg_indexes i
        JOIN pg_class c ON c.relname = i.indexname
        JOIN pg_index idx ON idx.indexrelid = c.oid
        JOIN pg_class t ON t.oid = idx.indrelid
        JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = ANY(idx.indkey)
        JOIN pg_am am ON am.oid = c.relam
        WHERE i.schemaname = $1
      `;
      
      const params = [schema];
      
      if (tableName) {
        query += ` AND i.tablename = $2`;
        params.push(tableName);
      }
      
      query += `
        GROUP BY i.indexname, i.tablename, idx.indisunique, idx.indisprimary, am.amname
        ORDER BY i.tablename, i.indexname
      `;
      
      const result = await client.query(query, params);
      return result.rows;
    } finally {
      client.release();
    }
  }

  async searchObjects(
    connectionId: string, 
    searchTerm: string, 
    objectTypes: string[] = ['table', 'view', 'function'],
    schema: string = 'public'
  ): Promise<any[]> {
    const client = await this.connectionManager.getConnection(connectionId);
    
    try {
      const results: any[] = [];
      const searchPattern = `%${searchTerm.toLowerCase()}%`;

      // Search tables
      if (objectTypes.includes('table')) {
        const tablesQuery = `
          SELECT 
            table_name as name,
            'table' as type,
            table_schema as schema
          FROM information_schema.tables 
          WHERE table_schema = $1 
          AND table_type = 'BASE TABLE'
          AND LOWER(table_name) LIKE $2
        `;
        
        const tablesResult = await client.query(tablesQuery, [schema, searchPattern]);
        results.push(...tablesResult.rows);
      }

      // Search views
      if (objectTypes.includes('view')) {
        const viewsQuery = `
          SELECT 
            table_name as name,
            'view' as type,
            table_schema as schema
          FROM information_schema.views 
          WHERE table_schema = $1 
          AND LOWER(table_name) LIKE $2
        `;
        
        const viewsResult = await client.query(viewsQuery, [schema, searchPattern]);
        results.push(...viewsResult.rows);
      }

      // Search functions
      if (objectTypes.includes('function')) {
        const functionsQuery = `
          SELECT 
            routine_name as name,
            routine_type as type,
            routine_schema as schema
          FROM information_schema.routines 
          WHERE routine_schema = $1 
          AND LOWER(routine_name) LIKE $2
        `;
        
        const functionsResult = await client.query(functionsQuery, [schema, searchPattern]);
        results.push(...functionsResult.rows);
      }

      return results.sort((a, b) => a.name.localeCompare(b.name));
    } finally {
      client.release();
    }
  }

  async getTableRelationships(connectionId: string, tableName: string, schema: string = 'public'): Promise<any[]> {
    const client = await this.connectionManager.getConnection(connectionId);
    
    try {
      const query = `
        SELECT 
          tc.constraint_name,
          tc.constraint_type,
          kcu.column_name,
          ccu.table_name as foreign_table_name,
          ccu.column_name as foreign_column_name,
          rc.update_rule,
          rc.delete_rule
        FROM information_schema.table_constraints tc
        JOIN information_schema.key_column_usage kcu ON tc.constraint_name = kcu.constraint_name
        LEFT JOIN information_schema.constraint_column_usage ccu ON tc.constraint_name = ccu.constraint_name
        LEFT JOIN information_schema.referential_constraints rc ON tc.constraint_name = rc.constraint_name
        WHERE tc.table_name = $1 AND tc.table_schema = $2
        ORDER BY tc.constraint_type, kcu.ordinal_position
      `;
      
      const result = await client.query(query, [tableName, schema]);
      return result.rows;
    } finally {
      client.release();
    }
  }

  /**
   * Get table size and statistics
   */
  async getTableStatistics(connectionId: string, tableName: string, schema: string = 'public'): Promise<any> {
    const client = await this.connectionManager.getConnection(connectionId);
    
    try {
      const query = `
        SELECT 
          schemaname,
          tablename,
          attname,
          n_distinct,
          most_common_vals,
          most_common_freqs,
          histogram_bounds,
          correlation,
          most_common_elems,
          most_common_elem_freqs,
          elem_count_histogram
        FROM pg_stats 
        WHERE schemaname = $1 AND tablename = $2
        ORDER BY attname
      `;
      
      const statsResult = await client.query(query, [schema, tableName]);

      // Get table size information
      const sizeQuery = `
        SELECT 
          pg_size_pretty(pg_total_relation_size($1::regclass)) as total_size,
          pg_size_pretty(pg_relation_size($1::regclass)) as table_size,
          pg_size_pretty(pg_total_relation_size($1::regclass) - pg_relation_size($1::regclass)) as index_size
      `;
      
      const sizeResult = await client.query(sizeQuery, [`${schema}.${tableName}`]);

      return {
        statistics: statsResult.rows,
        size: sizeResult.rows[0],
      };
    } finally {
      client.release();
    }
  }
}

export default SchemaService;