import { executeTSQL, QueryError, type ClickHouseSettings, type ExecuteTSQLOptions, type FieldMappings, type TSQLQueryResult, } from "@internal/clickhouse"; import type { CustomerQuerySource } from "@trigger.dev/database"; import { calculateTimeBucketInterval, intervalToSeconds, type TableSchema, type WhereClauseCondition, } from "@internal/tsql"; import { z } from "zod"; import { prisma } from "~/db.server"; import { env } from "~/env.server"; import { clickhouseFactory } from "./clickhouse/clickhouseFactoryInstance.server"; import type { ClientType } from "./clickhouse/clickhouseFactory.server"; import { queryConcurrencyLimiter, DEFAULT_ORG_CONCURRENCY_LIMIT, GLOBAL_CONCURRENCY_LIMIT, } from "./queryConcurrencyLimiter.server"; import { getLimit } from "./platform.v3.server"; import { timeFilters, timeFilterFromTo } from "~/components/runs/v3/SharedFilters"; import parse from "parse-duration"; import { querySchemas, type QueryScope } from "~/v3/querySchemas"; export type { TSQLQueryResult, QueryScope }; const scopeToEnum = { organization: "ORGANIZATION", project: "PROJECT", environment: "ENVIRONMENT", } as const; /** * Default ClickHouse settings for query protection * Based on PostHog's HogQL settings to prevent expensive queries */ function getDefaultClickhouseSettings(): ClickHouseSettings { return { // Query execution limits max_execution_time: env.QUERY_CLICKHOUSE_MAX_EXECUTION_TIME, timeout_overflow_mode: "throw", max_memory_usage: String(env.QUERY_CLICKHOUSE_MAX_MEMORY_USAGE), // AST complexity limits to prevent extremely complex queries max_ast_elements: String(env.QUERY_CLICKHOUSE_MAX_AST_ELEMENTS), max_expanded_ast_elements: String(env.QUERY_CLICKHOUSE_MAX_EXPANDED_AST_ELEMENTS), // Memory management for GROUP BY operations max_bytes_before_external_group_by: String( env.QUERY_CLICKHOUSE_MAX_BYTES_BEFORE_EXTERNAL_GROUP_BY ), // Safety settings format_csv_allow_double_quotes: 0, readonly: "1", // Ensure queries are read-only }; } export type ExecuteQueryOptions = Omit< ExecuteTSQLOptions, "tableSchema" | "fieldMappings" | "enforcedWhereClause" | "whereClauseFallback" | "schema" > & { organizationId: string; projectId: string; environmentId: string; /** * The scope of the query - determines tenant isolation. Callers that take it from * a request body must cap it against the credential first; see `v3/queryScope.ts`. */ scope: QueryScope; period?: string | null; from?: string | null; to?: string | null; /** Filter to specific task identifiers */ taskIdentifiers?: string[]; /** Filter to specific queues */ queues?: string[]; /** Filter to specific response models */ responseModels?: string[]; /** Filter to specific prompt slugs */ promptSlugs?: string[]; /** Filter to specific prompt versions */ promptVersions?: number[]; /** Filter to specific operations (e.g. ai.generateText.doGenerate) */ operations?: string[]; /** Filter to specific providers (e.g. openai.responses) */ providers?: string[]; /** History options for saving query to billing/audit */ history?: { /** Where the query originated from */ source: CustomerQuerySource; /** User ID (optional, null for API calls) */ userId?: string | null; /** Skip saving to history (e.g., when impersonating) */ skip?: boolean; }; /** Custom per-org concurrency limit (overrides default) */ customOrgConcurrencyLimit?: number; /** * Set when the caller wrote `query` themselves, as on the public query API and * the query editor. ClickHouse rejecting their SQL is then their mistake, so it * is logged as a warning instead of raising an alert. Leave unset for TRQL we * generate, where the same rejection is a bug worth alerting on. */ userAuthoredQuery?: boolean; }; /** * Extended result type that includes the optional queryId when saved to history */ export type ExecuteQueryResult = | { success: true; result: T; queryId: string | null; periodClipped: number | null; maxQueryPeriod: number; timeRange: { from: Date; to: Date }; } | { success: false; error: Error }; /** Own-property flag tagged on the transient "query concurrency exceeded" rejection (retryable). */ const QUERY_CONCURRENCY_REJECTION_FLAG = "__queryConcurrencyRejection"; /** True for the transient concurrency-limit rejection — a stable signal callers can retry on. */ export function isQueryConcurrencyRejection(error: unknown): boolean { return ( typeof error === "object" && error !== null && (error as Record)[QUERY_CONCURRENCY_REJECTION_FLAG] === true ); } function floorToSeconds(date: Date, alignSeconds: number): Date { const ms = alignSeconds * 1000; return new Date(Math.floor(date.getTime() / ms) * ms); } /** * ClickHouse client a table's reads run on. A table can name its own pool (`queryClient`) so a * heavy read family lands on its own service; everything else shares the query pool. */ function resolveQueryClientType(schema: TableSchema | undefined): ClientType { switch (schema?.queryClient) { case "queueMetrics": return "queueMetrics"; default: return "query"; } } /** * Swap a table for one of its rollups when the query's bucket interval is at least the * rollup's granularity. The rollup has identical logical columns, so only the physical * table (and therefore rows read) changes. */ function resolveRollup( schema: TableSchema, timeRange: { from: Date; to: Date }, minBucketSeconds?: number ): TableSchema { if (!schema.rollups || schema.rollups.length === 0) { return schema; } const interval = calculateTimeBucketInterval( timeRange.from, timeRange.to, schema.timeBucketThresholds, minBucketSeconds ); const intervalSeconds = intervalToSeconds(interval); const best = [...schema.rollups] .sort((a, b) => b.minIntervalSeconds - a.minIntervalSeconds) .find((r) => r.minIntervalSeconds <= intervalSeconds); return best ? { ...schema, clickhouseName: best.clickhouseName } : schema; } export async function getDefaultPeriod(organizationId: string): Promise { const idealDefaultPeriodDays = 7; const maxQueryPeriod = await getLimit(organizationId, "queryPeriodDays", 30); if (maxQueryPeriod < idealDefaultPeriodDays) { return `${maxQueryPeriod}d`; } return `${idealDefaultPeriodDays}d`; } /** * Execute a TSQL query against ClickHouse with tenant isolation * Handles building tenant options, field mappings, and optionally saves to history * Returns [error, result, queryId] where queryId is the CustomerQuery ID if saved to history */ export async function executeQuery( options: ExecuteQueryOptions ): Promise>[1], null>>> { const { period, from, to, scope, organizationId, projectId, environmentId, taskIdentifiers, queues, responseModels, promptSlugs, promptVersions, operations, providers, history, customOrgConcurrencyLimit, ...baseOptions } = options; // Generate unique request ID for concurrency tracking const requestId = crypto.randomUUID(); const orgLimit = customOrgConcurrencyLimit ?? DEFAULT_ORG_CONCURRENCY_LIMIT; // Acquire concurrency slot const acquireResult = await queryConcurrencyLimiter.acquire({ key: projectId, requestId, keyLimit: orgLimit, globalLimit: GLOBAL_CONCURRENCY_LIMIT, }); if (!acquireResult.success) { const errorMessage = acquireResult.reason === "key_limit" ? `You've exceeded your query concurrency of ${orgLimit} for this project. Please try again later.` : "We're experiencing a lot of queries at the moment. Please try again later."; const error = new QueryError(errorMessage, { query: options.query }); // Stable marker so callers can retry on a transient concurrency rejection without // matching the message text (which is free to change). Object.assign(error, { [QUERY_CONCURRENCY_REJECTION_FLAG]: true }); return { success: false, error }; } // Detect which table the query targets to determine the time column // Each table schema declares its primary time column via timeConstraint const matchedSchema = querySchemas.find((s) => new RegExp(`\\bFROM\\s+${s.name}\\b`, "i").test(options.query) ); const timeColumn = matchedSchema?.timeConstraint ?? "triggered_at"; // Build time filter fallback for the table's time column const defaultPeriod = await getDefaultPeriod(organizationId); const timeFilter = timeFilters({ period: period ?? undefined, from: from ?? undefined, to: to ?? undefined, defaultPeriod, }); // Align the time bounds so repeated auto-refresh queries produce identical query // params and can share ClickHouse query-cache entries (params are part of the key). const alignSeconds = matchedSchema?.queryCache?.alignSeconds; if (alignSeconds) { if (timeFilter.from) timeFilter.from = floorToSeconds(timeFilter.from, alignSeconds); if (timeFilter.to) timeFilter.to = floorToSeconds(timeFilter.to, alignSeconds); } // Calculate the effective "from" date the user is requesting (for period clipping check) // This is null only when the user specifies just a "to" date (rare case) let requestedFromDate: Date | null = null; if (timeFilter.from) { requestedFromDate = new Date(timeFilter.from); } else if (!timeFilter.to) { // Period specified (or default) - calculate from now const periodMs = parse(timeFilter.period ?? defaultPeriod) ?? 7 * 24 * 60 * 60 * 1000; requestedFromDate = new Date(Date.now() - periodMs); if (alignSeconds) { requestedFromDate = floorToSeconds(requestedFromDate, alignSeconds); } } // Build the fallback WHERE condition based on what the user specified let timeFallback: WhereClauseCondition; if (timeFilter.from && timeFilter.to) { timeFallback = { op: "between", low: timeFilter.from, high: timeFilter.to }; } else if (timeFilter.from) { timeFallback = { op: "gte", value: timeFilter.from }; } else if (timeFilter.to) { timeFallback = { op: "lte", value: timeFilter.to }; } else { timeFallback = { op: "gte", value: requestedFromDate! }; } const maxQueryPeriod = await getLimit(organizationId, "queryPeriodDays", 30); let maxQueryPeriodDate = new Date(Date.now() - maxQueryPeriod * 24 * 60 * 60 * 1000); if (alignSeconds) { maxQueryPeriodDate = floorToSeconds(maxQueryPeriodDate, alignSeconds); } // Check if the requested time period exceeds the plan limit const periodClipped = requestedFromDate !== null && requestedFromDate < maxQueryPeriodDate; // Force tenant isolation and time period limits // Global tables (no tenantColumns) skip tenant isolation — they contain anonymized cross-tenant data const isGlobalTable = matchedSchema != null && !matchedSchema.tenantColumns; const enforcedWhereClause = { ...(isGlobalTable ? {} : { organization_id: { op: "eq", value: organizationId }, project_id: scope === "project" || scope === "environment" ? { op: "eq", value: projectId } : undefined, environment_id: scope === "environment" ? { op: "eq", value: environmentId } : undefined, }), [timeColumn]: { op: "gte", value: maxQueryPeriodDate }, // Optional filters for tasks and queues task_identifier: taskIdentifiers && taskIdentifiers.length > 0 ? { op: "in", values: taskIdentifiers } : undefined, queue: queues && queues.length > 0 ? { op: "in", values: queues } : undefined, response_model: responseModels && responseModels.length > 0 ? { op: "in", values: responseModels } : undefined, prompt_slug: promptSlugs && promptSlugs.length > 0 ? { op: "in", values: promptSlugs } : undefined, prompt_version: promptVersions && promptVersions.length > 0 ? { op: "in", values: promptVersions } : undefined, operation_id: operations && operations.length > 0 ? { op: "in", values: operations } : undefined, gen_ai_system: providers && providers.length > 0 ? { op: "in", values: providers } : undefined, } satisfies Record; // Compute the effective time range for timeBucket() interval calculation const timeRange = timeFilterFromTo({ period: period ?? undefined, from: from ?? undefined, to: to ?? undefined, defaultPeriod, }); if (alignSeconds) { timeRange.from = floorToSeconds(timeRange.from, alignSeconds); timeRange.to = floorToSeconds(timeRange.to, alignSeconds); } try { // Build field mappings for project_ref → project_id and environment_id → slug translation const projects = await prisma.project.findMany({ where: { organizationId }, select: { id: true, externalRef: true }, }); const environments = await prisma.runtimeEnvironment.findMany({ where: { project: { organizationId } }, select: { id: true, slug: true }, }); const fieldMappings: FieldMappings = { project: Object.fromEntries(projects.map((p) => [p.id, p.externalRef])), environment: Object.fromEntries(environments.map((e) => [e.id, e.slug])), }; const queryClickhouse = await clickhouseFactory.getClickhouseForOrganization( organizationId, resolveQueryClientType(matchedSchema) ); // Serve coarse-bucket queries from the table's rollup when one qualifies. const effectiveSchemas = matchedSchema?.rollups ? querySchemas.map((s) => s === matchedSchema ? resolveRollup(s, timeRange, baseOptions.minBucketSeconds) : s ) : querySchemas; const queryCacheSettings: ClickHouseSettings = matchedSchema?.queryCache ? { use_query_cache: 1, query_cache_ttl: matchedSchema.queryCache.ttlSeconds } : {}; const result = await executeTSQL(queryClickhouse.reader, { ...baseOptions, schema: z.record(z.any()), tableSchema: effectiveSchemas, transformValues: true, enforcedWhereClause, fieldMappings, whereClauseFallback: { [timeColumn]: timeFallback, }, timeRange, clickhouseSettings: { ...getDefaultClickhouseSettings(), ...queryCacheSettings, ...baseOptions.clickhouseSettings, // Allow caller overrides if needed readonly: "1", // Not overridable: every query through here is read-only. }, querySettings: { maxRows: env.QUERY_CLICKHOUSE_MAX_RETURNED_ROWS, ...baseOptions.querySettings, // Allow caller overrides if needed }, }); // If query failed, return early with no queryId if (result[0] !== null) { return { success: false, error: result[0] }; } let queryId: string | null = null; // If query succeeded and history options provided, save to history // Skip history for EXPLAIN queries (admin debugging) and when explicitly skipped (e.g., impersonating) if (history && !history.skip && !baseOptions.explain) { // Check if this query is the same as the last one saved (avoid duplicate history entries) const lastQuery = await prisma.customerQuery.findFirst({ where: { organizationId, source: history.source, userId: history.userId ?? null, }, orderBy: { createdAt: "desc" }, select: { id: true, query: true, scope: true, filterPeriod: true, filterFrom: true, filterTo: true, }, }); // Save the effective period used for the query (timeFilters() handles defaults) // Only save period if no custom from/to range was specified const historyTimeFilter = { period: timeFilter.from || timeFilter.to ? undefined : timeFilter.period, from: timeFilter.from, to: timeFilter.to, }; const isDuplicate = lastQuery && lastQuery.query === options.query && lastQuery.scope === scopeToEnum[scope] && lastQuery.filterPeriod === (timeFilter?.period ?? null) && lastQuery.filterFrom?.getTime() === (timeFilter?.from?.getTime() ?? undefined) && lastQuery.filterTo?.getTime() === (timeFilter?.to?.getTime() ?? undefined); if (isDuplicate && lastQuery) { // Return the existing query's ID for duplicate queries queryId = lastQuery.id; } else { const created = await prisma.customerQuery.create({ data: { query: options.query, scope: scopeToEnum[scope], stats: { ...result[1].stats }, source: history.source, organizationId, projectId: scope === "project" || scope === "environment" ? projectId : null, environmentId: scope === "environment" ? environmentId : null, userId: history.userId ?? null, filterPeriod: historyTimeFilter?.period ?? null, filterFrom: historyTimeFilter?.from ?? null, filterTo: historyTimeFilter?.to ?? null, }, }); queryId = created.id; } } return { success: true, result: result[1], queryId, periodClipped: periodClipped ? maxQueryPeriod : null, maxQueryPeriod, timeRange, }; } finally { // Always release the concurrency slot await queryConcurrencyLimiter.release({ key: projectId, requestId, }); } }