Expose ClickHouse query stats from tsql queries

This commit is contained in:
Matt Aitken
2025-12-18 15:31:05 +00:00
parent 6c261ae369
commit 6ca631cbb5
6 changed files with 363 additions and 85 deletions
@@ -16,9 +16,11 @@ import type {
ClickhouseQueryBuilderFastFunction,
ClickhouseQueryBuilderFunction,
ClickhouseQueryFunction,
ClickhouseQueryWithStatsFunction,
ClickhouseReader,
ClickhouseWriter,
ColumnExpression,
QueryStats,
} from "./types.js";
import { generateErrorMessage } from "zod-error";
import { Logger, type LogLevel } from "@trigger.dev/core/logger";
@@ -229,6 +231,177 @@ export class ClickhouseClient implements ClickhouseReader, ClickhouseWriter {
};
}
public queryWithStats<TIn extends z.ZodSchema<any>, TOut extends z.ZodSchema<any>>(req: {
/**
* The name of the operation.
* This will be used to identify the operation in the span.
*/
name: string;
/**
* The SQL query to run.
* Use {paramName: Type} to define parameters
* Example: `SELECT * FROM table WHERE id = {id: String}`
*/
query: string;
/**
* The schema of the parameters
* Example: z.object({ id: z.string() })
*/
params?: TIn;
/**
* The schema of the output of each row
* Example: z.object({ id: z.string() })
*/
schema: TOut;
/**
* The settings to use for the query.
* These will be merged with the default settings.
*/
settings?: ClickHouseSettings;
}): ClickhouseQueryWithStatsFunction<z.input<TIn>, z.output<TOut>> {
return async (params, options) => {
const queryId = randomUUID();
return await startSpan(this.tracer, "queryWithStats", async (span) => {
this.logger.debug("Querying clickhouse with stats", {
name: req.name,
query: req.query.replace(/\s+/g, " "),
params,
settings: req.settings,
attributes: options?.attributes,
queryId,
});
span.setAttributes({
"clickhouse.clientName": this.name,
"clickhouse.operationName": req.name,
"clickhouse.queryId": queryId,
...flattenAttributes(req.settings, "clickhouse.settings"),
...flattenAttributes(options?.attributes),
});
const validParams = req.params?.safeParse(params);
if (validParams?.error) {
recordSpanError(span, validParams.error);
this.logger.error("Error parsing query params", {
name: req.name,
error: validParams.error,
query: req.query,
params,
queryId,
});
return [
new QueryError(`Bad params: ${generateErrorMessage(validParams.error.issues)}`, {
query: req.query,
}),
null,
];
}
let unparsedRows: Array<TOut> = [];
const [clickhouseError, res] = await tryCatch(
this.client.query({
query: req.query,
query_params: validParams?.data,
format: "JSONEachRow",
query_id: queryId,
...options?.params,
clickhouse_settings: {
...req.settings,
...options?.params?.clickhouse_settings,
},
})
);
if (clickhouseError) {
this.logger.error("Error querying clickhouse", {
name: req.name,
error: clickhouseError,
query: req.query,
params,
queryId,
});
recordClickhouseError(span, clickhouseError);
return [
new QueryError(`Unable to query clickhouse: ${clickhouseError.message}`, {
query: req.query,
}),
null,
];
}
unparsedRows = await res.json();
span.setAttributes({
"clickhouse.query_id": res.query_id,
...flattenAttributes(res.response_headers, "clickhouse.response_headers"),
});
// Parse the summary header to get stats
const summaryHeader = res.response_headers["x-clickhouse-summary"];
let stats: QueryStats = {
read_rows: "0",
read_bytes: "0",
written_rows: "0",
written_bytes: "0",
total_rows_to_read: "0",
result_rows: "0",
result_bytes: "0",
elapsed_ns: "0",
};
if (typeof summaryHeader === "string") {
const parsedSummary = JSON.parse(summaryHeader);
stats = {
read_rows: parsedSummary.read_rows ?? "0",
read_bytes: parsedSummary.read_bytes ?? "0",
written_rows: parsedSummary.written_rows ?? "0",
written_bytes: parsedSummary.written_bytes ?? "0",
total_rows_to_read: parsedSummary.total_rows_to_read ?? "0",
result_rows: parsedSummary.result_rows ?? "0",
result_bytes: parsedSummary.result_bytes ?? "0",
elapsed_ns: parsedSummary.elapsed_ns ?? "0",
};
span.setAttributes({
...flattenAttributes(parsedSummary, "clickhouse.summary"),
});
}
const parsed = z.array(req.schema).safeParse(unparsedRows);
if (parsed.error) {
this.logger.error("Error parsing clickhouse query result", {
name: req.name,
error: parsed.error,
query: req.query,
params,
queryId,
});
const queryError = new QueryError(generateErrorMessage(parsed.error.issues), {
query: req.query,
});
recordSpanError(span, queryError);
return [queryError, null];
}
span.setAttributes({
"clickhouse.rows": unparsedRows.length,
});
return [null, { rows: parsed.data, stats }];
});
};
}
public queryFast<TOut extends Record<string, any>, TParams extends Record<string, any>>(req: {
name: string;
query: string;
@@ -3,9 +3,10 @@ import { InsertError, QueryError } from "./errors.js";
import {
ClickhouseQueryBuilderFastFunction,
ClickhouseQueryBuilderFunction,
ClickhouseReader,
ClickhouseWriter,
QueryResultWithStats,
} from "./types.js";
import { ClickhouseReader } from "./types.js";
import { z } from "zod";
import { ClickHouseSettings, InsertResult } from "@clickhouse/client";
import { ClickhouseQueryBuilder, ClickhouseQueryFastBuilder } from "./queryBuilder.js";
@@ -51,6 +52,37 @@ export class NoopClient implements ClickhouseReader, ClickhouseWriter {
};
}
public queryWithStats<TIn extends z.ZodSchema<any>, TOut extends z.ZodSchema<any>>(req: {
query: string;
params?: TIn;
schema: TOut;
}): (params: z.input<TIn>) => Promise<Result<QueryResultWithStats<z.output<TOut>>, QueryError>> {
return async (params: z.input<TIn>) => {
const validParams = req.params?.safeParse(params);
if (validParams?.error) {
return [new QueryError(`Bad params: ${validParams.error.message}`, { query: "" }), null];
}
return [
null,
{
rows: [],
stats: {
read_rows: "0",
read_bytes: "0",
written_rows: "0",
written_bytes: "0",
total_rows_to_read: "0",
result_rows: "0",
result_bytes: "0",
elapsed_ns: "0",
},
},
];
};
}
public queryFast<TOut extends Record<string, any>, TParams extends Record<string, any>>(req: {
name: string;
query: string;
@@ -13,10 +13,12 @@ import {
type TableSchema,
type QuerySettings,
} from "@internal/tsql";
import type { ClickhouseReader } from "./types.js";
import type { ClickhouseReader, QueryStats } from "./types.js";
import { QueryError } from "./errors.js";
import type { OutputColumnMetadata } from "@internal/tsql";
export type { QueryStats };
export type { TableSchema, QuerySettings };
/**
@@ -56,6 +58,7 @@ export interface ExecuteTSQLOptions<TOut extends z.ZodSchema> {
export interface TSQLQuerySuccess<T> {
rows: T[];
columns: OutputColumnMetadata[];
stats: QueryStats;
}
/**
@@ -99,8 +102,8 @@ export async function executeTSQL<TOut extends z.ZodSchema>(
settings: options.querySettings,
});
// 2. Execute the query
const queryFn = reader.query({
// 2. Execute the query with stats
const queryFn = reader.queryWithStats({
name: options.name,
query: sql,
params: z.record(z.any()),
@@ -108,22 +111,24 @@ export async function executeTSQL<TOut extends z.ZodSchema>(
settings: options.clickhouseSettings,
});
const [error, rows] = await queryFn(params);
const [error, result] = await queryFn(params);
if (error) {
return [error, null];
}
const { rows, stats } = result;
// 3. Transform result values if enabled
if (shouldTransformValues && rows) {
const transformedRows = transformResults(
rows as Record<string, unknown>[],
options.tableSchema
);
return [null, { rows: transformedRows as z.output<TOut>[], columns }];
return [null, { rows: transformedRows as z.output<TOut>[], columns, stats }];
}
return [null, { rows: rows ?? [], columns }];
return [null, { rows: rows ?? [], columns, stats }];
} catch (error) {
if (error instanceof Error) {
return [new QueryError(error.message, { query: options.query }), null];
@@ -13,6 +13,36 @@ export type ClickhouseQueryFunction<TInput, TOutput> = (
}
) => Promise<Result<TOutput[], QueryError>>;
/**
* Query statistics returned by ClickHouse
*/
export interface QueryStats {
read_rows: string;
read_bytes: string;
written_rows: string;
written_bytes: string;
total_rows_to_read: string;
result_rows: string;
result_bytes: string;
elapsed_ns: string;
}
/**
* Result type for queries that include stats
*/
export interface QueryResultWithStats<TOutput> {
rows: TOutput[];
stats: QueryStats;
}
export type ClickhouseQueryWithStatsFunction<TInput, TOutput> = (
params: TInput,
options?: {
attributes?: Record<string, string | number | boolean>;
params?: BaseQueryParams;
}
) => Promise<Result<QueryResultWithStats<TOutput>, QueryError>>;
export type ClickhouseQueryBuilderFunction<TOutput> = (options?: {
settings?: ClickHouseSettings;
}) => ClickhouseQueryBuilder<TOutput>;
@@ -56,6 +86,39 @@ export interface ClickhouseReader {
settings?: ClickHouseSettings;
}): ClickhouseQueryFunction<z.input<TIn>, z.output<TOut>>;
/**
* Execute a query and return both rows and query statistics.
* Same as `query` but includes ClickHouse query stats in the result.
*/
queryWithStats<TIn extends z.ZodSchema<any>, TOut extends z.ZodSchema<any>>(req: {
/**
* The name of the operation.
* This will be used to identify the operation in the span.
*/
name: string;
/**
* The SQL query to run.
* Use {paramName: Type} to define parameters
* Example: `SELECT * FROM table WHERE id = {id: String}`
*/
query: string;
/**
* The schema of the parameters
* Example: z.object({ id: z.string() })
*/
params?: TIn;
/**
* The schema of the output of each row
* Example: z.object({ id: z.string() })
*/
schema: TOut;
/**
* The settings to use for the query.
* These will be merged with the default settings.
*/
settings?: ClickHouseSettings;
}): ClickhouseQueryWithStatsFunction<z.input<TIn>, z.output<TOut>>;
queryFast<TOut extends Record<string, any>, TParams extends Record<string, any>>(req: {
/**
* The name of the operation.
@@ -39,6 +39,7 @@ export {
type TableSchema,
type TSQLQueryResult,
type TSQLQuerySuccess,
type QueryStats,
} from "./client/tsql.js";
export type { OutputColumnMetadata } from "@internal/tsql";
+82 -78
View File
@@ -102,7 +102,7 @@ describe("TSQL Integration Tests", () => {
expect(insertError).toBeNull();
// Execute TSQL query
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-simple-select",
query: "SELECT run_id, status FROM task_runs",
schema: z.object({ run_id: z.string(), status: z.string() }),
@@ -113,14 +113,18 @@ describe("TSQL Integration Tests", () => {
});
expect(error).toBeNull();
expect(rows).toHaveLength(3);
expect(rows).toEqual(
expect(result?.rows).toHaveLength(3);
expect(result?.rows).toEqual(
expect.arrayContaining([
expect.objectContaining({ run_id: "run_test1", status: "COMPLETED_SUCCESSFULLY" }),
expect.objectContaining({ run_id: "run_test2", status: "PENDING" }),
expect.objectContaining({ run_id: "run_test3", status: "COMPLETED_SUCCESSFULLY" }),
])
);
// Verify stats are returned
expect(result?.stats).toBeDefined();
expect(result?.stats.read_rows).toBeDefined();
expect(result?.stats.elapsed_ns).toBeDefined();
});
clickhouseTest("should filter with WHERE clause", async ({ clickhouseContainer }) => {
@@ -137,7 +141,7 @@ describe("TSQL Integration Tests", () => {
createTaskRun({ run_id: "run_filter3", status: "COMPLETED_SUCCESSFULLY" }),
]);
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-where-clause",
query: "SELECT run_id, status FROM task_runs WHERE status = 'COMPLETED_SUCCESSFULLY'",
schema: z.object({ run_id: z.string(), status: z.string() }),
@@ -148,8 +152,8 @@ describe("TSQL Integration Tests", () => {
});
expect(error).toBeNull();
expect(rows).toHaveLength(2);
expect(rows?.every((r) => r.status === "COMPLETED_SUCCESSFULLY")).toBe(true);
expect(result?.rows).toHaveLength(2);
expect(result?.rows?.every((r) => r.status === "COMPLETED_SUCCESSFULLY")).toBe(true);
});
clickhouseTest("should enforce tenant isolation", async ({ clickhouseContainer }) => {
@@ -189,7 +193,7 @@ describe("TSQL Integration Tests", () => {
]);
// Query as tenant1 - should only see tenant1's data
const [error1, rows1] = await executeTSQL(client, {
const [error1, result1] = await executeTSQL(client, {
name: "test-tenant-isolation-1",
query: "SELECT run_id FROM task_runs",
schema: z.object({ run_id: z.string() }),
@@ -200,11 +204,11 @@ describe("TSQL Integration Tests", () => {
});
expect(error1).toBeNull();
expect(rows1).toHaveLength(2);
expect(rows1?.map((r) => r.run_id).sort()).toEqual(["run_tenant1_a", "run_tenant1_b"]);
expect(result1?.rows).toHaveLength(2);
expect(result1?.rows?.map((r) => r.run_id).sort()).toEqual(["run_tenant1_a", "run_tenant1_b"]);
// Query as tenant2 - should only see tenant2's data
const [error2, rows2] = await executeTSQL(client, {
const [error2, result2] = await executeTSQL(client, {
name: "test-tenant-isolation-2",
query: "SELECT run_id FROM task_runs",
schema: z.object({ run_id: z.string() }),
@@ -215,8 +219,8 @@ describe("TSQL Integration Tests", () => {
});
expect(error2).toBeNull();
expect(rows2).toHaveLength(2);
expect(rows2?.map((r) => r.run_id).sort()).toEqual(["run_tenant2_a", "run_tenant2_b"]);
expect(result2?.rows).toHaveLength(2);
expect(result2?.rows?.map((r) => r.run_id).sort()).toEqual(["run_tenant2_a", "run_tenant2_b"]);
});
clickhouseTest(
@@ -246,7 +250,7 @@ describe("TSQL Integration Tests", () => {
]);
// Attacker tries to access victim's data with OR 1=1
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-cross-tenant-attack",
query: "SELECT run_id, status FROM task_runs WHERE status = 'COMPLETED' OR 1=1",
schema: z.object({ run_id: z.string(), status: z.string() }),
@@ -258,9 +262,9 @@ describe("TSQL Integration Tests", () => {
expect(error).toBeNull();
// Should only get attacker's own data, not victim's
expect(rows).toHaveLength(1);
expect(rows?.[0].run_id).toBe("run_attacker");
expect(rows?.find((r) => r.run_id === "run_secret")).toBeUndefined();
expect(result?.rows).toHaveLength(1);
expect(result?.rows?.[0].run_id).toBe("run_attacker");
expect(result?.rows?.find((r) => r.run_id === "run_secret")).toBeUndefined();
}
);
@@ -279,7 +283,7 @@ describe("TSQL Integration Tests", () => {
createTaskRun({ run_id: "run_agg4", status: "FAILED" }),
]);
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-aggregation",
query:
"SELECT status, count(*) as cnt FROM task_runs GROUP BY status ORDER BY cnt DESC, status ASC",
@@ -291,10 +295,10 @@ describe("TSQL Integration Tests", () => {
});
expect(error).toBeNull();
expect(rows).toHaveLength(3);
expect(rows?.[0]).toEqual({ status: "COMPLETED_SUCCESSFULLY", cnt: 2 });
expect(result?.rows).toHaveLength(3);
expect(result?.rows?.[0]).toEqual({ status: "COMPLETED_SUCCESSFULLY", cnt: 2 });
// The remaining rows have cnt=1, check they're both present
expect(rows).toEqual(
expect(result?.rows).toEqual(
expect.arrayContaining([
{ status: "PENDING", cnt: 1 },
{ status: "FAILED", cnt: 1 },
@@ -317,7 +321,7 @@ describe("TSQL Integration Tests", () => {
createTaskRun({ run_id: "run_order3", created_at: now - 2000 }),
]);
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-order-limit",
query: "SELECT run_id FROM task_runs ORDER BY created_at DESC LIMIT 2",
schema: z.object({ run_id: z.string() }),
@@ -328,9 +332,9 @@ describe("TSQL Integration Tests", () => {
});
expect(error).toBeNull();
expect(rows).toHaveLength(2);
expect(rows?.[0].run_id).toBe("run_order2"); // Most recent
expect(rows?.[1].run_id).toBe("run_order3"); // Second most recent
expect(result?.rows).toHaveLength(2);
expect(result?.rows?.[0].run_id).toBe("run_order2"); // Most recent
expect(result?.rows?.[1].run_id).toBe("run_order3"); // Second most recent
});
clickhouseTest("should reject unknown tables", async ({ clickhouseContainer }) => {
@@ -339,7 +343,7 @@ describe("TSQL Integration Tests", () => {
url: clickhouseContainer.getConnectionUrl(),
});
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-unknown-table",
query: "SELECT * FROM unknown_table",
schema: z.object({ id: z.string() }),
@@ -351,7 +355,7 @@ describe("TSQL Integration Tests", () => {
expect(error).not.toBeNull();
expect(error?.message).toContain("unknown_table");
expect(rows).toBeNull();
expect(result).toBeNull();
});
clickhouseTest("should work with createTSQLExecutor", async ({ clickhouseContainer }) => {
@@ -370,7 +374,7 @@ describe("TSQL Integration Tests", () => {
// Create a reusable executor
const tsql = createTSQLExecutor(client, [taskRunsSchema]);
const [error, rows] = await tsql.execute({
const [error, result] = await tsql.execute({
name: "test-executor",
query: "SELECT run_id, status FROM task_runs WHERE status = 'PENDING'",
schema: z.object({ run_id: z.string(), status: z.string() }),
@@ -380,8 +384,8 @@ describe("TSQL Integration Tests", () => {
});
expect(error).toBeNull();
expect(rows).toHaveLength(1);
expect(rows?.[0]).toEqual({ run_id: "run_executor2", status: "PENDING" });
expect(result?.rows).toHaveLength(1);
expect(result?.rows?.[0]).toEqual({ run_id: "run_executor2", status: "PENDING" });
});
clickhouseTest("should handle string injection attempts", async ({ clickhouseContainer }) => {
@@ -398,7 +402,7 @@ describe("TSQL Integration Tests", () => {
]);
// Query with a "malicious" value that looks like SQL
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-injection",
query: "SELECT run_id, status FROM task_runs WHERE status = 'DROP TABLE task_runs'",
schema: z.object({ run_id: z.string(), status: z.string() }),
@@ -410,8 +414,8 @@ describe("TSQL Integration Tests", () => {
expect(error).toBeNull();
// Should find the row with the literal string value, not execute SQL
expect(rows).toHaveLength(1);
expect(rows?.[0].status).toBe("DROP TABLE task_runs");
expect(result?.rows).toHaveLength(1);
expect(result?.rows?.[0].status).toBe("DROP TABLE task_runs");
});
clickhouseTest("should handle IN queries", async ({ clickhouseContainer }) => {
@@ -429,7 +433,7 @@ describe("TSQL Integration Tests", () => {
createTaskRun({ run_id: "run_in4", status: "CANCELLED" }),
]);
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-in-query",
query:
"SELECT run_id, status FROM task_runs WHERE status IN ('COMPLETED_SUCCESSFULLY', 'FAILED')",
@@ -441,8 +445,8 @@ describe("TSQL Integration Tests", () => {
});
expect(error).toBeNull();
expect(rows).toHaveLength(2);
expect(rows?.map((r) => r.status).sort()).toEqual(["COMPLETED_SUCCESSFULLY", "FAILED"]);
expect(result?.rows).toHaveLength(2);
expect(result?.rows?.map((r) => r.status).sort()).toEqual(["COMPLETED_SUCCESSFULLY", "FAILED"]);
});
clickhouseTest("should handle LIKE queries", async ({ clickhouseContainer }) => {
@@ -459,7 +463,7 @@ describe("TSQL Integration Tests", () => {
createTaskRun({ run_id: "run_like3", task_identifier: "sms/send" }),
]);
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-like-query",
query: "SELECT run_id, task_identifier FROM task_runs WHERE task_identifier LIKE 'email%'",
schema: z.object({ run_id: z.string(), task_identifier: z.string() }),
@@ -470,8 +474,8 @@ describe("TSQL Integration Tests", () => {
});
expect(error).toBeNull();
expect(rows).toHaveLength(2);
expect(rows?.every((r) => r.task_identifier.startsWith("email"))).toBe(true);
expect(result?.rows).toHaveLength(2);
expect(result?.rows?.every((r) => r.task_identifier.startsWith("email"))).toBe(true);
});
});
@@ -522,7 +526,7 @@ describe("TSQL Optional Tenant Filter Tests", () => {
]);
// Query across all projects (omit projectId and environmentId)
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-cross-project-query",
query: "SELECT run_id FROM task_runs",
schema: z.object({ run_id: z.string() }),
@@ -532,8 +536,8 @@ describe("TSQL Optional Tenant Filter Tests", () => {
});
expect(error).toBeNull();
expect(rows).toHaveLength(4); // All runs from org_multi
expect(rows?.map((r) => r.run_id).sort()).toEqual([
expect(result?.rows).toHaveLength(4); // All runs from org_multi
expect(result?.rows?.map((r) => r.run_id).sort()).toEqual([
"run_proj1_a",
"run_proj1_b",
"run_proj2_a",
@@ -582,7 +586,7 @@ describe("TSQL Optional Tenant Filter Tests", () => {
]);
// Query across all environments (omit environmentId only)
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-cross-env-query",
query: "SELECT run_id FROM task_runs",
schema: z.object({ run_id: z.string() }),
@@ -593,8 +597,8 @@ describe("TSQL Optional Tenant Filter Tests", () => {
});
expect(error).toBeNull();
expect(rows).toHaveLength(3); // All runs from proj_envtest across all envs
expect(rows?.map((r) => r.run_id).sort()).toEqual([
expect(result?.rows).toHaveLength(3); // All runs from proj_envtest across all envs
expect(result?.rows?.map((r) => r.run_id).sort()).toEqual([
"run_dev_1",
"run_prod_1",
"run_staging_1",
@@ -641,7 +645,7 @@ describe("TSQL Optional Tenant Filter Tests", () => {
]);
// Query org1 across all projects - should NOT see org2's data
const [error1, rows1] = await executeTSQL(client, {
const [error1, result1] = await executeTSQL(client, {
name: "test-org-isolation-1",
query: "SELECT run_id FROM task_runs",
schema: z.object({ run_id: z.string() }),
@@ -651,11 +655,11 @@ describe("TSQL Optional Tenant Filter Tests", () => {
});
expect(error1).toBeNull();
expect(rows1).toHaveLength(2);
expect(rows1?.map((r) => r.run_id).sort()).toEqual(["run_org1_a", "run_org1_b"]);
expect(result1?.rows).toHaveLength(2);
expect(result1?.rows?.map((r) => r.run_id).sort()).toEqual(["run_org1_a", "run_org1_b"]);
// Query org2 across all projects - should NOT see org1's data
const [error2, rows2] = await executeTSQL(client, {
const [error2, result2] = await executeTSQL(client, {
name: "test-org-isolation-2",
query: "SELECT run_id FROM task_runs",
schema: z.object({ run_id: z.string() }),
@@ -665,8 +669,8 @@ describe("TSQL Optional Tenant Filter Tests", () => {
});
expect(error2).toBeNull();
expect(rows2).toHaveLength(2);
expect(rows2?.map((r) => r.run_id).sort()).toEqual(["run_org2_a", "run_org2_b"]);
expect(result2?.rows).toHaveLength(2);
expect(result2?.rows?.map((r) => r.run_id).sort()).toEqual(["run_org2_a", "run_org2_b"]);
}
);
@@ -698,7 +702,7 @@ describe("TSQL Optional Tenant Filter Tests", () => {
]);
// Attacker tries to use OR 1=1 to bypass org filter
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-or-bypass-attempt",
query: "SELECT run_id, status FROM task_runs WHERE status = 'COMPLETED' OR 1=1",
schema: z.object({ run_id: z.string(), status: z.string() }),
@@ -709,9 +713,9 @@ describe("TSQL Optional Tenant Filter Tests", () => {
expect(error).toBeNull();
// Should only get attacker's data, not victim's
expect(rows).toHaveLength(1);
expect(rows?.[0].run_id).toBe("run_attacker");
expect(rows?.find((r) => r.run_id === "run_victim")).toBeUndefined();
expect(result?.rows).toHaveLength(1);
expect(result?.rows?.[0].run_id).toBe("run_attacker");
expect(result?.rows?.find((r) => r.run_id === "run_victim")).toBeUndefined();
}
);
@@ -743,7 +747,7 @@ describe("TSQL Optional Tenant Filter Tests", () => {
const tsql = createTSQLExecutor(client, [taskRunsSchema]);
// Use executor with org-only filter
const [error, rows] = await tsql.execute({
const [error, result] = await tsql.execute({
name: "test-executor-optional",
query: "SELECT run_id FROM task_runs",
schema: z.object({ run_id: z.string() }),
@@ -752,8 +756,8 @@ describe("TSQL Optional Tenant Filter Tests", () => {
});
expect(error).toBeNull();
expect(rows).toHaveLength(2);
expect(rows?.map((r) => r.run_id).sort()).toEqual(["run_exec_1", "run_exec_2"]);
expect(result?.rows).toHaveLength(2);
expect(result?.rows?.map((r) => r.run_id).sort()).toEqual(["run_exec_1", "run_exec_2"]);
}
);
});
@@ -827,7 +831,7 @@ describe("TSQL Virtual Column Tests", () => {
}),
]);
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-virtual-column-select",
query: "SELECT run_id, execution_duration, usage_duration_seconds FROM task_runs",
schema: z.object({
@@ -842,11 +846,11 @@ describe("TSQL Virtual Column Tests", () => {
});
expect(error).toBeNull();
expect(rows).toHaveLength(1);
expect(result?.rows).toHaveLength(1);
// execution_duration should be approximately 5000ms (difference between started_at and completed_at)
expect(rows?.[0].execution_duration).toBeCloseTo(5000, -2); // within 100ms tolerance
expect(result?.rows?.[0].execution_duration).toBeCloseTo(5000, -2); // within 100ms tolerance
// usage_duration_seconds should be 3.5 (3500ms / 1000)
expect(rows?.[0].usage_duration_seconds).toBeCloseTo(3.5, 1);
expect(result?.rows?.[0].usage_duration_seconds).toBeCloseTo(3.5, 1);
}
);
@@ -881,7 +885,7 @@ describe("TSQL Virtual Column Tests", () => {
]);
// Query runs with execution_duration > 5000ms (5 seconds)
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-virtual-column-where",
query: "SELECT run_id FROM task_runs WHERE execution_duration > 5000",
schema: z.object({ run_id: z.string() }),
@@ -892,8 +896,8 @@ describe("TSQL Virtual Column Tests", () => {
});
expect(error).toBeNull();
expect(rows).toHaveLength(2);
expect(rows?.map((r) => r.run_id).sort()).toEqual(["run_long", "run_very_long"]);
expect(result?.rows).toHaveLength(2);
expect(result?.rows?.map((r) => r.run_id).sort()).toEqual(["run_long", "run_very_long"]);
}
);
@@ -923,7 +927,7 @@ describe("TSQL Virtual Column Tests", () => {
]);
// Order by usage_duration_seconds descending (virtual column)
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-virtual-column-order",
query:
"SELECT run_id, usage_duration_seconds FROM task_runs ORDER BY usage_duration_seconds DESC",
@@ -938,14 +942,14 @@ describe("TSQL Virtual Column Tests", () => {
});
expect(error).toBeNull();
expect(rows).toHaveLength(3);
expect(result?.rows).toHaveLength(3);
// Should be ordered by usage_duration_seconds DESC: b (3), c (2), a (1)
expect(rows?.[0].run_id).toBe("run_order_b");
expect(rows?.[0].usage_duration_seconds).toBeCloseTo(3.0, 1);
expect(rows?.[1].run_id).toBe("run_order_c");
expect(rows?.[1].usage_duration_seconds).toBeCloseTo(2.0, 1);
expect(rows?.[2].run_id).toBe("run_order_a");
expect(rows?.[2].usage_duration_seconds).toBeCloseTo(1.0, 1);
expect(result?.rows?.[0].run_id).toBe("run_order_b");
expect(result?.rows?.[0].usage_duration_seconds).toBeCloseTo(3.0, 1);
expect(result?.rows?.[1].run_id).toBe("run_order_c");
expect(result?.rows?.[1].usage_duration_seconds).toBeCloseTo(2.0, 1);
expect(result?.rows?.[2].run_id).toBe("run_order_a");
expect(result?.rows?.[2].usage_duration_seconds).toBeCloseTo(1.0, 1);
});
clickhouseTest(
@@ -966,7 +970,7 @@ describe("TSQL Virtual Column Tests", () => {
]);
// Use virtual column with custom alias
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-virtual-column-alias",
query: "SELECT run_id, usage_duration_seconds AS dur_sec FROM task_runs",
schema: z.object({
@@ -980,8 +984,8 @@ describe("TSQL Virtual Column Tests", () => {
});
expect(error).toBeNull();
expect(rows).toHaveLength(1);
expect(rows?.[0].dur_sec).toBeCloseTo(5.0, 1);
expect(result?.rows).toHaveLength(1);
expect(result?.rows?.[0].dur_sec).toBeCloseTo(5.0, 1);
}
);
@@ -1002,7 +1006,7 @@ describe("TSQL Virtual Column Tests", () => {
}),
]);
const [error, rows] = await executeTSQL(client, {
const [error, result] = await executeTSQL(client, {
name: "test-virtual-column-null",
query: "SELECT run_id, execution_duration FROM task_runs",
schema: z.object({
@@ -1016,9 +1020,9 @@ describe("TSQL Virtual Column Tests", () => {
});
expect(error).toBeNull();
expect(rows).toHaveLength(1);
expect(result?.rows).toHaveLength(1);
// execution_duration should be null when started_at or completed_at is null
expect(rows?.[0].execution_duration).toBeNull();
expect(result?.rows?.[0].execution_duration).toBeNull();
}
);
});