perf(cubejs): scheduler reads model code once per version, not every interval

services/cubejs/src/utils/__tests__/scheduledRefreshQuarantineCache.test.js

@@ -0,0 +1,41 @@
+import assert from "node:assert/strict";
+import { describe, it, mock } from "node:test";
+
+const UNSAFE = `cubes:
+  - name: events
+    sql: >
+      SELECT * FROM t WHERE {FILTER_PARAMS.events.partition.filter((c) => c)}
+`;
+const SAFE = `cubes:
+  - name: orders
+    sql: SELECT * FROM orders
+`;
+
+const dataSource = (id, versionId) => ({
+  id,
+  db_type: "clickhouse",
+  db_params: { host: "database.invalid" },
+  branches: [{ versions: [{ id: versionId, dataschemas: [{ id: `${versionId}-m`, name: "m.yml" }] }] }],
+});
+
+describe("scheduled refresh quarantine cache", () => {
+  it("reads model code once per version, not every interval", async () => {
+    let dataSources = [dataSource("ds-a", "v1"), dataSource("ds-b", "v2")];
+    const code = { v1: UNSAFE, v2: SAFE, v3: SAFE };
+    const getVersionModelCode = mock.fn(async (versionId) => [{ code: code[versionId] }]);
+    mock.module("../dataSourceHelpers.js", {
+      namedExports: { getDataSources: async () => dataSources, getVersionModelCode },
+    });
+    const { default: scheduledRefreshContexts } = await import("../scheduledRefreshContexts.js");
+    const ids = (contexts) => contexts.map((c) => c.securityContext.userScope.dataSource.dataSourceId);
+
+    assert.deepEqual(ids(await scheduledRefreshContexts()), ["ds-b"]);
+    assert.deepEqual(ids(await scheduledRefreshContexts()), ["ds-b"]);
+    assert.equal(getVersionModelCode.mock.callCount(), 2);
+
+    // ds-a publishes a fixed version: only the new version is read.
+    dataSources = [dataSource("ds-a", "v3"), dataSource("ds-b", "v2")];
+    assert.deepEqual(ids(await scheduledRefreshContexts()), ["ds-a", "ds-b"]);
+    assert.deepEqual(getVersionModelCode.mock.calls.map((c) => c.arguments[0]), ["v1", "v2", "v3"]);
+  });
+});

services/cubejs/src/utils/__tests__/scheduledRefreshStall.test.js

@@ -6,7 +6,10 @@ describe("scheduled refresh stall detection", () => {
     mock.timers.enable({ apis: ["Date"], now: 0 });
     // The interval's datasource lookup hangs, as it did on dev.
     mock.module("../dataSourceHelpers.js", {
-      namedExports: { getDataSources: () => new Promise(() => {}) },
+      namedExports: {
+        getDataSources: () => new Promise(() => {}),
+        getVersionModelCode: async () => [],
+      },
     });
     const { default: scheduledRefreshContexts, scheduledRefreshStalledFor } =
       await import("../scheduledRefreshContexts.js");

services/cubejs/src/utils/dataSourceHelpers.js

@@ -162,11 +162,32 @@ const userQuery = `
   }
 `;
 
+// The scheduler runs this every interval for every datasource, so it carries
+// model ids only; model code is read once per version (versionModelCodeQuery).
 const sourcesQuery = `
   {
     datasources {
       ${sourceFragment}
-      ${activeBranchModelsFragment}
+      branches(where: {status: {_eq: active}}) {
+        ${branchesFragment}
+        versions (limit: 1, order_by: {created_at: desc}) {
+          id
+          dataschemas {
+            id
+            name
+          }
+        }
+      }
+    }
+  }
+`;
+
+const versionModelCodeQuery = `
+  query ($versionId: uuid!) {
+    dataschemas(where: {version_id: {_eq: $versionId}}) {
+      id
+      name
+      code
     }
   }
 `;
@@ -351,6 +372,11 @@ export const getDataSources = async () => {
   return res;
 };
 
+export const getVersionModelCode = async (versionId) => {
+  const res = await fetchGraphQL(versionModelCodeQuery, { versionId });
+  return res?.data?.dataschemas || [];
+};
+
 /**
  * Legacy `versions_insert_input`-shaped entry point (generate-models,
  * smart-generate, reconcile-team) over {@link commitVersionFiles}. Keeps its

services/cubejs/src/utils/scheduledRefreshContexts.js

@@ -1,4 +1,4 @@
-import { getDataSources } from "./dataSourceHelpers.js";
+import { getDataSources, getVersionModelCode } from "./dataSourceHelpers.js";
 import buildSecurityContext from "./buildSecurityContext.js";
 import YAML from "yaml";
 
@@ -35,8 +35,12 @@ export const hasUnsafeScheduledRefreshPartitionFilter = (schema) => {
   }
 };
 
-const activeSchemas = (dataSource) =>
-  dataSource?.branches?.[0]?.versions?.[0]?.dataschemas || [];
+const activeVersion = (dataSource) => dataSource?.branches?.[0]?.versions?.[0];
+
+const activeSchemas = (dataSource) => activeVersion(dataSource)?.dataschemas || [];
+
+const countUnsafeSchemas = (schemas) =>
+  schemas.filter(hasUnsafeScheduledRefreshPartitionFilter).length;
 
 /**
  * Build refresh contexts while quarantining only the legacy model pattern
@@ -48,6 +52,7 @@ export const buildScheduledRefreshContexts = (
   {
     buildContext = buildSecurityContext,
     disabledDbTypes = new Set(),
+    unsafeSchemaCount = (dataSource) => countUnsafeSchemas(activeSchemas(dataSource)),
     warn = (summary) =>
       console.warn("Scheduled refresh context quarantined", summary),
   } = {},
@@ -64,9 +69,7 @@ export const buildScheduledRefreshContexts = (
       continue;
     }
 
-    const affected = activeSchemas(dataSource).filter(
-      hasUnsafeScheduledRefreshPartitionFilter,
-    ).length;
+    const affected = unsafeSchemaCount(dataSource);
 
     if (affected > 0) {
       excludedDataSources += 1;
@@ -118,13 +121,35 @@ export const scheduledRefreshStalledFor = (now = Date.now()) => {
  *
  * @returns {Promise<Array>} - A promise that resolves to an array of objects, where each object contains the security context for a data source.
  */
+// Versions are immutable, so a version's quarantine verdict never changes:
+// its model code is read once, not on every interval for every datasource.
+const unsafeSchemasByVersion = new Map();
+
 const scheduledRefreshContexts = async () => {
   lastIntervalStartedAt = Date.now();
-  const dataSources = await getDataSources();
+  const dataSources = (await getDataSources()) || [];
+
+  const activeVersionIds = new Set(
+    dataSources.map((dataSource) => activeVersion(dataSource)?.id).filter(Boolean),
+  );
+  for (const versionId of unsafeSchemasByVersion.keys()) {
+    if (!activeVersionIds.has(versionId)) unsafeSchemasByVersion.delete(versionId);
+  }
+  for (const versionId of activeVersionIds) {
+    if (!unsafeSchemasByVersion.has(versionId)) {
+      unsafeSchemasByVersion.set(
+        versionId,
+        countUnsafeSchemas(await getVersionModelCode(versionId)),
+      );
+    }
+  }
+
   return buildScheduledRefreshContexts(dataSources, {
     disabledDbTypes: parseDisabledScheduledRefreshDbTypes(
       process.env.CUBEJS_SCHEDULED_REFRESH_DISABLED_DB_TYPES,
     ),
+    unsafeSchemaCount: (dataSource) =>
+      unsafeSchemasByVersion.get(activeVersion(dataSource)?.id) ?? 0,
   });
 };