plugin-postgresql-connector/src/client/hooks/useQuery.ts

NocoBase plugin for connecting to external PostgreSQL databases

import { useState, useCallback, useEffect } from 'react';
import { useRequest } from '@nocobase/client';
import { message } from 'antd';

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

export interface SavedQuery {
  id: string;
  connectionId: string;
  name: string;
  query: string;
  queryType: 'SELECT' | 'INSERT' | 'UPDATE' | 'DELETE' | 'PROCEDURE' | 'FUNCTION';
  description?: string;
  category?: string;
  parameters?: any[];
  createdAt?: string;
  updatedAt?: string;
}

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

export const useQuery = (connectionId?: string) => {
  const [queryResult, setQueryResult] = useState<QueryResult | null>(null);
  const [savedQueries, setSavedQueries] = useState<SavedQuery[]>([]);
  const [queryHistory, setQueryHistory] = useState<string[]>([]);
  const [executionError, setExecutionError] = useState<string | null>(null);

  // Execute SQL query
  const { run: executeQuery, loading: executingQuery } = useRequest(
    (data: {
      connectionId: string;
      query: string;
      parameters?: any[];
      options?: QueryExecutionOptions;
    }) => ({
      url: '/postgresql-query/execute',
      method: 'POST',
      data,
    }),
    {
      manual: true,
      onSuccess: (data: any) => {
        setQueryResult(data.data);
        setExecutionError(null);
        addToHistory(data.data.query);
        message.success(`Query thực thi thành công! (${data.data.executionTime}ms)`);
      },
      onError: (error: any) => {
        setExecutionError(error.message);
        setQueryResult(null);
        message.error(`Query thất bại: ${error.message}`);
      },
    }
  );

  // Execute stored procedure
  const { run: executeProcedure, loading: executingProcedure } = useRequest(
    (data: {
      connectionId: string;
      procedureName: string;
      parameters?: any[];
    }) => ({
      url: '/postgresql-query/executeProcedure',
      method: 'POST',
      data,
    }),
    {
      manual: true,
      onSuccess: (data: any) => {
        setQueryResult(data.data);
        setExecutionError(null);
        message.success(`Procedure thực thi thành công! (${data.data.executionTime}ms)`);
      },
      onError: (error: any) => {
        setExecutionError(error.message);
        message.error(`Procedure thất bại: ${error.message}`);
      },
    }
  );

  // Execute function
  const { run: executeFunction, loading: executingFunction } = useRequest(
    (data: {
      connectionId: string;
      functionName: string;
      parameters?: any[];
    }) => ({
      url: '/postgresql-query/executeFunction',
      method: 'POST',
      data,
    }),
    {
      manual: true,
      onSuccess: (data: any) => {
        setQueryResult(data.data);
        setExecutionError(null);
        message.success(`Function thực thi thành công! (${data.data.executionTime}ms)`);
      },
      onError: (error: any) => {
        setExecutionError(error.message);
        message.error(`Function thất bại: ${error.message}`);
      },
    }
  );

  // Get table data
  const { run: getTableData, loading: loadingTableData } = useRequest(
    (data: {
      connectionId: string;
      tableName: string;
      limit?: number;
    }) => ({
      url: '/postgresql-query/getTableData',
      method: 'GET',
      params: data,
    }),
    {
      manual: true,
      onSuccess: (data: any) => {
        setQueryResult(data.data);
        setExecutionError(null);
      },
      onError: (error: any) => {
        setExecutionError(error.message);
        message.error(`Không thể tải dữ liệu table: ${error.message}`);
      },
    }
  );

  // Save query
  const { run: saveQuery, loading: savingQuery } = useRequest(
    (data: Omit<SavedQuery, 'id' | 'createdAt' | 'updatedAt'>) => ({
      url: '/postgresql-saved-queries',
      method: 'POST',
      data,
    }),
    {
      manual: true,
      onSuccess: () => {
        message.success('Query đã được lưu thành công!');
        if (connectionId) {
          fetchSavedQueries(connectionId);
        }
      },
      onError: (error: any) => {
        message.error(`Lưu query thất bại: ${error.message}`);
      },
    }
  );

  // Fetch saved queries
  const { run: fetchSavedQueries, loading: loadingSavedQueries } = useRequest(
    (connId: string) => ({
      url: '/postgresql-saved-queries',
      params: { connectionId: connId },
    }),
    {
      manual: true,
      onSuccess: (data: any) => {
        setSavedQueries(data.data || []);
      },
      onError: (error: any) => {
        message.error(`Không thể tải saved queries: ${error.message}`);
      },
    }
  );

  // Delete saved query
  const { run: deleteSavedQuery, loading: deletingQuery } = useRequest(
    (queryId: string) => ({
      url: `/postgresql-saved-queries/${queryId}`,
      method: 'DELETE',
    }),
    {
      manual: true,
      onSuccess: () => {
        message.success('Query đã được xóa!');
        if (connectionId) {
          fetchSavedQueries(connectionId);
        }
      },
      onError: (error: any) => {
        message.error(`Xóa query thất bại: ${error.message}`);
      },
    }
  );

  // Utility functions
  const addToHistory = useCallback((query: string) => {
    setQueryHistory(prev => {
      const newHistory = [query, ...prev.filter(q => q !== query)];
      return newHistory.slice(0, 50); // Keep last 50 queries
    });
  }, []);

  const clearHistory = useCallback(() => {
    setQueryHistory([]);
  }, []);

  const detectQueryType = useCallback((query: string): SavedQuery['queryType'] => {
    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.includes('CALL ')) return 'PROCEDURE';
    if (upperQuery.includes('SELECT ') && upperQuery.includes('(')) return 'FUNCTION';
    return 'SELECT';
  }, []);

  const formatQuery = useCallback((query: string): string => {
    // Basic SQL formatting
    return query
      .replace(/\s+/g, ' ')
      .replace(/,\s*/g, ',\n  ')
      .replace(/\bFROM\b/gi, '\nFROM')
      .replace(/\bWHERE\b/gi, '\nWHERE')
      .replace(/\bJOIN\b/gi, '\nJOIN')
      .replace(/\bORDER BY\b/gi, '\nORDER BY')
      .replace(/\bGROUP BY\b/gi, '\nGROUP BY')
      .replace(/\bHAVING\b/gi, '\nHAVING')
      .trim();
  }, []);

  const validateQuery = useCallback((query: string): string[] => {
    const errors: string[] = [];
    
    if (!query.trim()) {
      errors.push('Query không được để trống');
      return errors;
    }

    // Check for dangerous queries
    const upperQuery = query.toUpperCase();
    const dangerousKeywords = ['DROP', 'TRUNCATE', 'ALTER', 'CREATE'];
    
    dangerousKeywords.forEach(keyword => {
      if (upperQuery.includes(keyword)) {
        errors.push(`Query chứa từ khóa nguy hiểm: ${keyword}`);
      }
    });

    // Check for basic SQL syntax
    if (!upperQuery.match(/\b(SELECT|INSERT|UPDATE|DELETE|CALL)\b/)) {
      errors.push('Query không hợp lệ - thiếu keyword chính');
    }

    return errors;
  }, []);

  const exportQueryResult = useCallback((format: 'csv' | 'json' | 'excel') => {
    if (!queryResult?.rows?.length) {
      message.warning('Không có dữ liệu để export');
      return;
    }

    // This would integrate with export utilities
    message.info(`Export ${format.toUpperCase()} feature coming soon`);
  }, [queryResult]);

  const getQueryStats = useCallback(() => {
    if (!queryResult) return null;

    return {
      rowCount: queryResult.rowCount,
      columnCount: queryResult.fields?.length || 0,
      executionTime: queryResult.executionTime,
      dataSize: JSON.stringify(queryResult.rows).length,
    };
  }, [queryResult]);

  // Auto-fetch saved queries when connectionId changes
  useEffect(() => {
    if (connectionId) {
      fetchSavedQueries(connectionId);
    }
  }, [connectionId, fetchSavedQueries]);

  return {
    // State
    queryResult,
    savedQueries,
    queryHistory,
    executionError,

    // Loading states
    executingQuery,
    executingProcedure,
    executingFunction,
    loadingTableData,
    savingQuery,
    loadingSavedQueries,
    deletingQuery,

    // Actions
    executeQuery,
    executeProcedure,
    executeFunction,
    getTableData,
    saveQuery,
    fetchSavedQueries,
    deleteSavedQuery,

    // Utilities
    addToHistory,
    clearHistory,
    detectQueryType,
    formatQuery,
    validateQuery,
    exportQueryResult,
    getQueryStats,

    // State setters
    setQueryResult,
    setExecutionError,
  };
};

export default useQuery;