From 80e4a9bfdab542bcb1584bba1443354fc7699ed8 Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Sat, 20 Dec 2025 10:56:01 +0000 Subject: [PATCH] Added tests for arithmetic operations on expression columns --- apps/webapp/app/v3/querySchemas.ts | 13 +- internal-packages/clickhouse/src/tsql.test.ts | 315 ++++++++++++++++++ .../tsql/src/query/printer.test.ts | 145 ++++++++ 3 files changed, 467 insertions(+), 6 deletions(-) diff --git a/apps/webapp/app/v3/querySchemas.ts b/apps/webapp/app/v3/querySchemas.ts index 404827112..e62c5cfc7 100644 --- a/apps/webapp/app/v3/querySchemas.ts +++ b/apps/webapp/app/v3/querySchemas.ts @@ -36,9 +36,10 @@ export const runsSchema: TableSchema = { name: "run_id", clickhouseName: "friendly_id", ...column("String", { - description: "Run ID (e.g., run_abc123)", + description: + "A unique ID for a run. They always start with `run_`, e.g., run_cm1a2b3c4d5e6f7g8h9i", customRenderType: "runId", - example: "run_1234abcd5678efgh", + example: "run_cm1a2b3c4d5e6f7g8h9i", }), }, environment_id: { @@ -218,18 +219,18 @@ export const runsSchema: TableSchema = { compute_cost: { name: "compute_cost", ...column("Float64", { - description: "Compute cost in cents", + description: "Compute cost in dollars", customRenderType: "cost", - example: "0.0025", + example: "0.000676", }), expression: "cost_in_cents / 100.0", }, invocation_cost: { name: "invocation_cost", ...column("Float64", { - description: "Invocation cost in cents", + description: "Invocation cost in dollars", customRenderType: "cost", - example: "0.001", + example: "0.000025", }), expression: "base_cost_in_cents / 100.0", }, diff --git a/internal-packages/clickhouse/src/tsql.test.ts b/internal-packages/clickhouse/src/tsql.test.ts index edeff2822..abdd6216e 100644 --- a/internal-packages/clickhouse/src/tsql.test.ts +++ b/internal-packages/clickhouse/src/tsql.test.ts @@ -1025,4 +1025,319 @@ describe("TSQL Virtual Column Tests", () => { expect(result?.rows?.[0].execution_duration).toBeNull(); } ); + + /** + * Schema with cost-based virtual columns that use division expressions + * This mimics the real runsSchema where invocation_cost = base_cost_in_cents / 100.0 + */ + const costExpressionSchema: TableSchema = { + name: "task_runs", + clickhouseName: "trigger_dev.task_runs_v2", + columns: { + run_id: { name: "run_id", ...column("String") }, + friendly_id: { name: "friendly_id", ...column("String") }, + status: { name: "status", ...column("String") }, + task_identifier: { name: "task_identifier", ...column("String") }, + queue: { name: "queue", ...column("String") }, + environment_id: { name: "environment_id", ...column("String") }, + environment_type: { name: "environment_type", ...column("String") }, + organization_id: { name: "organization_id", ...column("String") }, + project_id: { name: "project_id", ...column("String") }, + created_at: { name: "created_at", ...column("DateTime") }, + updated_at: { name: "updated_at", ...column("DateTime") }, + is_test: { name: "is_test", ...column("Bool") }, + tags: { name: "tags", ...column("Array(String)") }, + // Raw cost columns stored in ClickHouse + cost_in_cents: { name: "cost_in_cents", ...column("Float64") }, + base_cost_in_cents: { name: "base_cost_in_cents", ...column("Float64") }, + // Virtual column: compute_cost = cost_in_cents / 100.0 + compute_cost: { + name: "compute_cost", + ...column("Float64"), + expression: "cost_in_cents / 100.0", + description: "Compute cost in dollars (cost_in_cents / 100)", + }, + // Virtual column: invocation_cost = base_cost_in_cents / 100.0 + invocation_cost: { + name: "invocation_cost", + ...column("Float64"), + expression: "base_cost_in_cents / 100.0", + description: "Invocation cost in dollars (base_cost_in_cents / 100)", + }, + }, + tenantColumns: { + organizationId: "organization_id", + projectId: "project_id", + environmentId: "environment_id", + }, + }; + + clickhouseTest( + "should filter by expression column with division in WHERE clause", + async ({ clickhouseContainer }) => { + const client = new ClickhouseClient({ + name: "test", + url: clickhouseContainer.getConnectionUrl(), + }); + + const insert = insertTaskRuns(client, { async_insert: 0 }); + + // Insert runs with different cost values + // invocation_cost = base_cost_in_cents / 100.0 + // So base_cost_in_cents of 500 = invocation_cost of 5.0 + // And base_cost_in_cents of 50 = invocation_cost of 0.5 + await insert([ + createTaskRun({ + run_id: "run_cheap", + base_cost_in_cents: 50, // invocation_cost = 0.5 + cost_in_cents: 100, + }), + createTaskRun({ + run_id: "run_medium", + base_cost_in_cents: 200, // invocation_cost = 2.0 + cost_in_cents: 300, + }), + createTaskRun({ + run_id: "run_expensive", + base_cost_in_cents: 500, // invocation_cost = 5.0 + cost_in_cents: 800, + }), + ]); + + // Query runs where invocation_cost > 1.0 (base_cost_in_cents > 100) + // Should return run_medium (2.0) and run_expensive (5.0) + const [error, result] = await executeTSQL(client, { + name: "test-expression-division-where", + query: "SELECT run_id, invocation_cost FROM task_runs WHERE invocation_cost > 1.0", + schema: z.object({ run_id: z.string(), invocation_cost: z.number() }), + organizationId: "org_tenant1", + projectId: "proj_tenant1", + environmentId: "env_tenant1", + tableSchema: [costExpressionSchema], + }); + + expect(error).toBeNull(); + expect(result?.rows).toHaveLength(2); + expect(result?.rows?.map((r) => r.run_id).sort()).toEqual(["run_expensive", "run_medium"]); + } + ); + + clickhouseTest( + "should filter by expression column with greater-than-or-equals comparison", + async ({ clickhouseContainer }) => { + const client = new ClickhouseClient({ + name: "test", + url: clickhouseContainer.getConnectionUrl(), + }); + + const insert = insertTaskRuns(client, { async_insert: 0 }); + + await insert([ + createTaskRun({ + run_id: "run_1", + base_cost_in_cents: 100, // invocation_cost = 1.0 + }), + createTaskRun({ + run_id: "run_2", + base_cost_in_cents: 200, // invocation_cost = 2.0 + }), + createTaskRun({ + run_id: "run_3", + base_cost_in_cents: 50, // invocation_cost = 0.5 + }), + ]); + + // Query runs where invocation_cost >= 1.0 + // Should return run_1 (1.0) and run_2 (2.0) + const [error, result] = await executeTSQL(client, { + name: "test-expression-gte-where", + query: "SELECT run_id, invocation_cost FROM task_runs WHERE invocation_cost >= 1.0", + schema: z.object({ run_id: z.string(), invocation_cost: z.number() }), + organizationId: "org_tenant1", + projectId: "proj_tenant1", + environmentId: "env_tenant1", + tableSchema: [costExpressionSchema], + }); + + expect(error).toBeNull(); + expect(result?.rows).toHaveLength(2); + expect(result?.rows?.map((r) => r.run_id).sort()).toEqual(["run_1", "run_2"]); + } + ); + + clickhouseTest( + "should filter by expression column with less-than comparison", + async ({ clickhouseContainer }) => { + const client = new ClickhouseClient({ + name: "test", + url: clickhouseContainer.getConnectionUrl(), + }); + + const insert = insertTaskRuns(client, { async_insert: 0 }); + + await insert([ + createTaskRun({ + run_id: "run_small", + base_cost_in_cents: 50, // invocation_cost = 0.5 + }), + createTaskRun({ + run_id: "run_large", + base_cost_in_cents: 300, // invocation_cost = 3.0 + }), + ]); + + // Query runs where invocation_cost < 1.0 + // Should return only run_small (0.5) + const [error, result] = await executeTSQL(client, { + name: "test-expression-lt-where", + query: "SELECT run_id, invocation_cost FROM task_runs WHERE invocation_cost < 1.0", + schema: z.object({ run_id: z.string(), invocation_cost: z.number() }), + organizationId: "org_tenant1", + projectId: "proj_tenant1", + environmentId: "env_tenant1", + tableSchema: [costExpressionSchema], + }); + + expect(error).toBeNull(); + expect(result?.rows).toHaveLength(1); + expect(result?.rows?.[0].run_id).toBe("run_small"); + } + ); + + clickhouseTest( + "should filter by expression column with BETWEEN comparison", + async ({ clickhouseContainer }) => { + const client = new ClickhouseClient({ + name: "test", + url: clickhouseContainer.getConnectionUrl(), + }); + + const insert = insertTaskRuns(client, { async_insert: 0 }); + + await insert([ + createTaskRun({ + run_id: "run_low", + base_cost_in_cents: 50, // invocation_cost = 0.5 + }), + createTaskRun({ + run_id: "run_mid", + base_cost_in_cents: 150, // invocation_cost = 1.5 + }), + createTaskRun({ + run_id: "run_high", + base_cost_in_cents: 500, // invocation_cost = 5.0 + }), + ]); + + // Query runs where invocation_cost is between 1.0 and 2.0 + // Should return only run_mid (1.5) + const [error, result] = await executeTSQL(client, { + name: "test-expression-between-where", + query: + "SELECT run_id, invocation_cost FROM task_runs WHERE invocation_cost BETWEEN 1.0 AND 2.0", + schema: z.object({ run_id: z.string(), invocation_cost: z.number() }), + organizationId: "org_tenant1", + projectId: "proj_tenant1", + environmentId: "env_tenant1", + tableSchema: [costExpressionSchema], + }); + + expect(error).toBeNull(); + expect(result?.rows).toHaveLength(1); + expect(result?.rows?.[0].run_id).toBe("run_mid"); + } + ); + + clickhouseTest( + "should use expression column in complex WHERE with AND/OR", + async ({ clickhouseContainer }) => { + const client = new ClickhouseClient({ + name: "test", + url: clickhouseContainer.getConnectionUrl(), + }); + + const insert = insertTaskRuns(client, { async_insert: 0 }); + + await insert([ + createTaskRun({ + run_id: "run_a", + status: "COMPLETED_SUCCESSFULLY", + base_cost_in_cents: 100, // invocation_cost = 1.0 + }), + createTaskRun({ + run_id: "run_b", + status: "FAILED", + base_cost_in_cents: 300, // invocation_cost = 3.0 + }), + createTaskRun({ + run_id: "run_c", + status: "COMPLETED_SUCCESSFULLY", + base_cost_in_cents: 300, // invocation_cost = 3.0 + }), + ]); + + // Query completed runs with invocation_cost > 2.0 + const [error, result] = await executeTSQL(client, { + name: "test-expression-complex-where", + query: + "SELECT run_id FROM task_runs WHERE status = 'COMPLETED_SUCCESSFULLY' AND invocation_cost > 2.0", + schema: z.object({ run_id: z.string() }), + organizationId: "org_tenant1", + projectId: "proj_tenant1", + environmentId: "env_tenant1", + tableSchema: [costExpressionSchema], + }); + + expect(error).toBeNull(); + expect(result?.rows).toHaveLength(1); + expect(result?.rows?.[0].run_id).toBe("run_c"); + } + ); + + clickhouseTest( + "should filter by expression column comparing with large integer value", + async ({ clickhouseContainer }) => { + const client = new ClickhouseClient({ + name: "test", + url: clickhouseContainer.getConnectionUrl(), + }); + + const insert = insertTaskRuns(client, { async_insert: 0 }); + + // Insert runs with large base_cost_in_cents values + // So invocation_cost = base_cost_in_cents / 100.0 + // To get invocation_cost > 100, we need base_cost_in_cents > 10000 + await insert([ + createTaskRun({ + run_id: "run_small_cost", + base_cost_in_cents: 500, // invocation_cost = 5.0 + }), + createTaskRun({ + run_id: "run_medium_cost", + base_cost_in_cents: 5000, // invocation_cost = 50.0 + }), + createTaskRun({ + run_id: "run_large_cost", + base_cost_in_cents: 15000, // invocation_cost = 150.0 + }), + ]); + + // Query runs where invocation_cost > 100 + // Should only return run_large_cost (150.0) + const [error, result] = await executeTSQL(client, { + name: "test-expression-large-integer-where", + query: "SELECT run_id, invocation_cost FROM task_runs WHERE invocation_cost > 100", + schema: z.object({ run_id: z.string(), invocation_cost: z.number() }), + organizationId: "org_tenant1", + projectId: "proj_tenant1", + environmentId: "env_tenant1", + tableSchema: [costExpressionSchema], + }); + + expect(error).toBeNull(); + expect(result?.rows).toHaveLength(1); + expect(result?.rows?.[0].run_id).toBe("run_large_cost"); + expect(result?.rows?.[0].invocation_cost).toBeCloseTo(150.0, 1); + } + ); }); diff --git a/internal-packages/tsql/src/query/printer.test.ts b/internal-packages/tsql/src/query/printer.test.ts index 57b4051b3..d8b1f1f99 100644 --- a/internal-packages/tsql/src/query/printer.test.ts +++ b/internal-packages/tsql/src/query/printer.test.ts @@ -1088,6 +1088,151 @@ describe("Virtual columns", () => { }); }); +describe("Expression columns with division (cost/invocation_cost pattern)", () => { + /** + * Schema mimicking the real runsSchema with cost-based expression columns + * These use division expressions like base_cost_in_cents / 100.0 + */ + const costExpressionSchema: TableSchema = { + name: "runs", + clickhouseName: "trigger_dev.task_runs_v2", + columns: { + run_id: { name: "run_id", ...column("String") }, + status: { name: "status", ...column("String") }, + cost_in_cents: { name: "cost_in_cents", ...column("Float64") }, + base_cost_in_cents: { name: "base_cost_in_cents", ...column("Float64") }, + // Virtual column: compute_cost = cost_in_cents / 100.0 + compute_cost: { + name: "compute_cost", + ...column("Float64"), + expression: "cost_in_cents / 100.0", + description: "Compute cost in dollars", + }, + // Virtual column: invocation_cost = base_cost_in_cents / 100.0 + invocation_cost: { + name: "invocation_cost", + ...column("Float64"), + expression: "base_cost_in_cents / 100.0", + description: "Invocation cost in dollars", + }, + organization_id: { name: "organization_id", ...column("String") }, + project_id: { name: "project_id", ...column("String") }, + environment_id: { name: "environment_id", ...column("String") }, + }, + tenantColumns: { + organizationId: "organization_id", + projectId: "project_id", + environmentId: "environment_id", + }, + }; + + function createCostExpressionContext() { + const schema = createSchemaRegistry([costExpressionSchema]); + return createPrinterContext({ + organizationId: "org_test", + projectId: "proj_test", + environmentId: "env_test", + schema, + }); + } + + describe("WHERE clause with division expression columns", () => { + it("should expand invocation_cost > 100 to (base_cost_in_cents / 100.0) > 100", () => { + const ctx = createCostExpressionContext(); + const { sql } = printQuery("SELECT * FROM runs WHERE invocation_cost > 100", ctx); + + // Virtual column should be expanded to its expression in the comparison + expect(sql).toContain("greater((base_cost_in_cents / 100.0), 100)"); + }); + + it("should expand invocation_cost >= 0.01 correctly", () => { + const ctx = createCostExpressionContext(); + const { sql } = printQuery("SELECT * FROM runs WHERE invocation_cost >= 0.01", ctx); + + expect(sql).toContain("greaterOrEquals((base_cost_in_cents / 100.0), 0.01)"); + }); + + it("should expand invocation_cost < 1.5 correctly", () => { + const ctx = createCostExpressionContext(); + const { sql } = printQuery("SELECT * FROM runs WHERE invocation_cost < 1.5", ctx); + + expect(sql).toContain("less((base_cost_in_cents / 100.0), 1.5)"); + }); + + it("should expand invocation_cost = 0 correctly", () => { + const ctx = createCostExpressionContext(); + const { sql } = printQuery("SELECT * FROM runs WHERE invocation_cost = 0", ctx); + + expect(sql).toContain("equals((base_cost_in_cents / 100.0), 0)"); + }); + + it("should expand compute_cost in BETWEEN correctly", () => { + const ctx = createCostExpressionContext(); + const { sql } = printQuery( + "SELECT * FROM runs WHERE compute_cost BETWEEN 1.0 AND 10.0", + ctx + ); + + expect(sql).toContain("(cost_in_cents / 100.0) BETWEEN 1 AND 10"); + }); + + it("should handle multiple expression columns in WHERE with AND", () => { + const ctx = createCostExpressionContext(); + const { sql } = printQuery( + "SELECT * FROM runs WHERE invocation_cost > 0.5 AND compute_cost < 5.0", + ctx + ); + + expect(sql).toContain("greater((base_cost_in_cents / 100.0), 0.5)"); + expect(sql).toContain("less((cost_in_cents / 100.0), 5)"); + }); + + it("should handle expression column with OR condition", () => { + const ctx = createCostExpressionContext(); + const { sql } = printQuery( + "SELECT * FROM runs WHERE invocation_cost > 10 OR status = 'COMPLETED'", + ctx + ); + + expect(sql).toContain("or("); + expect(sql).toContain("greater((base_cost_in_cents / 100.0), 10)"); + }); + }); + + describe("SELECT clause with division expression columns", () => { + it("should expand invocation_cost in SELECT with alias", () => { + const ctx = createCostExpressionContext(); + const { sql } = printQuery("SELECT run_id, invocation_cost FROM runs", ctx); + + expect(sql).toContain("(base_cost_in_cents / 100.0) AS invocation_cost"); + }); + + it("should expand both cost columns correctly", () => { + const ctx = createCostExpressionContext(); + const { sql } = printQuery("SELECT compute_cost, invocation_cost FROM runs", ctx); + + expect(sql).toContain("(cost_in_cents / 100.0) AS compute_cost"); + expect(sql).toContain("(base_cost_in_cents / 100.0) AS invocation_cost"); + }); + }); + + describe("ORDER BY clause with division expression columns", () => { + it("should expand invocation_cost in ORDER BY DESC", () => { + const ctx = createCostExpressionContext(); + const { sql } = printQuery("SELECT * FROM runs ORDER BY invocation_cost DESC", ctx); + + expect(sql).toContain("ORDER BY (base_cost_in_cents / 100.0) DESC"); + }); + + it("should expand compute_cost in ORDER BY ASC", () => { + const ctx = createCostExpressionContext(); + const { sql } = printQuery("SELECT * FROM runs ORDER BY compute_cost ASC", ctx); + + expect(sql).toContain("ORDER BY (cost_in_cents / 100.0) ASC"); + }); + }); +}); + describe("Column metadata", () => { /** * Schema with customRenderType for testing