sql-query.ts

import { getClickhouseExternalClient } from "@/lib/clickhouse";
import { SmartRequestAuth } from "@/route-handlers/smart-request";
import { tool } from "ai";
import { z } from "zod";

export function createSqlQueryTool(auth: SmartRequestAuth | null, targetProjectId?: string | null) {
  if (auth == null) {
    // Return null or throw - analytics queries require authentication
    return null;
  }

  const projectId = targetProjectId ?? auth.tenancy.project.id;
  const branchId = targetProjectId ? "main" : auth.tenancy.branchId;

  return tool({
    description: "Run a ClickHouse SQL query against the project's analytics database. Only SELECT queries are allowed. Project filtering is automatic.",
    inputSchema: z.object({
      query: z
        .string()
        .describe("The ClickHouse SQL query to execute. Only SELECT queries are allowed. Always include LIMIT clause."),
    }),
    execute: async ({ query }: { query: string }) => {
      const client = getClickhouseExternalClient();
      return await client.query({
        query,
        clickhouse_settings: {
          SQL_project_id: projectId,
          SQL_branch_id: branchId,
          max_execution_time: 5,
          readonly: "1",
          allow_ddl: 0,
          max_result_rows: "10000",
          max_result_bytes: (10 * 1024 * 1024).toString(),
          result_overflow_mode: "throw",
        },
        format: "JSONEachRow",
      })
        .then(async (resultSet) => {
          const rows = await resultSet.json<Record<string, unknown>[]>();
          const serialized = JSON.stringify(rows);
          if (serialized.length > 50_000) {
            return {
              success: false as const,
              error:
                `Result too large: ${rows.length} rows, ${serialized.length} characters (limit 50000). ` +
                `To fix: ` +
                `(1) Use aggregation (COUNT, uniqExact, GROUP BY, topK, quantile) instead of fetching rows. ` +
                `(2) If you need rows, add a WHERE clause or reduce LIMIT. ` +
                `(3) Select only the columns you need — avoid the 'data' column on events unless essential.`,
              rowCount: rows.length,
              characters: serialized.length,
              columnsReturned: rows.length > 0 ? Object.keys(rows[0]) : [],
            };
          }
          return {
            success: true as const,
            rowCount: rows.length,
            result: rows,
          };
        })
        .catch((error: unknown) => ({
          success: false as const,
          error: error instanceof Error ? error.message : "Query failed",
        }));
    },
  });
}