chore: migrate last span queries from postgres to clickhouse (feature flagged)

Author: geclos

packages/core/src/data-access/traces/countByDocument.ts

@@ -5,6 +5,10 @@ import { Workspace } from '../../schema/models/types/Workspace'
 import { Commit } from '../../schema/models/types/Commit'
 import { CommitsRepository } from '../../repositories'
 import { LogSources } from '@latitude-data/constants'
+import { isFeatureEnabledByName } from '../../services/workspaceFeatures/isFeatureEnabledByName'
+import { countDistinctTracesByDocument } from '../../queries/clickhouse/spans/countByDocument'
+
+const CLICKHOUSE_SPANS_READ_FLAG = 'clickhouse-spans-read'
 
 export async function countTracesByDocument(
   {
@@ -20,8 +24,26 @@ export async function countTracesByDocument(
   },
   db = database,
 ) {
+  const clickhouseEnabledResult = await isFeatureEnabledByName(
+    workspace.id,
+    CLICKHOUSE_SPANS_READ_FLAG,
+    db,
+  )
+  const shouldUseClickHouse =
+    clickhouseEnabledResult.ok && clickhouseEnabledResult.value
+
   const commitsRepo = new CommitsRepository(workspace.id, db)
   const commits = await commitsRepo.getCommitsHistory({ commit })
+
+  if (shouldUseClickHouse) {
+    return countDistinctTracesByDocument({
+      workspaceId: workspace.id,
+      documentUuid,
+      commitUuids: commits.map((c) => c.uuid),
+      logSources,
+    })
+  }
+
   let filters = [
     eq(spans.workspaceId, workspace.id),
     eq(spans.documentUuid, documentUuid),

packages/core/src/data-access/traces/hasProductionTraces.ts

@@ -2,6 +2,10 @@ import { and, eq, inArray, sql, SQL } from 'drizzle-orm'
 import { database } from '../../client'
 import { MAIN_SPAN_TYPES, RUN_SOURCES, RunSourceGroup } from '../../constants'
 import { spans } from '../../schema/models/spans'
+import { isFeatureEnabledByName } from '../../services/workspaceFeatures/isFeatureEnabledByName'
+import { hasProductionTraces as chHasProductionTraces } from '../../queries/clickhouse/spans/hasProductionTraces'
+
+const CLICKHOUSE_SPANS_READ_FLAG = 'clickhouse-spans-read'
 
 export async function hasProductionTraces(
   {
@@ -13,6 +17,18 @@ export async function hasProductionTraces(
   },
   db = database,
 ): Promise<boolean> {
+  const clickhouseEnabledResult = await isFeatureEnabledByName(
+    workspaceId,
+    CLICKHOUSE_SPANS_READ_FLAG,
+    db,
+  )
+  const shouldUseClickHouse =
+    clickhouseEnabledResult.ok && clickhouseEnabledResult.value
+
+  if (shouldUseClickHouse) {
+    return chHasProductionTraces({ workspaceId, projectId })
+  }
+
   const conditions: SQL<unknown>[] = [
     eq(spans.workspaceId, workspaceId),
     inArray(spans.type, Array.from(MAIN_SPAN_TYPES)),

packages/core/src/data-access/weeklyEmail/logs/index.ts

@@ -4,6 +4,7 @@ import {
   countDistinct,
   desc,
   eq,
+  inArray,
   isNotNull,
   sql,
 } from 'drizzle-orm'
@@ -16,6 +17,13 @@ import { Workspace } from '../../../schema/models/types/Workspace'
 import { SureDateRange } from '../../../constants'
 import { getDateRangeOrLastWeekRange } from '../utils'
 import { hasProductionTraces } from '../../traces/hasProductionTraces'
+import { isFeatureEnabledByName } from '../../../services/workspaceFeatures/isFeatureEnabledByName'
+import {
+  getGlobalLogsStats as chGetGlobalLogsStats,
+  getTopProjectsLogsStats as chGetTopProjectsLogsStats,
+} from '../../../queries/clickhouse/spans/weeklyEmailLogs'
+
+const CLICKHOUSE_SPANS_READ_FLAG = 'clickhouse-spans-read'
 
 async function getGlobalLogsStats(
   {
@@ -26,7 +34,22 @@ async function getGlobalLogsStats(
     range: SureDateRange
   },
   db = database,
+  useClickHouse = false,
 ) {
+  if (useClickHouse) {
+    const stats = await chGetGlobalLogsStats({
+      workspaceId: workspace.id,
+      from: range.from,
+      to: range.to,
+    })
+
+    return {
+      logsCount: stats.logsCount,
+      tokensSpent: stats.tokensSpent,
+      tokensCost: stats.totalCostInMillicents / 100000,
+    }
+  }
+
   const logsCountResult = await db
     .select({ count: countDistinct(spans.traceId) })
     .from(spans)
@@ -71,7 +94,42 @@ async function getTopProjectsLogsStats(
     projectsLimit: number
   },
   db = database,
+  useClickHouse = false,
 ) {
+  if (useClickHouse) {
+    const stats = await chGetTopProjectsLogsStats({
+      workspaceId: workspace.id,
+      from: range.from,
+      to: range.to,
+      projectsLimit,
+    })
+
+    if (stats.length === 0) return []
+
+    const namesByProjectId = await db
+      .select({ id: projects.id, name: projects.name })
+      .from(projects)
+      .where(
+        and(
+          eq(projects.workspaceId, workspace.id),
+          inArray(
+            projects.id,
+            stats.map((stat) => stat.projectId),
+          ),
+        ),
+      )
+      .then((rows) => new Map(rows.map((row) => [row.id, row.name])))
+
+    return stats.map((stat) => ({
+      projectId: stat.projectId,
+      projectName:
+        namesByProjectId.get(stat.projectId) ?? `Project ${stat.projectId}`,
+      logsCount: stat.logsCount,
+      tokensSpent: stat.tokensSpent,
+      tokensCost: stat.totalCostInMillicents / 100000,
+    }))
+  }
+
   // Pre-filter spans by date range in subqueries for better performance
   const spansInRangeSubquery = db
     .select({
@@ -133,17 +191,30 @@ export async function getLogsData(
   },
   db = database,
 ): Promise<LogStats> {
+  const clickhouseEnabledResult = await isFeatureEnabledByName(
+    workspace.id,
+    CLICKHOUSE_SPANS_READ_FLAG,
+    db,
+  )
+  const shouldUseClickHouse =
+    clickhouseEnabledResult.ok && clickhouseEnabledResult.value
+
   const usedInProduction = await hasProductionTraces(
     { workspaceId: workspace.id },
     db,
   )
 
   const range = getDateRangeOrLastWeekRange(dateRange)
 
-  const globalStats = await getGlobalLogsStats({ workspace, range }, db)
+  const globalStats = await getGlobalLogsStats(
+    { workspace, range },
+    db,
+    shouldUseClickHouse,
+  )
   const topProjects = await getTopProjectsLogsStats(
     { workspace, range, projectsLimit },
     db,
+    shouldUseClickHouse,
   )
 
   return {

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

@@ -0,0 +1,48 @@
+import { LogSources } from '@latitude-data/constants'
+import { clickhouseClient } from '../../../client/clickhouse'
+import { SPANS_TABLE } from '../../../clickhouse/models/spans'
+
+export async function countDistinctTracesByDocument({
+  workspaceId,
+  documentUuid,
+  commitUuids,
+  logSources,
+}: {
+  workspaceId: number
+  documentUuid: string
+  commitUuids: string[]
+  logSources?: LogSources[]
+}) {
+  if (commitUuids.length === 0) return 0
+
+  const queryParams: Record<string, unknown> = {
+    workspaceId,
+    documentUuid,
+    commitUuids,
+  }
+
+  const sourceFilter =
+    logSources && logSources.length > 0
+      ? 'AND source IN ({logSources: Array(String)})'
+      : ''
+
+  if (logSources && logSources.length > 0) {
+    queryParams.logSources = logSources
+  }
+
+  const result = await clickhouseClient().query({
+    query: `
+      SELECT count(DISTINCT trace_id) AS cnt
+      FROM ${SPANS_TABLE}
+      WHERE workspace_id = {workspaceId: UInt64}
+        AND document_uuid = {documentUuid: String}
+        AND commit_uuid IN ({commitUuids: Array(String)})
+        ${sourceFilter}
+    `,
+    format: 'JSONEachRow',
+    query_params: queryParams,
+  })
+
+  const rows = await result.json<{ cnt: string }>()
+  return Number(rows[0]?.cnt ?? 0)
+}

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

@@ -0,0 +1,31 @@
+import { MAIN_SPAN_TYPES } from '@latitude-data/constants'
+import { clickhouseClient } from '../../../client/clickhouse'
+import { SPANS_TABLE } from '../../../clickhouse/models/spans'
+import { toClickHouseDateTime } from '../../../clickhouse/insert'
+
+export async function countMainTypesSince({
+  workspaceId,
+  since,
+}: {
+  workspaceId: number
+  since: Date
+}) {
+  const result = await clickhouseClient().query({
+    query: `
+      SELECT count() AS cnt
+      FROM ${SPANS_TABLE}
+      WHERE workspace_id = {workspaceId: UInt64}
+        AND type IN ({spanTypes: Array(String)})
+        AND started_at >= {since: DateTime64(6, 'UTC')}
+    `,
+    format: 'JSONEachRow',
+    query_params: {
+      workspaceId,
+      spanTypes: Array.from(MAIN_SPAN_TYPES),
+      since: toClickHouseDateTime(since),
+    },
+  })
+
+  const rows = await result.json<{ cnt: string }>()
+  return Number(rows[0]?.cnt ?? 0)
+}

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

@@ -0,0 +1,40 @@
+import { SpanType } from '@latitude-data/constants'
+import { clickhouseClient } from '../../../client/clickhouse'
+import { SPANS_TABLE } from '../../../clickhouse/models/spans'
+import { toClickHouseDateTime } from '../../../clickhouse/insert'
+
+export async function getExperimentPromptSpansBefore({
+  workspaceId,
+  documentUuid,
+  before,
+  limit,
+}: {
+  workspaceId: number
+  documentUuid: string
+  before: Date
+  limit: number
+}) {
+  const result = await clickhouseClient().query({
+    query: `
+      SELECT span_id, trace_id
+      FROM ${SPANS_TABLE}
+      WHERE workspace_id = {workspaceId: UInt64}
+        AND document_uuid = {documentUuid: String}
+        AND started_at < {before: DateTime64(6, 'UTC')}
+        AND experiment_uuid IS NULL
+        AND type = {promptType: String}
+      ORDER BY started_at DESC
+      LIMIT {limit: UInt32}
+    `,
+    format: 'JSONEachRow',
+    query_params: {
+      workspaceId,
+      documentUuid,
+      before: toClickHouseDateTime(before),
+      promptType: SpanType.Prompt,
+      limit,
+    },
+  })
+
+  return result.json<{ span_id: string; trace_id: string }>()
+}

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

@@ -0,0 +1,68 @@
+import { SpanType } from '@latitude-data/constants'
+import { clickhouseClient } from '../../../client/clickhouse'
+import { SPANS_TABLE } from '../../../clickhouse/models/spans'
+
+export type ExperimentRunMetadataResult = {
+  count: number
+  totalDuration: number
+  totalCost: number
+  totalTokens: number
+}
+
+export async function getExperimentRunMetadata({
+  workspaceId,
+  experimentUuid,
+}: {
+  workspaceId: number
+  experimentUuid: string
+}): Promise<ExperimentRunMetadataResult> {
+  const result = await clickhouseClient().query({
+    query: `
+      WITH experiment_trace_ids AS (
+        SELECT DISTINCT trace_id
+        FROM ${SPANS_TABLE}
+        WHERE workspace_id = {workspaceId: UInt64}
+          AND experiment_uuid = {experimentUuid: String}
+      )
+      SELECT
+        countDistinctIf(trace_id, type IN ({runTypes: Array(String)})) AS trace_count,
+        coalesce(sumIf(duration_ms, type IN ({runTypes: Array(String)})), 0) AS total_duration,
+        coalesce(sumIf(cost, type = {completionType: String}), 0) AS total_cost,
+        coalesce(
+          sumIf(
+            coalesce(tokens_prompt, 0) +
+              coalesce(tokens_cached, 0) +
+              coalesce(tokens_reasoning, 0) +
+              coalesce(tokens_completion, 0),
+            type = {completionType: String}
+          ),
+          0
+        ) AS total_tokens
+      FROM ${SPANS_TABLE}
+      WHERE workspace_id = {workspaceId: UInt64}
+        AND trace_id IN (SELECT trace_id FROM experiment_trace_ids)
+    `,
+    format: 'JSONEachRow',
+    query_params: {
+      workspaceId,
+      experimentUuid,
+      runTypes: [SpanType.Prompt, SpanType.Chat],
+      completionType: SpanType.Completion,
+    },
+  })
+
+  const rows = await result.json<{
+    trace_count: string
+    total_duration: string
+    total_cost: string
+    total_tokens: string
+  }>()
+
+  const row = rows[0]
+  return {
+    count: Number(row?.trace_count ?? 0),
+    totalDuration: Number(row?.total_duration ?? 0),
+    totalCost: Number(row?.total_cost ?? 0),
+    totalTokens: Number(row?.total_tokens ?? 0),
+  }
+}