plugin-postgresql-connector/src/server/controllers/query.ts

NocoBase plugin for connecting to external PostgreSQL databases

import { Context } from '@nocobase/server';
import { ConnectionManager } from '../services/ConnectionManager';
import { QueryExecutor, QueryOptions } from '../services/QueryExecutor';
import { SavedQuery } from '../models/SavedQuery';
import Joi from 'joi';

const executeQuerySchema = Joi.object({
  connectionId: Joi.string().uuid().required(),
  query: Joi.string().min(1).required(),
  parameters: Joi.array().items(Joi.any()).default([]),
  options: Joi.object({
    timeout: Joi.number().integer().min(1000).max(300000).default(30000),
    maxRows: Joi.number().integer().min(1).max(10000).default(1000),
    formatQuery: Joi.boolean().default(false),
    includeMetadata: Joi.boolean().default(true),
  }).default({}),
});

const executeProcedureSchema = Joi.object({
  connectionId: Joi.string().uuid().required(),
  procedureName: Joi.string().min(1).required(),
  parameters: Joi.array().items(Joi.any()).default([]),
  options: Joi.object({
    timeout: Joi.number().integer().min(1000).max(300000).default(30000),
    formatQuery: Joi.boolean().default(false),
    includeMetadata: Joi.boolean().default(true),
  }).default({}),
});

const executeFunctionSchema = Joi.object({
  connectionId: Joi.string().uuid().required(),
  functionName: Joi.string().min(1).required(),
  parameters: Joi.array().items(Joi.any()).default([]),
  options: Joi.object({
    timeout: Joi.number().integer().min(1000).max(300000).default(30000),
    formatQuery: Joi.boolean().default(false),
    includeMetadata: Joi.boolean().default(true),
  }).default({}),
});

const getViewDataSchema = Joi.object({
  connectionId: Joi.string().uuid().required(),
  viewName: Joi.string().min(1).required(),
  limit: Joi.number().integer().min(1).max(10000).default(100),
  options: Joi.object({
    formatQuery: Joi.boolean().default(false),
    includeMetadata: Joi.boolean().default(true),
  }).default({}),
});

const getTableDataSchema = Joi.object({
  connectionId: Joi.string().uuid().required(),
  tableName: Joi.string().min(1).required(),
  limit: Joi.number().integer().min(1).max(10000).default(100),
  offset: Joi.number().integer().min(0).default(0),
  orderBy: Joi.string().optional(),
  options: Joi.object({
    formatQuery: Joi.boolean().default(false),
    includeMetadata: Joi.boolean().default(true),
  }).default({}),
});

const saveQuerySchema = Joi.object({
  connectionId: Joi.string().uuid().required(),
  name: Joi.string().min(1).max(255).required(),
  query: Joi.string().min(1).required(),
  queryType: Joi.string().valid('SELECT', 'INSERT', 'UPDATE', 'DELETE', 'PROCEDURE', 'FUNCTION', 'VIEW').required(),
  parameters: Joi.array().items(
    Joi.object({
      name: Joi.string().required(),
      value: Joi.any(),
      type: Joi.string().optional(),
    })
  ).default([]),
  description: Joi.string().max(1000).optional(),
  category: Joi.string().max(100).optional(),
  tags: Joi.array().items(Joi.string()).default([]),
});

export class QueryController {
  constructor(
    private connectionManager: ConnectionManager,
    private queryExecutor: QueryExecutor
  ) {}

  async execute(ctx: Context) {
    const { error, value } = executeQuerySchema.validate(ctx.request.body);
    if (error) {
      ctx.throw(400, `Validation error: ${error.details[0].message}`);
    }

    const { connectionId, query, parameters, options } = value;

    try {
      await this.validateConnection(connectionId);

      const queryAnalysis = this.queryExecutor.analyzeQuery(query);
      
      const startTime = Date.now();
      const result = await this.queryExecutor.executeQuery(connectionId, query, parameters, options);
      const totalTime = Date.now() - startTime;

      ctx.body = {
        success: true,
        data: {
          ...result,
          analysis: queryAnalysis,
          totalExecutionTime: totalTime,
        },
        meta: {
          executedAt: new Date().toISOString(),
          connectionId,
          queryLength: query.length,
          parameterCount: parameters.length,
        },
      };

      this.logQueryExecution(ctx, connectionId, query, result.executionTime, true);

    } catch (error) {
      this.logQueryExecution(ctx, connectionId, query, 0, false, error);
      ctx.throw(500, `Query execution failed: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  async executeProcedure(ctx: Context) {
    const { error, value } = executeProcedureSchema.validate(ctx.request.body);
    if (error) {
      ctx.throw(400, `Validation error: ${error.details[0].message}`);
    }

    const { connectionId, procedureName, parameters, options } = value;

    try {
      await this.validateConnection(connectionId);

      const result = await this.queryExecutor.executeProcedure(connectionId, procedureName, parameters, options);

      ctx.body = {
        success: true,
        data: result,
        meta: {
          executedAt: new Date().toISOString(),
          connectionId,
          procedureName,
          parameterCount: parameters.length,
        },
      };

      this.logQueryExecution(ctx, connectionId, `CALL ${procedureName}`, result.executionTime, true);

    } catch (error) {
      this.logQueryExecution(ctx, connectionId, `CALL ${procedureName}`, 0, false, error);
      ctx.throw(500, `Procedure execution failed: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  async executeFunction(ctx: Context) {
    const { error, value } = executeFunctionSchema.validate(ctx.request.body);
    if (error) {
      ctx.throw(400, `Validation error: ${error.details[0].message}`);
    }

    const { connectionId, functionName, parameters, options } = value;

    try {
      await this.validateConnection(connectionId);

      const result = await this.queryExecutor.executeFunction(connectionId, functionName, parameters, options);

      ctx.body = {
        success: true,
        data: result,
        meta: {
          executedAt: new Date().toISOString(),
          connectionId,
          functionName,
          parameterCount: parameters.length,
        },
      };

      this.logQueryExecution(ctx, connectionId, `SELECT ${functionName}()`, result.executionTime, true);

    } catch (error) {
      this.logQueryExecution(ctx, connectionId, `SELECT ${functionName}()`, 0, false, error);
      ctx.throw(500, `Function execution failed: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  async getViewData(ctx: Context) {
    const { error, value } = getViewDataSchema.validate({
      connectionId: ctx.query.connectionId,
      viewName: ctx.query.viewName,
      limit: ctx.query.limit ? parseInt(ctx.query.limit as string) : 100,
      options: ctx.query.options || {},
    });

    if (error) {
      ctx.throw(400, `Validation error: ${error.details[0].message}`);
    }

    const { connectionId, viewName, limit, options } = value;

    try {
      await this.validateConnection(connectionId);

      const result = await this.queryExecutor.getViewData(connectionId, viewName, limit, options);

      ctx.body = {
        success: true,
        data: result,
        meta: {
          executedAt: new Date().toISOString(),
          connectionId,
          viewName,
          limit,
        },
      };

    } catch (error) {
      ctx.throw(500, `Failed to fetch view data: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  async getTableData(ctx: Context) {
    const { error, value } = getTableDataSchema.validate({
      connectionId: ctx.query.connectionId,
      tableName: ctx.query.tableName,
      limit: ctx.query.limit ? parseInt(ctx.query.limit as string) : 100,
      offset: ctx.query.offset ? parseInt(ctx.query.offset as string) : 0,
      orderBy: ctx.query.orderBy,
      options: ctx.query.options || {},
    });

    if (error) {
      ctx.throw(400, `Validation error: ${error.details[0].message}`);
    }

    const { connectionId, tableName, limit, offset, orderBy, options } = value;

    try {
      await this.validateConnection(connectionId);

      const result = await this.queryExecutor.getTableData(connectionId, tableName, limit, offset, orderBy, options);

      ctx.body = {
        success: true,
        data: result,
        meta: {
          executedAt: new Date().toISOString(),
          connectionId,
          tableName,
          limit,
          offset,
          orderBy,
        },
      };

    } catch (error) {
      ctx.throw(500, `Failed to fetch table data: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  async saveQuery(ctx: Context) {
    const { error, value } = saveQuerySchema.validate(ctx.request.body);
    if (error) {
      ctx.throw(400, `Validation error: ${error.details[0].message}`);
    }

    const { connectionId, name, query, queryType, parameters, description, category, tags } = value;

    try {
      await this.validateConnection(connectionId);

      // Check if query name already exists for this connection
      const existingQuery = await SavedQuery.findOne({
        where: {
          connectionId,
          name,
          deletedAt: null,
        },
      });

      if (existingQuery) {
        ctx.throw(409, 'A query with this name already exists for this connection');
      }

      // Validate query by attempting to analyze it
      try {
        this.queryExecutor.analyzeQuery(query);
      } catch (analysisError) {
        ctx.throw(400, `Invalid query syntax: ${analysisError instanceof Error ? analysisError.message : 'Unknown error'}`);
      }

      const savedQuery = await SavedQuery.create({
        connectionId,
        name,
        query,
        queryType,
        parameters,
        description,
        category,
        tags,
        createdBy: ctx.state.currentUser?.id,
      });

      ctx.body = {
        success: true,
        data: {
          id: savedQuery.id,
          name: savedQuery.name,
          queryType: savedQuery.queryType,
          category: savedQuery.category,
          tags: savedQuery.tags,
          createdAt: savedQuery.createdAt,
        },
        message: 'Query saved successfully',
      };

    } catch (error) {
      if (error instanceof Error && error.message.includes('already exists')) {
        throw error;
      }
      ctx.throw(500, `Failed to save query: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  async getSavedQueries(ctx: Context) {
    const { connectionId, category, queryType, search, page = 1, limit = 20 } = ctx.query;

    try {
      const where: any = { deletedAt: null };

      if (connectionId) {
        where.connectionId = connectionId;
      }

      if (category) {
        where.category = category;
      }

      if (queryType) {
        where.queryType = queryType;
      }

      if (search) {
        where.name = { [require('sequelize').Op.iLike]: `%${search}%` };
      }

      const offset = (parseInt(page as string) - 1) * parseInt(limit as string);

      const { count, rows } = await SavedQuery.findAndCountAll({
        where,
        attributes: ['id', 'name', 'queryType', 'category', 'tags', 'description', 'lastExecutedAt', 'executionCount', 'createdAt', 'updatedAt'],
        order: [['updatedAt', 'DESC']],
        limit: parseInt(limit as string),
        offset,
      });

      ctx.body = {
        success: true,
        data: rows,
        meta: {
          total: count,
          page: parseInt(page as string),
          limit: parseInt(limit as string),
          totalPages: Math.ceil(count / parseInt(limit as string)),
        },
      };

    } catch (error) {
      ctx.throw(500, `Failed to fetch saved queries: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  async getSavedQuery(ctx: Context) {
    const { id } = ctx.params;

    try {
      const savedQuery = await SavedQuery.findOne({
        where: { id, deletedAt: null },
      });

      if (!savedQuery) {
        ctx.throw(404, 'Saved query not found');
      }

      ctx.body = {
        success: true,
        data: savedQuery,
      };

    } catch (error) {
      if (error instanceof Error && error.message.includes('not found')) {
        throw error;
      }
      ctx.throw(500, `Failed to fetch saved query: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  async deleteSavedQuery(ctx: Context) {
    const { id } = ctx.params;

    try {
      const savedQuery = await SavedQuery.findOne({
        where: { id, deletedAt: null },
      });

      if (!savedQuery) {
        ctx.throw(404, 'Saved query not found');
      }

      await savedQuery.update({
        deletedAt: new Date(),
        deletedBy: ctx.state.currentUser?.id,
      });

      ctx.body = {
        success: true,
        message: 'Query deleted successfully',
      };

    } catch (error) {
      if (error instanceof Error && error.message.includes('not found')) {
        throw error;
      }
      ctx.throw(500, `Failed to delete saved query: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  async executeSavedQuery(ctx: Context) {
    const { id } = ctx.params;
    const { parameters = [], options = {} } = ctx.request.body;

    try {
      const savedQuery = await SavedQuery.findOne({
        where: { id, deletedAt: null },
      });

      if (!savedQuery) {
        ctx.throw(404, 'Saved query not found');
      }

      await this.validateConnection(savedQuery.connectionId);

      let result;
      
      switch (savedQuery.queryType) {
        case 'SELECT':
        case 'INSERT':
        case 'UPDATE':
        case 'DELETE':
          result = await this.queryExecutor.executeQuery(savedQuery.connectionId, savedQuery.query, parameters, options);
          break;
        case 'PROCEDURE':
          const procName = this.extractProcedureName(savedQuery.query);
          result = await this.queryExecutor.executeProcedure(savedQuery.connectionId, procName, parameters, options);
          break;
        case 'FUNCTION':
          const funcName = this.extractFunctionName(savedQuery.query);
          result = await this.queryExecutor.executeFunction(savedQuery.connectionId, funcName, parameters, options);
          break;
        case 'VIEW':
          const viewName = this.extractViewName(savedQuery.query);
          result = await this.queryExecutor.getViewData(savedQuery.connectionId, viewName, options.maxRows || 100, options);
          break;
        default:
          result = await this.queryExecutor.executeQuery(savedQuery.connectionId, savedQuery.query, parameters, options);
      }

      // Update execution statistics
      await savedQuery.increment('executionCount');
      await savedQuery.update({ lastExecutedAt: new Date() });

      ctx.body = {
        success: true,
        data: result,
        meta: {
          executedAt: new Date().toISOString(),
          savedQueryId: id,
          savedQueryName: savedQuery.name,
          queryType: savedQuery.queryType,
        },
      };

      this.logQueryExecution(ctx, savedQuery.connectionId, savedQuery.query, result.executionTime, true);

    } catch (error) {
      if (error instanceof Error && error.message.includes('not found')) {
        throw error;
      }
      ctx.throw(500, `Failed to execute saved query: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  async getQueryStatistics(ctx: Context) {
    try {
      const stats = this.queryExecutor.getStatistics();
      
      ctx.body = {
        success: true,
        data: stats,
      };

    } catch (error) {
      ctx.throw(500, `Failed to get query statistics: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  async clearQueryCache(ctx: Context) {
    try {
      this.queryExecutor.clearCache();
      
      ctx.body = {
        success: true,
        message: 'Query cache cleared successfully',
      };

    } catch (error) {
      ctx.throw(500, `Failed to clear query cache: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  // Helper methods
  private async validateConnection(connectionId: string): Promise<void> {
    const { Connection } = require('../models/Connection');
    const connection = await Connection.findOne({
      where: { id: connectionId, isActive: true },
    });

    if (!connection) {
      throw new Error('Connection not found or inactive');
    }
  }

  private logQueryExecution(
    ctx: Context,
    connectionId: string,
    query: string,
    executionTime: number,
    success: boolean,
    error?: any
  ): void {
    const logEntry = {
      timestamp: new Date().toISOString(),
      userId: ctx.state.currentUser?.id,
      connectionId,
      query: query.substring(0, 1000),
      executionTime,
      success,
      error: error ? error.message : null,
      userAgent: ctx.get('User-Agent'),
      ip: ctx.ip,
    };

    console.log('Query Execution Log:', logEntry);
  }

  private extractProcedureName(query: string): string {
    const match = query.match(/CALL\s+([a-zA-Z_][a-zA-Z0-9_]*)/i);
    return match ? match[1] : '';
  }

  private extractFunctionName(query: string): string {
    const match = query.match(/SELECT\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*\(/i);
    return match ? match[1] : '';
  }

  private extractViewName(query: string): string {
    const match = query.match(/FROM\s+([a-zA-Z_][a-zA-Z0-9_]*)/i);
    return match ? match[1] : '';
  }
}

export default QueryController;