|
| 1 | +import { |
| 2 | + buildTraceStartTimeSortKeyPredicate, |
| 3 | + ClickHouse, |
| 4 | + type TaskEventV2Input, |
| 5 | +} from "@internal/clickhouse"; |
| 6 | +import { clickhouseTest } from "@internal/testcontainers"; |
| 7 | +import { describe, expect, vi } from "vitest"; |
| 8 | +import { z } from "zod"; |
| 9 | +import { |
| 10 | + ClickhouseEventRepository, |
| 11 | + convertDateToClickhouseDateTime, |
| 12 | +} from "~/v3/eventRepository/clickhouseEventRepository.server"; |
| 13 | +import { convertDateToNanoseconds } from "~/v3/eventRepository/common.server"; |
| 14 | + |
| 15 | +vi.setConfig({ testTimeout: 120_000 }); |
| 16 | + |
| 17 | +const TRACE_ID = "trace_start_time_end_test"; |
| 18 | +const ENV_ID = "env_start_time_end_test"; |
| 19 | +const START = new Date("2026-09-01T10:00:00.000Z"); |
| 20 | +// Sub-second end, so a rounding mistake in the sort-key bound would drop rows. |
| 21 | +const END = new Date("2026-09-01T10:00:00.500Z"); |
| 22 | +// The trace scope adds 60s to the end. |
| 23 | +const SCOPE_END_SECOND = "2026-09-01 10:01:00"; |
| 24 | +const NOISE_ROWS = 30_000; |
| 25 | +// query_log stores the SQL with params inlined as _CAST(value, 'Int64'). |
| 26 | +const SORT_KEY_START = |
| 27 | + /toUnixTimestamp\(start_time\) >= intDiv\(_CAST\(\d+, 'Int64'\), 1000000000\)/; |
| 28 | +const SORT_KEY_END = /toUnixTimestamp\(start_time\) <= intDiv\(_CAST\(\d+, 'Int64'\), 1000000000\)/; |
| 29 | + |
| 30 | +function clickhouseDate(value: Date) { |
| 31 | + return value.toISOString().replace("T", " ").replace("Z", ""); |
| 32 | +} |
| 33 | + |
| 34 | +function row( |
| 35 | + traceId: string, |
| 36 | + spanId: string, |
| 37 | + message: string, |
| 38 | + startTime: string, |
| 39 | + insertedAt = START |
| 40 | +): TaskEventV2Input { |
| 41 | + return { |
| 42 | + environment_id: ENV_ID, |
| 43 | + organization_id: "org_start_time_end_test", |
| 44 | + project_id: "project_start_time_end_test", |
| 45 | + task_identifier: "start-time-end-task", |
| 46 | + run_id: "run_start_time_end_test", |
| 47 | + start_time: startTime, |
| 48 | + inserted_at: convertDateToClickhouseDateTime(insertedAt), |
| 49 | + duration: "1000000", |
| 50 | + trace_id: traceId, |
| 51 | + span_id: spanId, |
| 52 | + parent_span_id: "", |
| 53 | + message, |
| 54 | + kind: "SPAN", |
| 55 | + status: "OK", |
| 56 | + attributes: {}, |
| 57 | + metadata: "{}", |
| 58 | + expires_at: clickhouseDate(new Date(Date.now() + 86_400_000)), |
| 59 | + }; |
| 60 | +} |
| 61 | + |
| 62 | +async function shutdownRepository(repository: ClickhouseEventRepository): Promise<void> { |
| 63 | + await Promise.all([ |
| 64 | + (repository as any)._flushScheduler.shutdown(), |
| 65 | + (repository as any)._llmMetricsFlushScheduler.shutdown(), |
| 66 | + (repository as any)._otlpMetricsFlushScheduler.shutdown(), |
| 67 | + ]); |
| 68 | +} |
| 69 | + |
| 70 | +function setup(url: string) { |
| 71 | + const clickhouse = new ClickHouse({ url, name: "test" }); |
| 72 | + const repository = new ClickhouseEventRepository({ |
| 73 | + clickhouse, |
| 74 | + version: "v2", |
| 75 | + insertStrategy: "insert", |
| 76 | + }); |
| 77 | + return { clickhouse, repository }; |
| 78 | +} |
| 79 | + |
| 80 | +async function readMessages(repository: ClickhouseEventRepository, end: Date | undefined) { |
| 81 | + const chunk = await repository.getTraceChunk( |
| 82 | + "taskEventPartitioned", |
| 83 | + ENV_ID, |
| 84 | + TRACE_ID, |
| 85 | + START, |
| 86 | + end, |
| 87 | + undefined, |
| 88 | + {} |
| 89 | + ); |
| 90 | + return (chunk?.events ?? []).map((e) => e.message).sort(); |
| 91 | +} |
| 92 | + |
| 93 | +describe("trace reads with a start_time end", () => { |
| 94 | + clickhouseTest("returns the same rows around the end second", async ({ clickhouseContainer }) => { |
| 95 | + const { clickhouse, repository } = setup(clickhouseContainer.getConnectionUrl()); |
| 96 | + |
| 97 | + try { |
| 98 | + const [insertError] = await clickhouse.taskEventsV2.insert([ |
| 99 | + row(TRACE_ID, "a", "second-start", `${SCOPE_END_SECOND}.000000000`), |
| 100 | + row(TRACE_ID, "b", "before-end", `${SCOPE_END_SECOND}.100000000`), |
| 101 | + row(TRACE_ID, "c", "at-end", `${SCOPE_END_SECOND}.500000000`), |
| 102 | + row(TRACE_ID, "d", "end-plus-1ns", `${SCOPE_END_SECOND}.500000001`), |
| 103 | + row(TRACE_ID, "e", "same-second-later", `${SCOPE_END_SECOND}.900000000`), |
| 104 | + row(TRACE_ID, "f", "next-second", "2026-09-01 10:01:01.000000000"), |
| 105 | + ]); |
| 106 | + expect(insertError).toBeNull(); |
| 107 | + |
| 108 | + expect(await readMessages(repository, END)).toEqual(["at-end", "before-end", "second-start"]); |
| 109 | + expect(await readMessages(repository, undefined)).toEqual([ |
| 110 | + "at-end", |
| 111 | + "before-end", |
| 112 | + "end-plus-1ns", |
| 113 | + "next-second", |
| 114 | + "same-second-later", |
| 115 | + "second-start", |
| 116 | + ]); |
| 117 | + } finally { |
| 118 | + await shutdownRepository(repository); |
| 119 | + } |
| 120 | + }); |
| 121 | + |
| 122 | + clickhouseTest( |
| 123 | + "narrows the primary key with the sort-key bounds alone", |
| 124 | + async ({ clickhouseContainer }) => { |
| 125 | + const { clickhouse, repository } = setup(clickhouseContainer.getConnectionUrl()); |
| 126 | + |
| 127 | + try { |
| 128 | + // Other traces in the same env, one per second before and after the run, so the |
| 129 | + // run's row sits mid-part with granules on both sides. |
| 130 | + const noise = Array.from({ length: NOISE_ROWS }, (_, i) => |
| 131 | + row( |
| 132 | + `trace_noise_${i}`, |
| 133 | + `noise_${i}`, |
| 134 | + "noise", |
| 135 | + clickhouseDate(new Date(START.getTime() + (i - NOISE_ROWS / 2) * 1000)) |
| 136 | + ) |
| 137 | + ); |
| 138 | + const [insertError] = await clickhouse.taskEventsV2.insert([ |
| 139 | + row(TRACE_ID, "a", "span", clickhouseDate(START)), |
| 140 | + ...noise, |
| 141 | + ]); |
| 142 | + expect(insertError).toBeNull(); |
| 143 | + |
| 144 | + // Some ClickHouse setups (prod) ignore the exact string bounds when choosing |
| 145 | + // granules, so check the sort-key clause without them. |
| 146 | + const { clause, params } = buildTraceStartTimeSortKeyPredicate({ |
| 147 | + startNs: convertDateToNanoseconds(START).toString(), |
| 148 | + endNs: convertDateToNanoseconds(END).toString(), |
| 149 | + }); |
| 150 | + const explain = clickhouse.reader.query({ |
| 151 | + name: "explainSortKeyBounds", |
| 152 | + query: `EXPLAIN indexes = 1 SELECT span_id FROM trigger_dev.task_events_v2 WHERE environment_id = {environmentId: String} AND ${clause}`, |
| 153 | + params: z.any(), |
| 154 | + schema: z.object({ explain: z.string() }), |
| 155 | + }); |
| 156 | + const [explainError, rows] = await explain({ ...params, environmentId: ENV_ID }); |
| 157 | + expect(explainError).toBeNull(); |
| 158 | + |
| 159 | + const lines = (rows ?? []).map((r) => r.explain.trim()); |
| 160 | + const primaryKey = lines.indexOf("PrimaryKey"); |
| 161 | + const granules = lines.slice(primaryKey).find((l) => l.startsWith("Granules:")); |
| 162 | + const [selected, total] = granules!.replace("Granules:", "").trim().split("/").map(Number); |
| 163 | + expect(total).toBeGreaterThan(2); |
| 164 | + expect(selected).toBe(1); |
| 165 | + } finally { |
| 166 | + await shutdownRepository(repository); |
| 167 | + } |
| 168 | + } |
| 169 | + ); |
| 170 | + |
| 171 | + clickhouseTest( |
| 172 | + "adds the sort-key bounds to trace chunk reads", |
| 173 | + async ({ clickhouseContainer }) => { |
| 174 | + const { clickhouse, repository } = setup(clickhouseContainer.getConnectionUrl()); |
| 175 | + |
| 176 | + try { |
| 177 | + const [insertError] = await clickhouse.taskEventsV2.insert([ |
| 178 | + row(TRACE_ID, "a", "span", clickhouseDate(START)), |
| 179 | + ]); |
| 180 | + expect(insertError).toBeNull(); |
| 181 | + |
| 182 | + // The container is shared across tests, so only look at queries from here on. |
| 183 | + const serverNow = clickhouse.reader.query({ |
| 184 | + name: "serverNowMicros", |
| 185 | + query: "SELECT toString(toUnixTimestamp64Micro(now64(6))) AS since", |
| 186 | + schema: z.object({ since: z.string() }), |
| 187 | + }); |
| 188 | + const [nowError, nowRows] = await serverNow({}); |
| 189 | + expect(nowError).toBeNull(); |
| 190 | + const since = nowRows![0].since; |
| 191 | + |
| 192 | + // Completed run first, then live run, so query_log order identifies them. |
| 193 | + expect(await readMessages(repository, END)).toEqual(["span"]); |
| 194 | + expect(await readMessages(repository, undefined)).toEqual(["span"]); |
| 195 | + |
| 196 | + const flushLogs = clickhouse.writer.command({ |
| 197 | + name: "flushLogs", |
| 198 | + query: "SYSTEM FLUSH LOGS", |
| 199 | + }); |
| 200 | + const [flushError] = await flushLogs({}); |
| 201 | + expect(flushError).toBeNull(); |
| 202 | + |
| 203 | + const loggedQueries = clickhouse.reader.query({ |
| 204 | + name: "traceChunkQueryLog", |
| 205 | + query: `SELECT query FROM system.query_log |
| 206 | + WHERE type = 'QueryFinish' |
| 207 | + AND toUnixTimestamp64Micro(event_time_microseconds) > {since: Int64} |
| 208 | + AND query LIKE 'SELECT%FROM trigger_dev.task_events_v2%' |
| 209 | + ORDER BY event_time_microseconds`, |
| 210 | + params: z.object({ since: z.string() }), |
| 211 | + schema: z.object({ query: z.string() }), |
| 212 | + }); |
| 213 | + const [logError, rows] = await loggedQueries({ since }); |
| 214 | + expect(logError).toBeNull(); |
| 215 | + |
| 216 | + const queries = (rows ?? []).map((r) => r.query); |
| 217 | + expect(queries).toHaveLength(2); |
| 218 | + const [completed, live] = queries; |
| 219 | + |
| 220 | + expect(completed).toMatch(SORT_KEY_START); |
| 221 | + expect(completed).toMatch(SORT_KEY_END); |
| 222 | + expect(live).toMatch(SORT_KEY_START); |
| 223 | + expect(live).not.toMatch(SORT_KEY_END); |
| 224 | + } finally { |
| 225 | + await shutdownRepository(repository); |
| 226 | + } |
| 227 | + } |
| 228 | + ); |
| 229 | +}); |
0 commit comments