diff --git a/services/cubejs/index.js b/services/cubejs/index.js index 5a09747d..79524aa0 100644 --- a/services/cubejs/index.js +++ b/services/cubejs/index.js @@ -86,7 +86,7 @@ const options = { scheduledRefreshContexts, externalDbType: "cubestore", externalDriverFactory, - cacheAndQueueDriver: "cubestore", + cacheAndQueueDriver: process.env.CUBEJS_CACHE_AND_QUEUE_DRIVER || "cubestore", logger: logging, // sql server diff --git a/services/cubejs/src/__tests__/partitionPruning.test.js b/services/cubejs/src/__tests__/partitionPruning.test.js new file mode 100644 index 00000000..d91991dd --- /dev/null +++ b/services/cubejs/src/__tests__/partitionPruning.test.js @@ -0,0 +1,117 @@ +import assert from "node:assert/strict"; +import { describe, it } from "node:test"; + +import { applyPartitionPruning } from "../utils/queryRewrite.js"; + +const policy = { semantic_events: { dimension: "timestamp", marginDays: 31 } }; +const map = new Map([ + [ + "simple_stays", + { + sourceTable: "semantic_events", + dimensions: new Set(["timestamp", "started", "poi"]), + partitionDimension: null, + }, + ], + [ + "bookings", + { sourceTable: "bookings", dimensions: new Set(["created"]), partitionDimension: null }, + ], + [ + "no_ts", + { sourceTable: "semantic_events", dimensions: new Set(["started"]), partitionDimension: null }, + ], +]); + +describe("partition pruning rewrite", () => { + it("adds a widened timestamp window from an explicit timeDimensions range", () => { + const q = applyPartitionPruning( + { + measures: ["simple_stays.poi_stay_count"], + dimensions: ["simple_stays.poi"], + timeDimensions: [ + { dimension: "simple_stays.started", dateRange: ["2026-06-01", "2026-06-07"] }, + ], + }, + map, + policy, + ); + assert.deepEqual(q.filters, [ + { + member: "simple_stays.timestamp", + operator: "inDateRange", + values: ["2026-05-01", "2026-07-08"], + }, + ]); + }); + + it("derives the window from inDateRange filters too, taking the union", () => { + const q = applyPartitionPruning( + { + measures: ["simple_stays.count"], + filters: [ + { member: "simple_stays.started", operator: "inDateRange", values: ["2026-01-10", "2026-01-20"] }, + { member: "simple_stays.started", operator: "inDateRange", values: ["2026-03-01", "2026-03-05"] }, + ], + }, + map, + policy, + ); + const added = q.filters.at(-1); + assert.equal(added.member, "simple_stays.timestamp"); + assert.deepEqual(added.values, ["2025-12-10", "2026-04-05"]); + }); + + it("leaves queries alone when the partition column is already constrained", () => { + const q = { + measures: ["simple_stays.count"], + timeDimensions: [ + { dimension: "simple_stays.timestamp", dateRange: ["2026-06-01", "2026-06-07"] }, + ], + }; + applyPartitionPruning(q, map, policy); + assert.equal(q.filters, undefined); + }); + + it("skips relative ranges, ungoverned tables, and cubes without the dimension", () => { + for (const query of [ + { + measures: ["simple_stays.count"], + timeDimensions: [{ dimension: "simple_stays.started", dateRange: "last 7 days" }], + }, + { + measures: ["bookings.count"], + timeDimensions: [{ dimension: "bookings.created", dateRange: ["2026-06-01", "2026-06-07"] }], + }, + { + measures: ["no_ts.count"], + timeDimensions: [{ dimension: "no_ts.started", dateRange: ["2026-06-01", "2026-06-07"] }], + }, + ]) { + applyPartitionPruning(query, map, policy); + assert.equal(query.filters, undefined); + } + }); + + it("honours meta.partition_dimension per cube", () => { + const custom = new Map([ + [ + "events_v2", + { + sourceTable: "semantic_events", + dimensions: new Set(["event_time", "occurred"]), + partitionDimension: "event_time", + }, + ], + ]); + const q = applyPartitionPruning( + { + measures: ["events_v2.count"], + timeDimensions: [{ dimension: "events_v2.occurred", dateRange: ["2026-02-01", "2026-02-02"] }], + }, + custom, + policy, + ); + assert.equal(q.filters[0].member, "events_v2.event_time"); + }); +}); diff --git a/services/cubejs/src/utils/queryRewrite.js b/services/cubejs/src/utils/queryRewrite.js index 34c0beb1..84e42e4a 100644 --- a/services/cubejs/src/utils/queryRewrite.js +++ b/services/cubejs/src/utils/queryRewrite.js @@ -11,6 +11,133 @@ const getColumnsArray = (cube) => [ ...(cube?.segments || []), ]; +// --- Partition pruning policy (owner rule, 2026-09-03) --- +// A cube over a partitioned event table must never reach the warehouse +// without a predicate on the partition time column. `cst.semantic_events` +// is partitioned by (partition, toStartOfMonth(timestamp)); a row-type cube +// whose time dimension is a JSON payload path (e.g. `properties.started_at`) +// makes ClickHouse read the JSON column for every row of every month — +// measured 10 GiB / 18.6M rows for a one-week POI query, 13 MiB / 119 ms +// once the timestamp window is present. The rewrite derives that window from +// the query's explicit date ranges, widened by `marginDays` on both sides so +// events recorded later than the payload time (a stay ends days after it +// starts) are never excluded. Policy per source table; override with +// CUBEJS_PARTITION_PRUNING (JSON: { table: { dimension, marginDays } }), and +// per cube with `meta.partition_dimension`. +const DEFAULT_PARTITION_PRUNING = { + semantic_events: { dimension: "timestamp", marginDays: 31 }, +}; + +function loadPartitionPruningPolicy() { + const raw = process.env.CUBEJS_PARTITION_PRUNING; + if (!raw) return DEFAULT_PARTITION_PRUNING; + try { + const parsed = JSON.parse(raw); + return parsed && typeof parsed === "object" + ? { ...DEFAULT_PARTITION_PRUNING, ...parsed } + : DEFAULT_PARTITION_PRUNING; + } catch (err) { + console.error( + "[queryRewrite] CUBEJS_PARTITION_PRUNING is not valid JSON:", + err.message, + ); + return DEFAULT_PARTITION_PRUNING; + } +} + +const PARTITION_PRUNING = loadPartitionPruningPolicy(); + +const DAY_MS = 24 * 60 * 60 * 1000; + +function parseExplicitDate(value) { + if (typeof value !== "string") return null; + // Cube accepts "YYYY-MM-DD" and ISO datetimes; relative expressions + // ("last 7 days") are left to Cube and yield no pruning window. + if (!/^\d{4}-\d{2}-\d{2}/.test(value)) return null; + const ms = Date.parse(value.length === 10 ? `${value}T00:00:00Z` : value); + return Number.isFinite(ms) ? ms : null; +} + +function formatDay(ms) { + return new Date(ms).toISOString().slice(0, 10); +} + +/** + * Explicit [from, to] windows the query states for members of `cubeName`, + * from timeDimensions.dateRange and inDateRange filters. + */ +function explicitDateWindows(query, cubeName) { + const windows = []; + const consider = (member, range) => { + if (typeof member !== "string" || !member.startsWith(`${cubeName}.`)) { + return; + } + const pair = Array.isArray(range) ? range : [range, range]; + if (pair.length !== 2) return; + const from = parseExplicitDate(pair[0]); + const to = parseExplicitDate(pair[1]); + if (from === null || to === null) return; + windows.push([from, to]); + }; + for (const td of query.timeDimensions || []) { + if (td?.dateRange) consider(td.dimension, td.dateRange); + } + for (const f of query.filters || []) { + if (f?.operator === "inDateRange") consider(f.member, f.values); + } + return windows; +} + +function mentionsMember(query, member) { + return ( + (query.timeDimensions || []).some((td) => td?.dimension === member) || + (query.filters || []).some((f) => f?.member === member) || + (query.dimensions || []).includes(member) + ); +} + +/** + * Add the partition-column window for every query cube whose source table + * is governed by the pruning policy. Mutates and returns `query`. + * Exported for tests. + */ +export function applyPartitionPruning( + query, + cubeToTable, + policy = PARTITION_PRUNING, +) { + if (!query || !cubeToTable || cubeToTable.size === 0) return query; + if (process.env.CUBEJS_PARTITION_PRUNING_DEBUG === "1") { + console.log("[queryRewrite] pruning map", { + cubes: [...cubeToTable.entries()].map(([k, v]) => `${k}→${v.sourceTable}`), + query: extractCubeNames(query), + }); + } + for (const cubeName of extractCubeNames(query)) { + const info = cubeToTable.get(cubeName); + if (!info?.sourceTable) continue; + const rule = policy[info.sourceTable]; + if (!rule) continue; + const dimension = info.partitionDimension || rule.dimension; + if (!dimension || !info.dimensions?.has(dimension)) continue; + const member = `${cubeName}.${dimension}`; + // Already constrained on the partition column — nothing to add. + if (mentionsMember(query, member)) continue; + const windows = explicitDateWindows(query, cubeName); + if (windows.length === 0) continue; + const margin = (Number(rule.marginDays) || 0) * DAY_MS; + const from = Math.min(...windows.map((w) => w[0])) - margin; + const to = Math.max(...windows.map((w) => w[1])) + margin; + if (!query.filters) query.filters = []; + query.filters.push({ + member, + operator: "inDateRange", + values: [formatDay(from), formatDay(to)], + }); + } + return query; +} + // --- Rule cache with 60-second TTL --- let rulesCache = null; let rulesCacheTime = 0; @@ -136,7 +263,11 @@ async function buildCubeToTableMap(schemaVersion, fileIds) { const dims = new Set( (cube.dimensions || []).map((d) => d.name).filter(Boolean), ); - mapping.set(cube.name, { sourceTable, dimensions: dims }); + mapping.set(cube.name, { + sourceTable, + dimensions: dims, + partitionDimension: cube.meta?.partition_dimension || null, + }); } } } catch (err) { @@ -214,16 +345,17 @@ const queryRewrite = async (query, { securityContext }) => { } } + // --- Step 0b: Partition pruning (applies to ALL roles, before any rule) --- + const { schemaVersion, files } = userScope.dataSource; + const cubeToTable = await buildCubeToTableMap(schemaVersion, files); + applyPartitionPruning(query, cubeToTable); + // --- Step 1: Rule-based row filtering (applies to ALL roles) --- const rules = await loadRules(); if (rules.length > 0) { const queryCubeNames = extractCubeNames(query); - // Build cube → source table mapping from active schemas - const { schemaVersion, files } = userScope.dataSource; - const cubeToTable = await buildCubeToTableMap(schemaVersion, files); - // Index rules by table name for fast lookup const rulesByTable = new Map(); for (const rule of rules) {