aiquery2/api/app/controllers/report/clickhouse.js

'use strict';

const {
   CLICKHOUSE_IDENTIFIER_PATTERN,
} = require('./constants');
const { ensureObjectPayload } = require('./helpers');
const { ensureDataSource } = require('./repository');
const {
   validateStructuredQueryDslAgainstFields,
   buildStructuredPreviewSql,
} = require('./structuredQuery');

const ensureClickHouseIdentifier = (value, label = 'identifier') => {
   const normalized = String(value || '').trim();
   if (!normalized) throw `missing param: ${label}`;
   if (!CLICKHOUSE_IDENTIFIER_PATTERN.test(normalized)) throw `invalid param: ${label}`;
   return normalized;
};

const parseClickHouseRows = jsonResult => {
   if (Array.isArray(jsonResult)) return jsonResult.map(r => ({ ...r, id: r.name }));
   if (Array.isArray(jsonResult?.data)) return jsonResult.data.map(r => ({ ...r, id: r.name }));
   return [];
};

const resolveClickHouseDataSourceContextFromDataSource = async (ctx, dataSource, transaction = null) => {
   if (String(dataSource?.type || '').trim().toLowerCase() !== 'clickhouse') {
      throw 'data source type must be clickhouse';
   }
   const connectionConfig = ensureObjectPayload(dataSource.connectionConfig || {}, 'connectionConfig');
   const clientKey = String(connectionConfig.clientKey || '').trim();
   const database = String(connectionConfig.database || '').trim();
   if (!clientKey) throw 'missing param: connectionConfig.clientKey';
   if (!database) throw 'missing param: connectionConfig.database';

   const client = ctx.app?.fs?.clickHouse?.[clientKey];
   if (!client) throw `clickhouse client not found: ${clientKey}`;

   const configuredDbs = Array.isArray(ctx.app?.fs?.config?.clickHouse?.db)
      ? ctx.app.fs.config.clickHouse.db
      : [];
   const matched = configuredDbs.find(item => String(item?.name || '').trim() === clientKey);
   if (!matched) throw `clickhouse client not configured: ${clientKey}`;
   if (String(matched.db || '').trim() !== database) {
      throw `database mismatch for client: ${clientKey}`;
   }

   return {
      dataSource,
      clientKey,
      database,
      client,
      transaction,
   };
};

const resolveClickHouseDataSourceContext = async (ctx, dataSourceId, transaction = null) => {
   const dataSource = await ensureDataSource(ctx, dataSourceId, transaction);
   return resolveClickHouseDataSourceContextFromDataSource(ctx, dataSource, transaction);
};

const queryClickHouseRows = async (client, query, queryParams = {}) => {
   const normalizedQuery = String(query || '').trim().replace(/;+\s*$/, '');
   const finalQuery = normalizedQuery.replace(/\s+FORMAT\s+[A-Za-z0-9_]+\s*$/i, '');

   if (process.env.NODE_ENV === 'development') {
      console.log('finalQuery:', finalQuery);
   }

   const result = await client.query({
      query: finalQuery,
      query_params: queryParams,
      format: 'JSON',
   });
   const json = await result.json();
   return parseClickHouseRows(json);
};

const getClickHouseTableColumns = async (client, database, tableName) => {
   const query = `
      SELECT
         name,
         type,
         default_kind AS "defaultKind",
         default_expression AS "defaultExpression",
         is_in_primary_key AS "isInPrimaryKey",
         is_in_sorting_key AS "isInSortingKey",
         position
      FROM system.columns
      WHERE database = {database:String}
        AND table = {table:String}
      ORDER BY position ASC
   `;
   return queryClickHouseRows(client, query, { database, table: tableName });
};

const validateStructuredQueryDsl = async (dataSourceContext, rawDsl) => {
   const tableName = ensureClickHouseIdentifier(rawDsl?.table, 'table');
   const columns = await getClickHouseTableColumns(
      dataSourceContext.client,
      dataSourceContext.database,
      tableName
   );
   if (!columns.length) throw `table not found: ${tableName}`;

   const fieldSet = new Set(columns.map(item => item.name));
   return validateStructuredQueryDslAgainstFields(rawDsl, fieldSet, {
      tableParser: ensureClickHouseIdentifier,
      fieldParser: ensureClickHouseIdentifier,
      aliasParser: ensureClickHouseIdentifier,
   });
};

const listClickHouseTables = async (ctx, dataSource, transaction = null) => {
   const dataSourceContext = await resolveClickHouseDataSourceContextFromDataSource(ctx, dataSource, transaction);
   const query = `
      SELECT
         name,
         engine
      FROM system.tables
      WHERE database = {database:String}
      ORDER BY name ASC
   `;
   const rows = await queryClickHouseRows(
      dataSourceContext.client,
      query,
      { database: dataSourceContext.database }
   );

   return {
      dataSourceId: dataSource.id,
      clientKey: dataSourceContext.clientKey,
      database: dataSourceContext.database,
      rows,
   };
};

const getClickHouseTableColumnsPayload = async (ctx, dataSource, tableName, transaction = null) => {
   const safeTableName = ensureClickHouseIdentifier(tableName, 'tableName');
   const dataSourceContext = await resolveClickHouseDataSourceContextFromDataSource(ctx, dataSource, transaction);
   const rows = await getClickHouseTableColumns(
      dataSourceContext.client,
      dataSourceContext.database,
      safeTableName
   );
   if (!rows.length) throw `table not found: ${safeTableName}`;

   return {
      dataSourceId: dataSource.id,
      clientKey: dataSourceContext.clientKey,
      database: dataSourceContext.database,
      table: safeTableName,
      rows,
   };
};

const previewClickHouseQuery = async (ctx, dataSource, rawDsl, transaction = null) => {
   const dataSourceContext = await resolveClickHouseDataSourceContextFromDataSource(ctx, dataSource, transaction);
   const normalizedDsl = await validateStructuredQueryDsl(dataSourceContext, rawDsl);
   const queryBundle = buildStructuredPreviewSql(normalizedDsl);
   const rows = await queryClickHouseRows(
      dataSourceContext.client,
      queryBundle.query,
      queryBundle.queryParams
   );

   let total = null;
   if (!queryBundle.hasGroupBy) {
      const totalSql = `SELECT count() AS total FROM ${queryBundle.tableSql}${queryBundle.whereClause} FORMAT JSON`;
      const totalRows = await queryClickHouseRows(
         dataSourceContext.client,
         totalSql,
         queryBundle.queryParams
      );
      if (totalRows.length && totalRows[0]?.total !== undefined) {
         total = Number(totalRows[0].total);
      }
   }

   return {
      columns: queryBundle.resultColumns,
      rows,
      total,
      sqlPreview: queryBundle.sqlPreview,
   };
};

module.exports = {
   ensureClickHouseIdentifier,
   resolveClickHouseDataSourceContext,
   resolveClickHouseDataSourceContextFromDataSource,
   queryClickHouseRows,
   getClickHouseTableColumns,
   validateStructuredQueryDsl,
   listClickHouseTables,
   getClickHouseTableColumnsPayload,
   previewClickHouseQuery,
};