Skip to content

Commit af0daca

Browse files
devinat1devinat1
committed
draft: HQL tracing via datadog
1 parent 5c2b176 commit af0daca

File tree

3 files changed

+167
-2
lines changed

3 files changed

+167
-2
lines changed

‎valhalla/jawn/src/controllers/public/heliconeSqlController.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import {
2424
HQL_FEATURE_FLAG,
2525
} from "../../lib/utils/featureFlags";
2626
import { HqlQueryManager } from "../../managers/HqlQueryManager";
27+
import { Trace } from "../../utils/traceDecorator";
2728

2829
// --- Response Types ---
2930
export interface ClickHouseTableSchema {
@@ -89,6 +90,7 @@ export class HeliconeSqlController extends Controller {
8990
* Get ClickHouse schema (tables and columns)
9091
*/
9192
@Get("schema")
93+
@Trace("hql.controller.getClickHouseSchema")
9294
public async getClickHouseSchema(
9395
@Request() request: JawnAuthenticatedRequest
9496
): Promise<Result<ClickHouseTableSchema[], string>> {
@@ -104,6 +106,7 @@ export class HeliconeSqlController extends Controller {
104106
}
105107

106108
@Post("execute")
109+
@Trace("hql.controller.executeSql")
107110
public async executeSql(
108111
@Body() requestBody: ExecuteSqlRequest,
109112
@Request() request: JawnAuthenticatedRequest
@@ -139,6 +142,7 @@ export class HeliconeSqlController extends Controller {
139142
}
140143

141144
@Post("download")
145+
@Trace("hql.controller.downloadCsv")
142146
public async downloadCsv(
143147
@Body() requestBody: ExecuteSqlRequest,
144148
@Request() request: JawnAuthenticatedRequest
@@ -174,6 +178,7 @@ export class HeliconeSqlController extends Controller {
174178
}
175179

176180
@Get("saved-queries")
181+
@Trace("hql.controller.getSavedQueries")
177182
public async getSavedQueries(
178183
@Request() request: JawnAuthenticatedRequest
179184
): Promise<Result<Array<HqlSavedQuery>, string>> {
@@ -190,6 +195,7 @@ export class HeliconeSqlController extends Controller {
190195
}
191196

192197
@Get("saved-query/{queryId}")
198+
@Trace("hql.controller.getSavedQuery")
193199
public async getSavedQuery(
194200
@Path() queryId: string,
195201
@Request() request: JawnAuthenticatedRequest
@@ -207,6 +213,7 @@ export class HeliconeSqlController extends Controller {
207213
}
208214

209215
@Delete("saved-query/{queryId}")
216+
@Trace("hql.controller.deleteSavedQuery")
210217
public async deleteSavedQuery(
211218
@Path() queryId: string,
212219
@Request() request: JawnAuthenticatedRequest
@@ -224,6 +231,7 @@ export class HeliconeSqlController extends Controller {
224231
}
225232

226233
@Post("saved-queries/bulk-delete")
234+
@Trace("hql.controller.bulkDeleteSavedQueries")
227235
public async bulkDeleteSavedQueries(
228236
@Body() requestBody: BulkDeleteSavedQueriesRequest,
229237
@Request() request: JawnAuthenticatedRequest
@@ -239,6 +247,7 @@ export class HeliconeSqlController extends Controller {
239247
}
240248

241249
@Post("saved-query")
250+
@Trace("hql.controller.createSavedQuery")
242251
public async createSavedQuery(
243252
@Body() requestBody: CreateSavedQueryRequest,
244253
@Request() request: JawnAuthenticatedRequest
@@ -256,6 +265,7 @@ export class HeliconeSqlController extends Controller {
256265
}
257266

258267
@Put("saved-query")
268+
@Trace("hql.controller.updateSavedQuery")
259269
public async updateSavedQuery(
260270
@Body() requestBody: UpdateSavedQueryRequest,
261271
@Request() request: JawnAuthenticatedRequest

‎valhalla/jawn/src/managers/HeliconeSqlManager.ts‎

Lines changed: 51 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import { HqlError, HqlErrorCode, hqlError, parseClickhouseError } from "../lib/e
99
import { AST, Parser } from "node-sql-parser";
1010
import { HqlStore } from "../lib/stores/HqlStore";
1111
import { z } from "zod";
12+
import tracer from "../tracer";
1213

1314
export const CLICKHOUSE_TABLES = ["request_response_rmt"];
1415
const MAX_LIMIT = 300000;
@@ -119,7 +120,10 @@ export class HeliconeSqlManager {
119120
async getClickhouseSchema(): Promise<
120121
Result<ClickHouseTableSchema[], HqlError>
121122
> {
123+
const span = tracer.startSpan("hql.getClickhouseSchema");
122124
try {
125+
span.setTag("resource.name", "hql.getClickhouseSchema");
126+
span.setTag("span.type", "custom");
123127
const schemaPromises = CLICKHOUSE_TABLES.map(async (table_name) => {
124128
const columns = await clickhouseDb.dbQuery<ClickHouseTableRow>(
125129
`DESCRIBE TABLE ${table_name}`,
@@ -149,13 +153,18 @@ export class HeliconeSqlManager {
149153
});
150154

151155
const schema = await Promise.all(schemaPromises);
156+
span.setTag("hql.table_count", schema.length);
152157
return ok(schema);
153158
} catch (e) {
154159
const errorMessage = e instanceof Error ? e.message : String(e);
160+
span.setTag("error", true);
161+
span.setTag("error.message", errorMessage);
155162
return hqlError(
156163
HqlErrorCode.SCHEMA_FETCH_FAILED,
157164
errorMessage
158165
);
166+
} finally {
167+
span.finish();
159168
}
160169
}
161170

@@ -167,12 +176,21 @@ export class HeliconeSqlManager {
167176
sql: string,
168177
limit: number = 100
169178
): Promise<Result<ExecuteSqlResponse, HqlError>> {
179+
const span = tracer.startSpan("hql.executeSql");
180+
span.setTag("resource.name", "hql.executeSql");
181+
span.setTag("span.type", "custom");
182+
span.setTag("hql.limit", limit);
170183
try {
171184
// Parse SQL to validate and add limit
172185
let ast;
173186
try {
174187
ast = parser.astify(sql, { database: "Postgresql" });
175188
} catch (parseError) {
189+
span.setTag("error", true);
190+
span.setTag(
191+
"error.message",
192+
parseError instanceof Error ? parseError.message : String(parseError)
193+
);
176194
return hqlError(
177195
HqlErrorCode.SYNTAX_ERROR,
178196
parseError instanceof Error ? parseError.message : String(parseError)
@@ -187,6 +205,11 @@ export class HeliconeSqlManager {
187205
try {
188206
limitedAst = addLimit(normalizedAst, limit);
189207
} catch (limitError) {
208+
span.setTag("error", true);
209+
span.setTag(
210+
"error.message",
211+
limitError instanceof Error ? limitError.message : String(limitError)
212+
);
190213
return hqlError(
191214
HqlErrorCode.SYNTAX_ERROR,
192215
`Failed to apply limit: ${limitError instanceof Error ? limitError.message : String(limitError)}`
@@ -198,6 +221,8 @@ export class HeliconeSqlManager {
198221
// Validate SQL for security
199222
const validatedSql = validateSql(firstSql);
200223
if (isError(validatedSql)) {
224+
span.setTag("error", true);
225+
span.setTag("error.message", validatedSql.error.message);
201226
return validatedSql;
202227
}
203228

@@ -211,34 +236,53 @@ export class HeliconeSqlManager {
211236
});
212237

213238
const elapsedMilliseconds = Date.now() - start;
239+
span.setTag("hql.elapsed_ms", elapsedMilliseconds);
240+
span.setTag("hql.organization_id", this.authParams.organizationId);
214241

215242
if (isError(result)) {
216243
const errorCode = parseClickhouseError(result.error);
244+
span.setTag("error", true);
245+
span.setTag("error.message", result.error);
246+
span.setTag("hql.error_code", errorCode);
217247
return hqlError(errorCode, result.error);
218248
}
219249

250+
const rows = result.data ?? [];
251+
const size = Buffer.byteLength(JSON.stringify(rows), "utf8");
252+
span.setTag("hql.row_count", rows.length);
253+
span.setTag("hql.size_bytes", size);
220254
return ok({
221255
rows: result.data ?? [],
222256
elapsedMilliseconds,
223-
size: Buffer.byteLength(JSON.stringify(result.data), "utf8"),
224-
rowCount: result.data?.length ?? 0,
257+
size,
258+
rowCount: rows.length,
225259
});
226260
} catch (e) {
227261
const errorMessage = e instanceof Error ? e.message : String(e);
262+
span.setTag("error", true);
263+
span.setTag("error.message", errorMessage);
228264
return hqlError(
229265
HqlErrorCode.UNEXPECTED_ERROR,
230266
errorMessage
231267
);
268+
} finally {
269+
span.finish();
232270
}
233271
}
234272

235273
async downloadCsv(sql: string): Promise<Result<string, HqlError>> {
274+
const span = tracer.startSpan("hql.downloadCsv");
275+
span.setTag("resource.name", "hql.downloadCsv");
276+
span.setTag("span.type", "custom");
236277
const result = await this.executeSql(sql, MAX_LIMIT);
237278
if (isError(result)) {
279+
span.setTag("error", true);
280+
span.setTag("error.message", result.error.message);
238281
return result;
239282
}
240283

241284
if (!result.data?.rows?.length) {
285+
span.setTag("hql.no_data", true);
242286
return hqlError(HqlErrorCode.NO_DATA_RETURNED);
243287
}
244288

@@ -254,16 +298,21 @@ export class HeliconeSqlManager {
254298
);
255299

256300
if (isError(uploadResult)) {
301+
span.setTag("error", true);
302+
span.setTag("error.message", uploadResult.error);
257303
return hqlError(
258304
HqlErrorCode.CSV_UPLOAD_FAILED,
259305
uploadResult.error
260306
);
261307
}
262308

263309
if (!uploadResult.data) {
310+
span.setTag("error", true);
311+
span.setTag("error.message", "CSV URL not returned");
264312
return hqlError(HqlErrorCode.CSV_URL_NOT_RETURNED);
265313
}
266314

315+
span.setTag("hql.csv_url_generated", true);
267316
return ok(uploadResult.data);
268317
}
269318
}

‎valhalla/jawn/src/utils/traceDecorator.ts‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
import tracer from "../tracer";
2+
3+
type TraceConfig = {
4+
name?: string;
5+
resource?: string;
6+
type?: string;
7+
tags?: Record<string, unknown>;
8+
onStart?: (ctx: { args: unknown[]; that: any }) => Record<string, unknown> | void;
9+
onSuccess?: (ctx: { args: unknown[]; that: any; result: unknown }) => Record<string, unknown> | void;
10+
onError?: (ctx: { args: unknown[]; that: any; error: unknown }) => Record<string, unknown> | void;
11+
};
12+
13+
export function Trace(configOrName?: string | TraceConfig) {
14+
return function (
15+
_target: any,
16+
propertyKey: string,
17+
descriptor: PropertyDescriptor
18+
) {
19+
const original = descriptor.value;
20+
21+
descriptor.value = function (...args: unknown[]) {
22+
const that = this as any;
23+
const className = that?.constructor?.name ?? "UnknownClass";
24+
const cfg: TraceConfig =
25+
typeof configOrName === "string"
26+
? { name: configOrName }
27+
: configOrName ?? {};
28+
29+
const spanName = cfg.name ?? `${className}.${String(propertyKey)}`;
30+
const span = tracer.startSpan(spanName);
31+
if (cfg.resource) span.setTag("resource.name", cfg.resource);
32+
if (cfg.type) span.setTag("span.type", cfg.type);
33+
if (cfg.tags) {
34+
for (const [k, v] of Object.entries(cfg.tags)) {
35+
span.setTag(k, v as any);
36+
}
37+
}
38+
39+
try {
40+
const startTags = cfg.onStart?.({ args, that });
41+
if (startTags) {
42+
for (const [k, v] of Object.entries(startTags)) {
43+
span.setTag(k, v as any);
44+
}
45+
}
46+
47+
const maybePromise = original.apply(this, args);
48+
if (maybePromise && typeof maybePromise.then === "function") {
49+
return (maybePromise as Promise<unknown>)
50+
.then((result) => {
51+
const successTags = cfg.onSuccess?.({ args, that, result });
52+
if (successTags) {
53+
for (const [k, v] of Object.entries(successTags)) {
54+
span.setTag(k, v as any);
55+
}
56+
}
57+
return result;
58+
})
59+
.catch((error) => {
60+
span.setTag("error", true);
61+
span.setTag("error.message", error instanceof Error ? error.message : String(error));
62+
const errorTags = cfg.onError?.({ args, that, error });
63+
if (errorTags) {
64+
for (const [k, v] of Object.entries(errorTags)) {
65+
span.setTag(k, v as any);
66+
}
67+
}
68+
throw error;
69+
})
70+
.finally(() => {
71+
span.finish();
72+
});
73+
} else {
74+
const result = maybePromise;
75+
const successTags = cfg.onSuccess?.({ args, that, result });
76+
if (successTags) {
77+
for (const [k, v] of Object.entries(successTags)) {
78+
span.setTag(k, v as any);
79+
}
80+
}
81+
return result;
82+
}
83+
} catch (error) {
84+
span.setTag("error", true);
85+
span.setTag("error.message", error instanceof Error ? error.message : String(error));
86+
const errorTags = cfg.onError?.({ args, that, error });
87+
if (errorTags) {
88+
for (const [k, v] of Object.entries(errorTags)) {
89+
span.setTag(k, v as any);
90+
}
91+
}
92+
throw error;
93+
} finally {
94+
// If it was a Promise we already finished in finally above
95+
// This covers sync functions
96+
try {
97+
span.finish();
98+
} catch {}
99+
}
100+
};
101+
102+
return descriptor;
103+
};
104+
}
105+
106+

0 commit comments

Comments
 (0)