Skip to content

Commit f47b422

Browse files
geclosgeclos
committed
refactor(core): scope clickhouse span queries
Wrap span query exports with scopedQuery for consistent workspace-scoped query definitions and shared error unwrapping semantics.
1 parent 874c8ee commit f47b422

22 files changed

+1136
-949
lines changed

‎apps/web/src/app/api/traces/[traceId]/spans/[spanId]/route.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,9 @@ import {
44
SpanMetadatasRepository,
55
SpansRepository,
66
} from '@latitude-data/core/repositories'
7-
import { notFound } from 'next/navigation'
87
import { NextRequest, NextResponse } from 'next/server'
98
import { Workspace } from '@latitude-data/core/schema/models/types/Workspace'
9+
import { NotFoundError } from '@latitude-data/constants/errors'
1010

1111
export const GET = errorHandler(
1212
authHandler(
@@ -30,7 +30,7 @@ export const GET = errorHandler(
3030
.get({ traceId, spanId })
3131
.then((r) => r.unwrap())
3232

33-
if (!span) return notFound()
33+
if (!span) throw new NotFoundError(`Span with id ${spanId} not found`)
3434

3535
const metadatasRepository = new SpanMetadatasRepository(workspace.id)
3636
const metadata = await metadatasRepository

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

Lines changed: 65 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -2,35 +2,40 @@ import { SpanType } from '@latitude-data/constants'
22
import { clickhouseClient } from '../../../client/clickhouse'
33
import { SPANS_TABLE } from '../../../clickhouse/models/spans'
44
import { TracesAggregations } from '../../../schema/models/types/Span'
5+
import { scopedQuery } from '../../scope'
56

6-
export async function computeDocumentTracesAggregations({
7-
workspaceId,
8-
documentUuid,
9-
commitUuid,
10-
}: {
11-
workspaceId: number
12-
documentUuid: string
13-
commitUuid?: string
14-
}): Promise<TracesAggregations> {
15-
const params: Record<string, unknown> = {
16-
workspaceId,
17-
documentUuid,
18-
completionType: SpanType.Completion,
19-
}
7+
export const computeDocumentTracesAggregations = scopedQuery(
8+
async function computeDocumentTracesAggregations(
9+
{
10+
workspaceId,
11+
documentUuid,
12+
commitUuid,
13+
}: {
14+
workspaceId: number
15+
documentUuid: string
16+
commitUuid?: string
17+
},
18+
_db,
19+
): Promise<TracesAggregations> {
20+
const params: Record<string, unknown> = {
21+
workspaceId,
22+
documentUuid,
23+
completionType: SpanType.Completion,
24+
}
2025

21-
const conditions = [
22-
`workspace_id = {workspaceId: UInt64}`,
23-
`document_uuid = {documentUuid: UUID}`,
24-
]
26+
const conditions = [
27+
`workspace_id = {workspaceId: UInt64}`,
28+
`document_uuid = {documentUuid: UUID}`,
29+
]
2530

26-
if (commitUuid) {
27-
params.commitUuid = commitUuid
28-
conditions.push(`commit_uuid = {commitUuid: UUID}`)
29-
}
31+
if (commitUuid) {
32+
params.commitUuid = commitUuid
33+
conditions.push(`commit_uuid = {commitUuid: UUID}`)
34+
}
3035

31-
// Query all metrics in a single pass
32-
const aggregationResult = await clickhouseClient().query({
33-
query: `
36+
// Query all metrics in a single pass
37+
const aggregationResult = await clickhouseClient().query({
38+
query: `
3439
SELECT
3540
countDistinct(trace_id) AS total_count,
3641
coalesce(
@@ -89,40 +94,41 @@ export async function computeDocumentTracesAggregations({
8994
FROM ${SPANS_TABLE}
9095
WHERE ${conditions.join(' AND ')}
9196
`,
92-
format: 'JSONEachRow',
93-
query_params: params,
94-
})
97+
format: 'JSONEachRow',
98+
query_params: params,
99+
})
95100

96-
const aggRows = await aggregationResult.json<{
97-
total_count: string
98-
total_tokens: string
99-
total_cost_in_millicents: string
100-
average_tokens: string
101-
average_cost_in_millicents: string
102-
median_cost_in_millicents: string
103-
average_duration: string
104-
median_duration: string
105-
}>()
101+
const aggRows = await aggregationResult.json<{
102+
total_count: string
103+
total_tokens: string
104+
total_cost_in_millicents: string
105+
average_tokens: string
106+
average_cost_in_millicents: string
107+
median_cost_in_millicents: string
108+
average_duration: string
109+
median_duration: string
110+
}>()
106111

107-
const row = aggRows[0] ?? {
108-
total_count: '0',
109-
total_tokens: '0',
110-
total_cost_in_millicents: '0',
111-
average_tokens: '0',
112-
average_cost_in_millicents: '0',
113-
median_cost_in_millicents: '0',
114-
average_duration: '0',
115-
median_duration: '0',
116-
}
112+
const row = aggRows[0] ?? {
113+
total_count: '0',
114+
total_tokens: '0',
115+
total_cost_in_millicents: '0',
116+
average_tokens: '0',
117+
average_cost_in_millicents: '0',
118+
median_cost_in_millicents: '0',
119+
average_duration: '0',
120+
median_duration: '0',
121+
}
117122

118-
return {
119-
totalCount: Number(row.total_count),
120-
totalTokens: Number(row.total_tokens),
121-
totalCostInMillicents: Number(row.total_cost_in_millicents),
122-
averageTokens: Number(row.average_tokens),
123-
averageCostInMillicents: Number(row.average_cost_in_millicents),
124-
medianCostInMillicents: Number(row.median_cost_in_millicents),
125-
averageDuration: Number(row.average_duration),
126-
medianDuration: Number(row.median_duration),
127-
}
128-
}
123+
return {
124+
totalCount: Number(row.total_count),
125+
totalTokens: Number(row.total_tokens),
126+
totalCostInMillicents: Number(row.total_cost_in_millicents),
127+
averageTokens: Number(row.average_tokens),
128+
averageCostInMillicents: Number(row.average_cost_in_millicents),
129+
medianCostInMillicents: Number(row.median_cost_in_millicents),
130+
averageDuration: Number(row.average_duration),
131+
medianDuration: Number(row.median_duration),
132+
}
133+
},
134+
)

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

Lines changed: 48 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -2,45 +2,50 @@ import { subDays } from 'date-fns'
22
import { clickhouseClient } from '../../../client/clickhouse'
33
import { SPANS_TABLE } from '../../../clickhouse/models/spans'
44
import { toClickHouseDateTime } from '../../../clickhouse/insert'
5+
import { scopedQuery } from '../../scope'
56

67
export type DailyCount = {
78
date: string
89
count: number
910
}
1011

11-
export async function computeDocumentTracesDailyCount({
12-
workspaceId,
13-
documentUuid,
14-
commitUuid,
15-
days = 30,
16-
}: {
17-
workspaceId: number
18-
documentUuid: string
19-
commitUuid?: string
20-
days?: number
21-
}): Promise<DailyCount[]> {
22-
const now = new Date()
23-
const startDate = subDays(now, days)
24-
25-
const params: Record<string, unknown> = {
26-
workspaceId,
27-
documentUuid,
28-
startDate: toClickHouseDateTime(startDate),
29-
}
30-
31-
const conditions = [
32-
`workspace_id = {workspaceId: UInt64}`,
33-
`document_uuid = {documentUuid: UUID}`,
34-
`started_at >= {startDate: DateTime64(6, 'UTC')}`,
35-
]
36-
37-
if (commitUuid) {
38-
params.commitUuid = commitUuid
39-
conditions.push(`commit_uuid = {commitUuid: UUID}`)
40-
}
41-
42-
const result = await clickhouseClient().query({
43-
query: `
12+
export const computeDocumentTracesDailyCount = scopedQuery(
13+
async function computeDocumentTracesDailyCount(
14+
{
15+
workspaceId,
16+
documentUuid,
17+
commitUuid,
18+
days = 30,
19+
}: {
20+
workspaceId: number
21+
documentUuid: string
22+
commitUuid?: string
23+
days?: number
24+
},
25+
_db,
26+
): Promise<DailyCount[]> {
27+
const now = new Date()
28+
const startDate = subDays(now, days)
29+
30+
const params: Record<string, unknown> = {
31+
workspaceId,
32+
documentUuid,
33+
startDate: toClickHouseDateTime(startDate),
34+
}
35+
36+
const conditions = [
37+
`workspace_id = {workspaceId: UInt64}`,
38+
`document_uuid = {documentUuid: UUID}`,
39+
`started_at >= {startDate: DateTime64(6, 'UTC')}`,
40+
]
41+
42+
if (commitUuid) {
43+
params.commitUuid = commitUuid
44+
conditions.push(`commit_uuid = {commitUuid: UUID}`)
45+
}
46+
47+
const result = await clickhouseClient().query({
48+
query: `
4449
SELECT
4550
toDate(started_at) AS date,
4651
count(DISTINCT trace_id) AS count
@@ -49,14 +54,15 @@ export async function computeDocumentTracesDailyCount({
4954
GROUP BY date
5055
ORDER BY date
5156
`,
52-
format: 'JSONEachRow',
53-
query_params: params,
54-
})
57+
format: 'JSONEachRow',
58+
query_params: params,
59+
})
5560

56-
const rows = await result.json<{ date: string; count: string }>()
61+
const rows = await result.json<{ date: string; count: string }>()
5762

58-
return rows.map((row) => ({
59-
date: row.date,
60-
count: Number(row.count),
61-
}))
62-
}
63+
return rows.map((row) => ({
64+
date: row.date,
65+
count: Number(row.count),
66+
}))
67+
},
68+
)

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

Lines changed: 38 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,48 +1,54 @@
11
import { LogSources } from '@latitude-data/constants'
22
import { clickhouseClient } from '../../../client/clickhouse'
33
import { SPANS_TABLE } from '../../../clickhouse/models/spans'
4+
import { scopedQuery } from '../../scope'
45

5-
export async function countDistinctTracesByDocument({
6-
workspaceId,
7-
documentUuid,
8-
commitUuids,
9-
logSources,
10-
}: {
11-
workspaceId: number
12-
documentUuid: string
13-
commitUuids: string[]
14-
logSources?: LogSources[]
15-
}) {
16-
if (commitUuids.length === 0) return 0
6+
export const countDistinctTracesByDocument = scopedQuery(
7+
async function countDistinctTracesByDocument(
8+
{
9+
workspaceId,
10+
documentUuid,
11+
commitUuids,
12+
logSources,
13+
}: {
14+
workspaceId: number
15+
documentUuid: string
16+
commitUuids: string[]
17+
logSources?: LogSources[]
18+
},
19+
_db,
20+
) {
21+
if (commitUuids.length === 0) return 0
1722

18-
const queryParams: Record<string, unknown> = {
19-
workspaceId,
20-
documentUuid,
21-
commitUuids,
22-
}
23+
const queryParams: Record<string, unknown> = {
24+
workspaceId,
25+
documentUuid,
26+
commitUuids,
27+
}
2328

24-
const sourceFilter =
25-
logSources && logSources.length > 0
26-
? 'AND source IN ({logSources: Array(String)})'
27-
: ''
29+
const sourceFilter =
30+
logSources && logSources.length > 0
31+
? 'AND source IN ({logSources: Array(String)})'
32+
: ''
2833

29-
if (logSources && logSources.length > 0) {
30-
queryParams.logSources = logSources
31-
}
34+
if (logSources && logSources.length > 0) {
35+
queryParams.logSources = logSources
36+
}
3237

33-
const result = await clickhouseClient().query({
34-
query: `
38+
const result = await clickhouseClient().query({
39+
query: `
3540
SELECT count(DISTINCT trace_id) AS cnt
3641
FROM ${SPANS_TABLE}
3742
WHERE workspace_id = {workspaceId: UInt64}
3843
AND document_uuid = {documentUuid: UUID}
3944
AND commit_uuid IN ({commitUuids: Array(UUID)})
4045
${sourceFilter}
4146
`,
42-
format: 'JSONEachRow',
43-
query_params: queryParams,
44-
})
47+
format: 'JSONEachRow',
48+
query_params: queryParams,
49+
})
4550

46-
const rows = await result.json<{ cnt: string }>()
47-
return Number(rows[0]?.cnt ?? 0)
48-
}
51+
const rows = await result.json<{ cnt: string }>()
52+
return Number(rows[0]?.cnt ?? 0)
53+
},
54+
)

0 commit comments

Comments
 (0)