@fabricorg/databricks-testkit 0.1.0 → 0.2.0

package/dist/index.cjs

@@ -500,7 +500,7 @@ function looselyEqual(a, b) {
 
 // src/mock-client.ts
 function createMockDatabricksClient() {
-  const queues = { GET: [], POST: [] };
+  const queues = { GET: [], POST: [], DELETE: [] };
   const calls = [];
   function take(method, path) {
     const index = queues[method].findIndex((q) => path.startsWith(q.pathPrefix));
@@ -524,6 +524,12 @@ function createMockDatabricksClient() {
       if (queued.error) throw queued.error;
       return queued.response;
     },
+    async request(method, path, body, query) {
+      calls.push({ method, path, body, query });
+      const queued = take(method, path);
+      if (queued.error) throw queued.error;
+      return queued.response;
+    },
     queue(method, pathPrefix, response) {
       queues[method].push({ pathPrefix, response });
       return client;
@@ -692,6 +698,26 @@ function sqlRoundTripCheck() {
     }
   };
 }
+function restrictedPrincipalCheck() {
+  return {
+    id: "restricted-principal",
+    description: "Live checks use the designated least-privilege service principal",
+    configured: (env) => Boolean(env.DBX_TEST_RESTRICTED_PRINCIPAL_ID),
+    async run({ env }) {
+      const expected = env.DBX_TEST_RESTRICTED_PRINCIPAL_ID?.toLowerCase();
+      const actual = env.DATABRICKS_CLIENT_ID?.toLowerCase();
+      if (!expected || !actual) {
+        throw new Error(
+          "Restricted-principal check requires DBX_TEST_RESTRICTED_PRINCIPAL_ID and DATABRICKS_CLIENT_ID"
+        );
+      }
+      if (actual !== expected) {
+        throw new Error(`Live suite is using principal '${actual}', expected '${expected}'`);
+      }
+      return { principalId: actual };
+    }
+  };
+}
 function assertDeltaContract(table, rows, columns) {
   for (const expected of columns) {
     const actual = rows.find((row) => String(row.col_name ?? row.column_name) === expected.name);
@@ -783,17 +809,21 @@ function notebookSubmitCheck() {
   return {
     id: "notebook",
     description: "Databricks notebook submit run reaches SUCCESS",
-    configured: (env) => Boolean(env.DBX_TEST_NOTEBOOK_PATH && env.DBX_TEST_CLUSTER_ID),
+    configured: (env) => Boolean(
+      env.DBX_TEST_NOTEBOOK_PATH && (env.DBX_TEST_CLUSTER_ID || env.DBX_TEST_NOTEBOOK_SERVERLESS === "1")
+    ),
     async run({ env, client }) {
-      if (!env.DBX_TEST_NOTEBOOK_PATH || !env.DBX_TEST_CLUSTER_ID) {
-        throw new Error("Notebook submit requires DBX_TEST_NOTEBOOK_PATH and DBX_TEST_CLUSTER_ID");
+      if (!env.DBX_TEST_NOTEBOOK_PATH || !env.DBX_TEST_CLUSTER_ID && env.DBX_TEST_NOTEBOOK_SERVERLESS !== "1") {
+        throw new Error(
+          "Notebook submit requires DBX_TEST_NOTEBOOK_PATH and either DBX_TEST_CLUSTER_ID or DBX_TEST_NOTEBOOK_SERVERLESS=1"
+        );
       }
       const submitted = await client.post("/api/2.1/jobs/runs/submit", {
         run_name: `fabric-experiments-notebook-${Date.now()}`,
         tasks: [
           {
             task_key: "notebook",
-            existing_cluster_id: env.DBX_TEST_CLUSTER_ID,
+            ...env.DBX_TEST_CLUSTER_ID ? { existing_cluster_id: env.DBX_TEST_CLUSTER_ID } : {},
             notebook_task: { notebook_path: env.DBX_TEST_NOTEBOOK_PATH }
           }
         ]
@@ -813,29 +843,59 @@ function dbtBuildCheck() {
   return {
     id: "dbt",
     description: "Databricks dbt task reaches SUCCESS",
-    configured: (env) => Boolean(env.DBX_TEST_DBT_GIT_URL && env.DATABRICKS_WAREHOUSE_ID),
+    configured: (env) => Boolean(
+      (env.DBX_TEST_DBT_GIT_URL || env.DBX_TEST_DBT_PROJECT_DIR) && env.DATABRICKS_WAREHOUSE_ID
+    ),
     async run({ env, client }) {
-      if (!env.DBX_TEST_DBT_GIT_URL || !env.DATABRICKS_WAREHOUSE_ID) {
-        throw new Error("dbt build requires DBX_TEST_DBT_GIT_URL and DATABRICKS_WAREHOUSE_ID");
+      if (!env.DBX_TEST_DBT_GIT_URL && !env.DBX_TEST_DBT_PROJECT_DIR || !env.DATABRICKS_WAREHOUSE_ID) {
+        throw new Error(
+          "dbt build requires DBX_TEST_DBT_GIT_URL or DBX_TEST_DBT_PROJECT_DIR, plus DATABRICKS_WAREHOUSE_ID"
+        );
       }
-      const submitted = await client.post("/api/2.1/jobs/runs/submit", {
+      const sourceVars = JSON.stringify({
+        source_catalog: env.DBX_TEST_DBT_SOURCE_CATALOG ?? env.DBX_TEST_CATALOG,
+        source_schema: env.DBX_TEST_DBT_SOURCE_SCHEMA
+      });
+      const serverless = env.DBX_TEST_DBT_SERVERLESS === "1";
+      const request = {
         run_name: `fabric-experiments-dbt-${Date.now()}`,
-        git_source: {
-          git_url: env.DBX_TEST_DBT_GIT_URL,
-          git_provider: env.DBX_TEST_DBT_GIT_PROVIDER ?? "gitHub",
-          git_branch: env.DBX_TEST_DBT_GIT_BRANCH ?? "main"
-        },
+        ...env.DBX_TEST_DBT_GIT_URL ? {
+          git_source: {
+            git_url: env.DBX_TEST_DBT_GIT_URL,
+            git_provider: env.DBX_TEST_DBT_GIT_PROVIDER ?? "gitHub",
+            git_branch: env.DBX_TEST_DBT_GIT_BRANCH ?? "main"
+          }
+        } : {},
         tasks: [
           {
             task_key: "dbt_build",
+            ...serverless ? { environment_key: "dbt" } : {},
             dbt_task: {
-              commands: ["dbt deps", "dbt build"],
+              commands: ["dbt deps", `dbt build --vars '${sourceVars}'`],
               warehouse_id: env.DATABRICKS_WAREHOUSE_ID,
+              ...env.DBX_TEST_CATALOG ? { catalog: env.DBX_TEST_CATALOG } : {},
+              ...env.DBX_TEST_SCHEMA ? { schema: env.DBX_TEST_SCHEMA } : {},
+              ...env.DBX_TEST_DBT_GIT_URL ? {} : { source: "WORKSPACE" },
               ...env.DBX_TEST_DBT_PROJECT_DIR ? { project_directory: env.DBX_TEST_DBT_PROJECT_DIR } : {}
             }
           }
-        ]
-      });
+        ],
+        ...serverless ? {
+          environments: [
+            {
+              environment_key: "dbt",
+              spec: {
+                environment_version: "2",
+                dependencies: ["dbt-databricks>=1.8,<2"]
+              }
+            }
+          ]
+        } : {}
+      };
+      const submitted = await client.post(
+        "/api/2.1/jobs/runs/submit",
+        request
+      );
       if (!submitted.run_id) throw new Error("dbt submit returned no run_id");
       const run = await databricks.waitForDatabricksRun(client, submitted.run_id, {
         timeoutMs: Number(env.DBX_TEST_TIMEOUT_MS ?? 9e5),
@@ -843,7 +903,10 @@ function dbtBuildCheck() {
       });
       const result = run.state?.result_state;
       if (result !== "SUCCESS") throw new Error(`dbt run ended in ${result ?? "UNKNOWN"}`);
-      return { runId: submitted.run_id, gitUrl: env.DBX_TEST_DBT_GIT_URL };
+      return {
+        runId: submitted.run_id,
+        source: env.DBX_TEST_DBT_GIT_URL ?? env.DBX_TEST_DBT_PROJECT_DIR
+      };
     }
   };
 }
@@ -922,6 +985,99 @@ function lakebaseRoundTripCheck(options = {}) {
   };
 }
 
+// src/lakebase-branch.ts
+init_workspace_context();
+var DEFAULT_POLL_INTERVAL_MS = 2e3;
+var DEFAULT_TIMEOUT_MS = 5 * 6e4;
+function resolveLakebaseProject(env, explicit) {
+  const project = explicit ?? env.DBX_TEST_LAKEBASE_PROJECT;
+  if (!project || !/^projects\/[a-z0-9][a-z0-9-]{0,62}$/.test(project)) {
+    throw new Error(
+      `Lakebase project must look like projects/fabric-experiments; got '${project ?? ""}'`
+    );
+  }
+  return project;
+}
+function toLakebaseBranchId(environmentName) {
+  const normalized = environmentName.toLowerCase().replace(/[^a-z0-9-]+/g, "-").replace(/^-+|-+$/g, "");
+  const prefixed = normalized.startsWith("fx-") ? normalized : `fx-${normalized}`;
+  const trimmed = prefixed.slice(0, 63).replace(/-+$/g, "");
+  if (!/^fx-[a-z0-9](?:[a-z0-9-]*[a-z0-9])?$/.test(trimmed)) {
+    throw new Error(`Cannot derive a safe Lakebase branch ID from '${environmentName}'`);
+  }
+  return trimmed;
+}
+async function createEphemeralLakebaseBranch(options) {
+  const env = options.env ?? process.env;
+  const client = options.client ?? createClientFromEnv(env);
+  const project = resolveLakebaseProject(env, options.project);
+  const branchId = toLakebaseBranchId(options.branchId);
+  const ttlSeconds = Math.max(60, Math.ceil(options.ttlMs / 1e3));
+  const operation = await client.post(
+    `/api/2.0/postgres/${project}/branches`,
+    { spec: { ttl: `${ttlSeconds}s` } },
+    { branch_id: branchId, replace_existing: false }
+  );
+  const branch = await waitForOperation(client, operation, options);
+  const branchName = branch?.name ?? `${project}/branches/${branchId}`;
+  const [endpointResult, databaseResult] = await Promise.all([
+    client.get(`/api/2.0/postgres/${branchName}/endpoints`),
+    client.get(`/api/2.0/postgres/${branchName}/databases`)
+  ]);
+  const endpoint = endpointResult.endpoints?.find((item) => item.endpoint_id === "primary") ?? endpointResult.endpoints?.[0];
+  const database = databaseResult.databases?.[0];
+  const host = endpoint?.status?.hosts?.host;
+  const postgresDatabase = database?.status?.postgres_database;
+  if (!endpoint?.name || !host || !database?.name || !postgresDatabase) {
+    throw new Error(`Lakebase branch '${branchName}' did not expose an endpoint and database`);
+  }
+  const createdAt = branch?.create_time ?? (options.now?.() ?? /* @__PURE__ */ new Date()).toISOString();
+  return {
+    project,
+    branch: branchName,
+    branchId,
+    endpoint: endpoint.name,
+    host,
+    databaseResource: database.name,
+    database: postgresDatabase,
+    createdAt,
+    expiresAt: new Date(Date.parse(createdAt) + ttlSeconds * 1e3).toISOString()
+  };
+}
+async function deleteEphemeralLakebaseBranch(branch, options = {}) {
+  if (!/^projects\/[a-z0-9-]+\/branches\/fx-[a-z0-9-]+$/.test(branch.branch)) {
+    throw new Error(`Refusing to delete non-ephemeral Lakebase branch '${branch.branch}'`);
+  }
+  const env = options.env ?? process.env;
+  const client = options.client ?? createClientFromEnv(env);
+  const operation = await client.request(
+    "DELETE",
+    `/api/2.0/postgres/${branch.branch}`,
+    void 0,
+    { purge: true }
+  );
+  await waitForOperation(client, operation, options);
+}
+async function waitForOperation(client, initial, options) {
+  let operation = initial;
+  const deadline = Date.now() + (options.timeoutMs ?? DEFAULT_TIMEOUT_MS);
+  while (!operation.done) {
+    if (!operation.name) throw new Error("Lakebase operation did not return an operation name");
+    if (Date.now() >= deadline) throw new Error(`Lakebase operation '${operation.name}' timed out`);
+    await delay(options.pollIntervalMs ?? DEFAULT_POLL_INTERVAL_MS);
+    operation = await client.get(`/api/2.0/postgres/${operation.name}`);
+  }
+  if (operation.error) {
+    throw new Error(
+      `Lakebase operation failed${operation.error.code ? ` (${operation.error.code})` : ""}: ${operation.error.message ?? "unknown error"}`
+    );
+  }
+  return operation.response;
+}
+function delay(ms) {
+  return new Promise((resolve2) => setTimeout(resolve2, ms));
+}
+
 // src/databricks-http.ts
 init_workspace_context();
 async function databricksFetch(env, pathOrUrl, init = {}, fetchImpl = fetch) {
@@ -1115,6 +1271,7 @@ exports.autoLoaderIngestionCheck = autoLoaderIngestionCheck;
 exports.compareToGolden = compareToGolden;
 exports.createClientFromEnv = createClientFromEnv;
 exports.createDuckDbContext = createDuckDbContext;
+exports.createEphemeralLakebaseBranch = createEphemeralLakebaseBranch;
 exports.createExecutionContext = createExecutionContext;
 exports.createLakebaseTestProvider = createLakebaseTestProvider;
 exports.createMockDatabricksClient = createMockDatabricksClient;
@@ -1122,6 +1279,7 @@ exports.createWorkspaceContext = createWorkspaceContext;
 exports.customLiveCheck = customLiveCheck;
 exports.dbtBuildCheck = dbtBuildCheck;
 exports.defineLiveSuite = defineLiveSuite;
+exports.deleteEphemeralLakebaseBranch = deleteEphemeralLakebaseBranch;
 exports.deleteVolumeFile = deleteVolumeFile;
 exports.deltaContractCheck = deltaContractCheck;
 exports.ensureVolumeDirectory = ensureVolumeDirectory;
@@ -1138,11 +1296,14 @@ exports.parseStatementRows = parseStatementRows;
 exports.pipelineRefreshCheck = pipelineRefreshCheck;
 exports.renderJUnit = renderJUnit;
 exports.resolveFixtureRows = resolveFixtureRows;
+exports.resolveLakebaseProject = resolveLakebaseProject;
 exports.resolveProfile = resolveProfile;
 exports.resolveTokenProvider = resolveTokenProvider;
+exports.restrictedPrincipalCheck = restrictedPrincipalCheck;
 exports.runLiveSuite = runLiveSuite;
 exports.secretScopeCheck = secretScopeCheck;
 exports.sqlRoundTripCheck = sqlRoundTripCheck;
+exports.toLakebaseBranchId = toLakebaseBranchId;
 exports.toNamedParameters = toNamedParameters;
 exports.unityCatalogAccessCheck = unityCatalogAccessCheck;
 exports.uploadVolumeFile = uploadVolumeFile;