feat: migrate rest of span queries to clickhouse

Author: geclos

apps/web/src/app/api/projects/[projectId]/commits/[commitUuid]/documents/[documentUuid]/traces/aggregations/route.ts

@@ -34,6 +34,7 @@ export const GET = errorHandler(
         .then((r) => r.unwrap())
 
       const result = await computeDocumentTracesAggregations({
+        workspaceId: workspace.id,
         documentUuid: document.documentUuid,
         commitUuid: headCommit?.uuid === commitUuid ? undefined : commitUuid,
       }).then((r) => r.unwrap())

apps/web/src/app/api/projects/[projectId]/commits/[commitUuid]/documents/[documentUuid]/traces/daily-count/route.ts

@@ -38,6 +38,7 @@ export const GET = errorHandler(
         .then((r) => r.unwrap())
 
       const result = await computeDocumentTracesDailyCount({
+        workspaceId: workspace.id,
         documentUuid: document.documentUuid,
         commitUuid: headCommit?.uuid === commitUuid ? undefined : commitUuid,
         days,

packages/core/src/data-access/conversations/fetchConversation.ts

@@ -5,6 +5,10 @@ import { OkType, Result } from '../../lib/Result'
 import { spans } from '../../schema/models/spans'
 import { Workspace } from '../../schema/models/types/Workspace'
 import { conversationAggregateFields } from './shared'
+import { isFeatureEnabledByName } from '../../services/workspaceFeatures/isFeatureEnabledByName'
+import { fetchConversation as chFetchConversation } from '../../queries/clickhouse/spans/fetchConversation'
+
+const CLICKHOUSE_SPANS_READ_FLAG = 'clickhouse-spans-read'
 
 export type Conversation = OkType<typeof fetchConversation>
 
@@ -20,6 +24,28 @@ export async function fetchConversation(
   },
   db = database,
 ) {
+  const clickhouseEnabledResult = await isFeatureEnabledByName(
+    workspace.id,
+    CLICKHOUSE_SPANS_READ_FLAG,
+    db,
+  )
+  const shouldUseClickHouse =
+    clickhouseEnabledResult.ok && clickhouseEnabledResult.value
+
+  if (shouldUseClickHouse) {
+    const result = await chFetchConversation({
+      workspaceId: workspace.id,
+      documentLogUuid,
+      documentUuid,
+    })
+
+    if (!result) {
+      return Result.error(new NotFoundError('Conversation not found'))
+    }
+
+    return Result.ok(result)
+  }
+
   const conditions = [
     eq(spans.workspaceId, workspace.id),
     eq(spans.documentLogUuid, documentLogUuid),

packages/core/src/data-access/conversations/fetchConversations.ts

@@ -21,6 +21,10 @@ import {
   shouldFallbackToAllTime,
 } from '../../services/spans/defaultCreatedAtWindow'
 import { conversationAggregateFields } from './shared'
+import { isFeatureEnabledByName } from '../../services/workspaceFeatures/isFeatureEnabledByName'
+import { fetchConversations as chFetchConversations } from '../../queries/clickhouse/spans/fetchConversations'
+
+const CLICKHOUSE_SPANS_READ_FLAG = 'clickhouse-spans-read'
 
 export type ConversationFilters = {
   commitUuids: string[]
@@ -186,6 +190,58 @@ export async function fetchConversations(
   }: FetchConversationsParams,
   db = database,
 ) {
+  const clickhouseEnabledResult = await isFeatureEnabledByName(
+    workspace.id,
+    CLICKHOUSE_SPANS_READ_FLAG,
+    db,
+  )
+  const shouldUseClickHouse =
+    clickhouseEnabledResult.ok && clickhouseEnabledResult.value
+
+  if (shouldUseClickHouse) {
+    const normalizedCreatedAt = normalizeCreatedAtRange(filters.createdAt)
+    const defaultCreatedAt = applyDefaultSpansCreatedAtRange({
+      createdAt: normalizedCreatedAt,
+      hasCursor: Boolean(from),
+    })
+
+    const queryParams = {
+      workspaceId: workspace.id,
+      documentUuid,
+      filters,
+      from,
+      limit,
+    }
+
+    const firstPage = await chFetchConversations({
+      ...queryParams,
+      createdAt: defaultCreatedAt,
+    })
+
+    if (
+      !shouldFallbackToAllTime({
+        hasCursor: Boolean(from),
+        normalizedCreatedAt,
+        itemCount: firstPage.items.length,
+      })
+    ) {
+      return Result.ok<FetchConversationsResult>({
+        ...firstPage,
+        didFallbackToAllTime: undefined,
+      })
+    }
+
+    const allTime = await chFetchConversations({
+      ...queryParams,
+      createdAt: undefined,
+    })
+
+    return Result.ok<FetchConversationsResult>({
+      ...allTime,
+      didFallbackToAllTime: true,
+    })
+  }
+
   const normalizedCreatedAt = normalizeCreatedAtRange(filters.createdAt)
   const defaultCreatedAt = applyDefaultSpansCreatedAtRange({
     createdAt: normalizedCreatedAt,

packages/core/src/queries/clickhouse/spans/computeDocumentTracesAggregations.ts

@@ -0,0 +1,114 @@
+import { clickhouseClient } from '../../../client/clickhouse'
+import { SPANS_TABLE } from '../../../clickhouse/models/spans'
+import { TracesAggregations } from '../../../schema/models/types/Span'
+
+export async function computeDocumentTracesAggregations({
+  workspaceId,
+  documentUuid,
+  commitUuid,
+}: {
+  workspaceId: number
+  documentUuid: string
+  commitUuid?: string
+}): Promise<TracesAggregations> {
+  const params: Record<string, unknown> = {
+    workspaceId,
+    documentUuid,
+  }
+
+  const countConditions = [
+    `workspace_id = {workspaceId: UInt64}`,
+    `document_uuid = {documentUuid: String}`,
+  ]
+
+  const aggregationConditions = [
+    `workspace_id = {workspaceId: UInt64}`,
+    `document_uuid = {documentUuid: String}`,
+    `type = 'completion'`,
+  ]
+
+  if (commitUuid) {
+    params.commitUuid = commitUuid
+    countConditions.push(`commit_uuid = {commitUuid: String}`)
+    aggregationConditions.push(`commit_uuid = {commitUuid: String}`)
+  }
+
+  // Query for total count of distinct traces
+  const countResult = await clickhouseClient().query({
+    query: `
+      SELECT count(DISTINCT trace_id) AS total_count
+      FROM ${SPANS_TABLE}
+      WHERE ${countConditions.join(' AND ')}
+    `,
+    format: 'JSONEachRow',
+    query_params: params,
+  })
+
+  const countRows = await countResult.json<{ total_count: string }>()
+  const totalCount = Number(countRows[0]?.total_count ?? 0)
+
+  // Query for aggregations with median calculations using quantile(0.5)
+  const aggregationResult = await clickhouseClient().query({
+    query: `
+      SELECT
+        coalesce(
+          sum(
+            coalesce(tokens_prompt, 0) +
+            coalesce(tokens_cached, 0) +
+            coalesce(tokens_reasoning, 0) +
+            coalesce(tokens_completion, 0)
+          ),
+          0
+        ) AS total_tokens,
+        coalesce(sum(cost), 0) AS total_cost_in_millicents,
+        coalesce(
+          avg(
+            coalesce(tokens_prompt, 0) +
+            coalesce(tokens_cached, 0) +
+            coalesce(tokens_reasoning, 0) +
+            coalesce(tokens_completion, 0)
+          ),
+          0
+        ) AS average_tokens,
+        coalesce(avg(cost), 0) AS average_cost_in_millicents,
+        coalesce(quantile(0.5)(cost), 0) AS median_cost_in_millicents,
+        coalesce(avg(duration_ms), 0) AS average_duration,
+        coalesce(quantile(0.5)(duration_ms), 0) AS median_duration
+      FROM ${SPANS_TABLE}
+      WHERE ${aggregationConditions.join(' AND ')}
+    `,
+    format: 'JSONEachRow',
+    query_params: params,
+  })
+
+  const aggRows = await aggregationResult.json<{
+    total_tokens: string
+    total_cost_in_millicents: string
+    average_tokens: string
+    average_cost_in_millicents: string
+    median_cost_in_millicents: string
+    average_duration: string
+    median_duration: string
+  }>()
+
+  const row = aggRows[0] ?? {
+    total_tokens: '0',
+    total_cost_in_millicents: '0',
+    average_tokens: '0',
+    average_cost_in_millicents: '0',
+    median_cost_in_millicents: '0',
+    average_duration: '0',
+    median_duration: '0',
+  }
+
+  return {
+    totalCount,
+    totalTokens: Number(row.total_tokens),
+    totalCostInMillicents: Number(row.total_cost_in_millicents),
+    averageTokens: Number(row.average_tokens),
+    averageCostInMillicents: Number(row.average_cost_in_millicents),
+    medianCostInMillicents: Number(row.median_cost_in_millicents),
+    averageDuration: Number(row.average_duration),
+    medianDuration: Number(row.median_duration),
+  }
+}

packages/core/src/queries/clickhouse/spans/computeDocumentTracesDailyCount.ts

@@ -0,0 +1,62 @@
+import { subDays } from 'date-fns'
+import { clickhouseClient } from '../../../client/clickhouse'
+import { SPANS_TABLE } from '../../../clickhouse/models/spans'
+import { toClickHouseDateTime } from '../../../clickhouse/insert'
+
+export type DailyCount = {
+  date: string
+  count: number
+}
+
+export async function computeDocumentTracesDailyCount({
+  workspaceId,
+  documentUuid,
+  commitUuid,
+  days = 30,
+}: {
+  workspaceId: number
+  documentUuid: string
+  commitUuid?: string
+  days?: number
+}): Promise<DailyCount[]> {
+  const now = new Date()
+  const startDate = subDays(now, days)
+
+  const params: Record<string, unknown> = {
+    workspaceId,
+    documentUuid,
+    startDate: toClickHouseDateTime(startDate),
+  }
+
+  const conditions = [
+    `workspace_id = {workspaceId: UInt64}`,
+    `document_uuid = {documentUuid: String}`,
+    `started_at >= {startDate: DateTime64(6, 'UTC')}`,
+  ]
+
+  if (commitUuid) {
+    params.commitUuid = commitUuid
+    conditions.push(`commit_uuid = {commitUuid: String}`)
+  }
+
+  const result = await clickhouseClient().query({
+    query: `
+      SELECT
+        toDate(started_at) AS date,
+        count(DISTINCT trace_id) AS count
+      FROM ${SPANS_TABLE}
+      WHERE ${conditions.join(' AND ')}
+      GROUP BY date
+      ORDER BY date
+    `,
+    format: 'JSONEachRow',
+    query_params: params,
+  })
+
+  const rows = await result.json<{ date: string; count: string }>()
+
+  return rows.map((row) => ({
+    date: row.date,
+    count: Number(row.count),
+  }))
+}