|
| 1 | +import { ClickHouse, type TaskEventV2Input } from "@internal/clickhouse"; |
| 2 | +import { clickhouseTest } from "@internal/testcontainers"; |
| 3 | +import { describe, expect, vi } from "vitest"; |
| 4 | +import { z } from "zod"; |
| 5 | +import { ClickhouseEventRepository } from "~/v3/eventRepository/clickhouseEventRepository.server"; |
| 6 | + |
| 7 | +vi.setConfig({ testTimeout: 60_000 }); |
| 8 | + |
| 9 | +const TRACE_ID = "trace_reader_limits"; |
| 10 | +const ENV_ID = "env_reader_limits"; |
| 11 | +const BASE = new Date("2026-09-01T10:00:00.000Z"); |
| 12 | + |
| 13 | +function clickhouseDate(value: Date) { |
| 14 | + return value.toISOString().replace("T", " ").replace("Z", ""); |
| 15 | +} |
| 16 | + |
| 17 | +const EVENT: TaskEventV2Input = { |
| 18 | + environment_id: ENV_ID, |
| 19 | + organization_id: "org_reader_limits", |
| 20 | + project_id: "project_reader_limits", |
| 21 | + task_identifier: "limits-task", |
| 22 | + run_id: "run_reader_limits", |
| 23 | + start_time: clickhouseDate(BASE), |
| 24 | + duration: "1000000", |
| 25 | + trace_id: TRACE_ID, |
| 26 | + span_id: "root", |
| 27 | + parent_span_id: "", |
| 28 | + message: "root", |
| 29 | + kind: "SPAN", |
| 30 | + status: "OK", |
| 31 | + attributes: {}, |
| 32 | + metadata: "{}", |
| 33 | + expires_at: clickhouseDate(new Date(Date.now() + 86_400_000)), |
| 34 | +}; |
| 35 | + |
| 36 | +async function setup(url: string) { |
| 37 | + const clickhouse = new ClickHouse({ |
| 38 | + writerUrl: url, |
| 39 | + readerUrl: url, |
| 40 | + readerClickhouseSettings: { max_execution_time: 30 }, |
| 41 | + }); |
| 42 | + const repository = new ClickhouseEventRepository({ |
| 43 | + clickhouse, |
| 44 | + version: "v2", |
| 45 | + exportMaxExecutionTime: 300, |
| 46 | + }); |
| 47 | + |
| 48 | + const [insertError] = await clickhouse.taskEventsV2.insert([EVENT]); |
| 49 | + if (insertError) throw insertError; |
| 50 | + |
| 51 | + return { clickhouse, repository }; |
| 52 | +} |
| 53 | + |
| 54 | +async function shutdown(clickhouse: ClickHouse, repository: ClickhouseEventRepository) { |
| 55 | + await Promise.all([ |
| 56 | + (repository as any)._flushScheduler.shutdown(), |
| 57 | + (repository as any)._llmMetricsFlushScheduler.shutdown(), |
| 58 | + (repository as any)._otlpMetricsFlushScheduler.shutdown(), |
| 59 | + ]); |
| 60 | + await clickhouse.close(); |
| 61 | +} |
| 62 | + |
| 63 | +async function serverNow(url: string): Promise<string> { |
| 64 | + const inspector = new ClickHouse({ url, name: "inspector" }); |
| 65 | + const [error, rows] = await inspector.reader.query({ |
| 66 | + name: "read-server-now", |
| 67 | + query: "SELECT toString(now64(6)) AS now", |
| 68 | + schema: z.object({ now: z.string() }), |
| 69 | + })({}); |
| 70 | + await inspector.close(); |
| 71 | + if (error) throw error; |
| 72 | + |
| 73 | + return rows[0].now; |
| 74 | +} |
| 75 | + |
| 76 | +async function eventReadLimits(url: string, since: string): Promise<string[]> { |
| 77 | + const inspector = new ClickHouse({ url, name: "inspector" }); |
| 78 | + const [flushError] = await inspector.writer.command({ |
| 79 | + name: "flush-logs", |
| 80 | + query: "SYSTEM FLUSH LOGS", |
| 81 | + })({}); |
| 82 | + if (flushError) throw flushError; |
| 83 | + |
| 84 | + const [error, rows] = await inspector.reader.query({ |
| 85 | + name: "read-event-read-limits", |
| 86 | + query: `SELECT Settings['max_execution_time'] AS max_execution_time |
| 87 | + FROM system.query_log |
| 88 | + WHERE type = 'QueryFinish' |
| 89 | + AND query_kind = 'Select' |
| 90 | + AND has(tables, 'trigger_dev.task_events_v2') |
| 91 | + AND query_start_time_microseconds > toDateTime64({since: String}, 6)`, |
| 92 | + params: z.object({ since: z.string() }), |
| 93 | + schema: z.object({ max_execution_time: z.string() }), |
| 94 | + })({ since }); |
| 95 | + await inspector.close(); |
| 96 | + if (error) throw error; |
| 97 | + |
| 98 | + return rows.map((row) => row.max_execution_time); |
| 99 | +} |
| 100 | + |
| 101 | +describe("event read limits", () => { |
| 102 | + clickhouseTest("export reads use the export limit", async ({ clickhouseContainer }) => { |
| 103 | + const url = clickhouseContainer.getConnectionUrl(); |
| 104 | + const { clickhouse, repository } = await setup(url); |
| 105 | + const since = await serverNow(url); |
| 106 | + |
| 107 | + const exported: string[] = []; |
| 108 | + for await (const event of repository.streamTraceEvents( |
| 109 | + "taskEventPartitioned", |
| 110 | + ENV_ID, |
| 111 | + TRACE_ID, |
| 112 | + new Date(BASE.getTime() - 60_000), |
| 113 | + new Date(BASE.getTime() + 60_000) |
| 114 | + )) { |
| 115 | + exported.push(event.spanId); |
| 116 | + } |
| 117 | + |
| 118 | + expect(exported).toEqual(["root"]); |
| 119 | + expect(await eventReadLimits(url, since)).toEqual(["300"]); |
| 120 | + |
| 121 | + await shutdown(clickhouse, repository); |
| 122 | + }); |
| 123 | + |
| 124 | + clickhouseTest("interactive reads use the reader limit", async ({ clickhouseContainer }) => { |
| 125 | + const url = clickhouseContainer.getConnectionUrl(); |
| 126 | + const { clickhouse, repository } = await setup(url); |
| 127 | + const since = await serverNow(url); |
| 128 | + |
| 129 | + const chunk = await repository.getTraceChunk( |
| 130 | + "taskEventPartitioned", |
| 131 | + ENV_ID, |
| 132 | + TRACE_ID, |
| 133 | + new Date(BASE.getTime() - 60_000), |
| 134 | + new Date(BASE.getTime() + 60_000), |
| 135 | + undefined, |
| 136 | + { includeDebugLogs: true, limit: 100 } |
| 137 | + ); |
| 138 | + |
| 139 | + expect(chunk?.events.map((event) => event.spanId)).toEqual(["root"]); |
| 140 | + expect(new Set(await eventReadLimits(url, since))).toEqual(new Set(["30"])); |
| 141 | + |
| 142 | + await shutdown(clickhouse, repository); |
| 143 | + }); |
| 144 | +}); |
0 commit comments