MyGoNavi/frontend/src/utils/dataSourceCapabilities.ts

import type { ConnectionConfig } from '../types';
import {
  isConnectionDataEditRestricted,
  isConnectionStructureEditRestricted,
} from './connectionReadOnly';
import { normalizeOceanBaseProtocol } from './oceanBaseProtocol';

type ConnectionLike = Pick<
  ConnectionConfig,
  'type' | 'driver' | 'oceanBaseProtocol' | 'readOnly' | 'protection'
> | null | undefined;

const normalizeDataSourceToken = (raw: string): string => {
  const normalized = String(raw || '').trim().toLowerCase();
  switch (normalized) {
    case 'doris':
      return 'diros';
    case 'starrocks':
      return 'starrocks';
    case 'postgresql':
      return 'postgres';
    case 'opengauss':
    case 'open_gauss':
    case 'open-gauss':
      return 'opengauss';
    case 'gaussdb':
    case 'gauss_db':
    case 'gauss-db':
      return 'gaussdb';
    case 'goldendb':
    case 'greatdb':
    case 'gdb':
      return 'goldendb';
    case 'dm':
      return 'dameng';
    case 'elastic':
    case 'elasticsearch':
      return 'elasticsearch';
    case 'chromadb':
    case 'chroma-db':
      return 'chroma';
    case 'qdrantdb':
    case 'qdrant-db':
      return 'qdrant';
    case 'rocketmq':
    case 'rocket-mq':
    case 'rocket_mq':
    case 'apache-rocketmq':
    case 'apache_rocketmq':
    case 'rmq':
      return 'rocketmq';
    case 'mqtt':
    case 'mqtts':
      return 'mqtt';
    case 'apache-iotdb':
    case 'apache_iotdb':
      return 'iotdb';
    case 'kafka':
    case 'apache-kafka':
    case 'apache_kafka':
      return 'kafka';
    case 'rabbitmq':
    case 'rabbit-mq':
    case 'rabbit_mq':
      return 'rabbitmq';
    case 'intersystems':
    case 'intersystemsiris':
    case 'inter-systems':
    case 'inter-systems-iris':
      return 'iris';
    default:
      return normalized;
  }
};

export const resolveDataSourceType = (config: ConnectionLike): string => {
  if (!config) return '';
  const type = normalizeDataSourceToken(String(config.type || ''));
  if (type === 'custom') {
    const driver = normalizeDataSourceToken(String(config.driver || ''));
    if (driver === 'oceanbase' && normalizeOceanBaseProtocol(config.oceanBaseProtocol) === 'oracle') {
      return 'oracle';
    }
    return driver || 'custom';
  }
  if (type === 'oceanbase' && normalizeOceanBaseProtocol(config.oceanBaseProtocol) === 'oracle') {
    return 'oracle';
  }
  return type;
};

export const shouldShowOceanBaseRowNumberColumn = (config: ConnectionLike): boolean => {
  if (!config) return false;
  const type = normalizeDataSourceToken(String(config.type || ''));
  const driver = normalizeDataSourceToken(String(config.driver || ''));
  return type === 'oceanbase' || driver === 'oceanbase';
};

const SQL_QUERY_EXPORT_TYPES = new Set([
  'mysql',
  'goldendb',
  'mariadb',
  'oceanbase',
  'diros',
  'starrocks',
  'sphinx',
  'postgres',
  'kingbase',
  'highgo',
  'vastbase',
  'opengauss',
  'gaussdb',
  'sqlserver',
  'iris',
  'sqlite',
  'duckdb',
  'oracle',
  'dameng',
  'tdengine',
  'clickhouse',
  'trino',
]);

const COPY_INSERT_TYPES = new Set([
  'mysql',
  'goldendb',
  'mariadb',
  'oceanbase',
  'diros',
  'starrocks',
  'sphinx',
  'postgres',
  'kingbase',
  'highgo',
  'vastbase',
  'opengauss',
  'gaussdb',
  'sqlserver',
  'iris',
  'sqlite',
  'duckdb',
  'oracle',
  'dameng',
  'tdengine',
  'clickhouse',
  'trino',
]);

const QUERY_EDITOR_DISABLED_TYPES = new Set(['redis']);
const FORCE_READ_ONLY_QUERY_TYPES = new Set(['tdengine', 'iotdb', 'clickhouse', 'rocketmq', 'mqtt', 'kafka', 'rabbitmq']);
const FORCE_READ_ONLY_STRUCTURE_DESIGNER_TYPES = new Set(['elasticsearch', 'mongodb', 'redis', 'iotdb']);
const MESSAGE_PUBLISH_TYPES = new Set(['rocketmq', 'mqtt', 'kafka', 'rabbitmq']);
const MANUAL_TOTAL_COUNT_TYPES = new Set(['duckdb', 'oracle', 'rocketmq', 'mqtt']);
const APPROXIMATE_TABLE_COUNT_TYPES = new Set(['duckdb', 'oracle']);
const APPROXIMATE_TOTAL_PAGE_TYPES = new Set(['duckdb']);

export type DataSourceCapabilities = {
  type: string;
  supportsQueryEditor: boolean;
  supportsSqlQueryExport: boolean;
  supportsCopyInsert: boolean;
  supportsCreateDatabase: boolean;
  supportsRenameDatabase: boolean;
  supportsDropDatabase: boolean;
  supportsMessagePublish: boolean;
  forceReadOnlyQueryResult: boolean;
  forceReadOnlyStructureDesigner: boolean;
  preferManualTotalCount: boolean;
  supportsApproximateTableCount: boolean;
  supportsApproximateTotalPages: boolean;
};

const CREATE_DATABASE_TYPES = new Set([
  'mysql',
  'goldendb',
  'mariadb',
  'oceanbase',
  'diros',
  'starrocks',
  'postgres',
  'kingbase',
  'highgo',
  'vastbase',
  'opengauss',
  'gaussdb',
  'sqlserver',
  'tdengine',
  'clickhouse',
]);

const RENAME_DATABASE_TYPES = new Set([
  'diros',
  'postgres',
  'kingbase',
  'highgo',
  'vastbase',
  'opengauss',
  'gaussdb',
]);

const DROP_DATABASE_TYPES = new Set([
  'mysql',
  'goldendb',
  'mariadb',
  'oceanbase',
  'diros',
  'starrocks',
  'postgres',
  'kingbase',
  'highgo',
  'vastbase',
  'opengauss',
  'gaussdb',
  'tdengine',
  'clickhouse',
]);

export const getDataSourceCapabilities = (config: ConnectionLike): DataSourceCapabilities => {
  const type = resolveDataSourceType(config);
  const dataEditRestricted = isConnectionDataEditRestricted(config);
  const structureEditRestricted = isConnectionStructureEditRestricted(config);
  return {
    type,
    supportsQueryEditor: !QUERY_EDITOR_DISABLED_TYPES.has(type),
    supportsSqlQueryExport: SQL_QUERY_EXPORT_TYPES.has(type),
    supportsCopyInsert: COPY_INSERT_TYPES.has(type),
    supportsCreateDatabase: !structureEditRestricted && CREATE_DATABASE_TYPES.has(type),
    supportsRenameDatabase: !structureEditRestricted && RENAME_DATABASE_TYPES.has(type),
    supportsDropDatabase: !structureEditRestricted && DROP_DATABASE_TYPES.has(type),
    supportsMessagePublish: !dataEditRestricted && MESSAGE_PUBLISH_TYPES.has(type),
    forceReadOnlyQueryResult: dataEditRestricted || FORCE_READ_ONLY_QUERY_TYPES.has(type),
    forceReadOnlyStructureDesigner:
      structureEditRestricted || FORCE_READ_ONLY_STRUCTURE_DESIGNER_TYPES.has(type),
    preferManualTotalCount: MANUAL_TOTAL_COUNT_TYPES.has(type),
    supportsApproximateTableCount: APPROXIMATE_TABLE_COUNT_TYPES.has(type),
    supportsApproximateTotalPages: APPROXIMATE_TOTAL_PAGE_TYPES.has(type),
  };
};