Add an end-to-end test for the asset bundles demo

.github/workflows/acceptance.yml

@@ -127,6 +127,24 @@ jobs:
       - name: Install hatch
         run: pip install hatch==1.9.4
 
+      - name: Install Databricks CLI
+        run: |
+          curl -fsSL https://raw.githubusercontent.com/databricks/setup-cli/main/install.sh | sh
+          databricks --version
+
+      - name: Azure login (OIDC)
+        uses: azure/login@v2
+        with:
+          client-id: ${{ secrets.ARM_CLIENT_ID }}
+          tenant-id: ${{ secrets.ARM_TENANT_ID }}
+          allow-no-subscriptions: true
+
+      - name: Set env vars for Azure CLI auth
+        run: |
+          val=$(az keyvault secret show --id "${{ secrets.VAULT_URI }}/secrets/DATABRICKS-HOST" --query value -o tsv)
+          echo "DATABRICKS_HOST=$val" >> $GITHUB_ENV
+          echo "DATABRICKS_AUTH_TYPE=azure-cli" >> $GITHUB_ENV
+
       - name: Run e2e tests
         uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.4
         with:
@@ -161,6 +179,24 @@ jobs:
       - name: Install hatch
         run: pip install hatch==1.9.4
 
+      - name: Install Databricks CLI
+        run: |
+          curl -fsSL https://raw.githubusercontent.com/databricks/setup-cli/main/install.sh | sh
+          databricks --version
+
+      - name: Azure login (OIDC)
+        uses: azure/login@v2
+        with:
+          client-id: ${{ secrets.ARM_CLIENT_ID }}
+          tenant-id: ${{ secrets.ARM_TENANT_ID }}
+          allow-no-subscriptions: true
+
+      - name: Set env vars for Azure CLI auth
+        run: |
+          val=$(az keyvault secret show --id "${{ secrets.VAULT_URI }}/secrets/DATABRICKS-HOST" --query value -o tsv)
+          echo "DATABRICKS_HOST=$val" >> $GITHUB_ENV
+          echo "DATABRICKS_AUTH_TYPE=azure-cli" >> $GITHUB_ENV
+
       - name: Run e2e tests on serverless cluster
         uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.4
         with:

.github/workflows/nightly.yml

@@ -116,6 +116,24 @@ jobs:
       - name: Install hatch
         run: pip install hatch==1.9.4
 
+      - name: Install Databricks CLI
+        run: |
+          curl -fsSL https://raw.githubusercontent.com/databricks/setup-cli/main/install.sh | sh
+          databricks --version
+
+      - name: Azure login (OIDC)
+        uses: azure/login@v2
+        with:
+          client-id: ${{ secrets.ARM_CLIENT_ID }}
+          tenant-id: ${{ secrets.ARM_TENANT_ID }}
+          allow-no-subscriptions: true
+
+      - name: Set env vars for Azure CLI auth
+        run: |
+          val=$(az keyvault secret show --id "${{ secrets.VAULT_URI }}/secrets/DATABRICKS-HOST" --query value -o tsv)
+          echo "DATABRICKS_HOST=$val" >> $GITHUB_ENV
+          echo "DATABRICKS_AUTH_TYPE=azure-cli" >> $GITHUB_ENV
+
       - name: Run e2e tests
         uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.4
         with:
@@ -150,6 +168,24 @@ jobs:
       - name: Install hatch
         run: pip install hatch==1.9.4
 
+      - name: Install Databricks CLI
+        run: |
+          curl -fsSL https://raw.githubusercontent.com/databricks/setup-cli/main/install.sh | sh
+          databricks --version
+
+      - name: Azure login (OIDC)
+        uses: azure/login@v2
+        with:
+          client-id: ${{ secrets.ARM_CLIENT_ID }}
+          tenant-id: ${{ secrets.ARM_TENANT_ID }}
+          allow-no-subscriptions: true
+
+      - name: Set env vars for Azure CLI auth
+        run: |
+          val=$(az keyvault secret show --id "${{ secrets.VAULT_URI }}/secrets/DATABRICKS-HOST" --query value -o tsv)
+          echo "DATABRICKS_HOST=$val" >> $GITHUB_ENV
+          echo "DATABRICKS_AUTH_TYPE=azure-cli" >> $GITHUB_ENV
+
       - name: Run e2e tests on serverless cluster
         uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.4
         with:

demos/dqx_demo_asset_bundle/databricks.yml

@@ -14,13 +14,20 @@ variables:  # Sets the default value of several notebook parameters
   sensor_rules_file:
     description: Name of the DQX rules file for the sensor dataset
     default: sensor_data_quality_rules.yml
+  library_ref:
+    description: PyPi package name for installing DQX
+    default: databricks-labs-dqx
 
 resources:
   jobs:
     dqx_demo_job:
       name: "[${workspace.current_user.userName}] DQX Demo Job"
       tasks:
         - task_key: dqx_demo_notebook
+          job_cluster_key: dqx_demo_cluster
+          libraries:
+            - pypi:
+                package: "${var.library_ref}"
           notebook_task:
             notebook_path: ./dqx_demo_notebook.py
             base_parameters:  # Parameters are passed to the notebook; Default values are set in the `variables` section
@@ -29,12 +36,13 @@ resources:
               maintenance_rules_file: "${var.maintenance_rules_file}"
               sensor_rules_file: "${var.sensor_rules_file}"
             source: WORKSPACE
-          environment_key: Default
+      job_clusters:
+        - job_cluster_key: dqx_demo_cluster
+          new_cluster:
+            num_workers: 1
+            node_type_id: Standard_DS3_v2
+            spark_version: 15.4.x-scala2.12
+            data_security_mode: SINGLE_USER
+
       queue:
-        enabled: true
-      environments:
-        - environment_key: Default
-          spec:
-            client: "2"
-            dependencies:
-              - databricks-labs-dqx
+        enabled: true

demos/dqx_demo_asset_bundle/dqx_demo_notebook.py

@@ -29,11 +29,6 @@
 print(f"Selected Catalog for Demo Dataset: {database}")
 print(f"Selected Schema for Demo Dataset: {schema}")
 
-spark.sql(f"CREATE CATALOG IF NOT EXISTS {database}")
-spark.sql(f"USE CATALOG {database}")
-spark.sql(f"CREATE SCHEMA IF NOT EXISTS {schema}")
-spark.sql(f"USE SCHEMA {schema}")
-
 sensor_table = f"{database}.{schema}.sensor_data"
 maintenance_table = f"{database}.{schema}.maintenance_data"
 
@@ -228,7 +223,7 @@
             "SEN-001",
             "MCH-001",
             "temperature",
-            735,
+            735.0,
             datetime.strptime("2025-04-29 14:32:00", "%Y-%m-%d %H:%M:%S"),
             datetime.strptime("2025-04-01", "%Y-%m-%d").date(),
             80,
@@ -282,7 +277,7 @@
             "SEN-002",
             "MCH-001",
             "temperature",
-            735,
+            735.0,
             datetime.strptime("2026-04-29 14:32:00", "%Y-%m-%d %H:%M:%S"),
             datetime.strptime("2025-04-01", "%Y-%m-%d").date(),
             80,
@@ -337,7 +332,7 @@
             "SEN004",
             "MCH-001",
             "temperature",
-            724,
+            724.0,
             datetime.strptime("2025-04-28 14:32:00", "%Y-%m-%d %H:%M:%S"),
             datetime.strptime("2025-04-30", "%Y-%m-%d").date(),
             85,
@@ -351,7 +346,7 @@
             "SEN004",
             "MCH-001",
             "temperature",
-            724,
+            724.0,
             datetime.strptime("2025-04-28 14:32:00", "%Y-%m-%d %H:%M:%S"),
             datetime.strptime("2025-04-30", "%Y-%m-%d").date(),
             85,

tests/e2e/test_run_demos.py

@@ -1,4 +1,6 @@
 import logging
+import shutil
+import subprocess
 
 from datetime import timedelta
 from pathlib import Path
@@ -229,6 +231,33 @@ def test_run_dqx_streaming_demo_diy(make_notebook, make_job, tmp_path, library_r
     logging.info(f"Job run {run.run_id} completed successfully for dqx_streaming_demo")
 
 
+def test_run_dqx_demo_asset_bundle(make_schema, library_ref):
+    cli_path = shutil.which("databricks")
+    path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_asset_bundle"
+    catalog = "main"
+    schema = make_schema(catalog_name=catalog).name
+
+    try:
+        subprocess.run([cli_path, "bundle", "validate"], check=True, capture_output=True, cwd=path)
+        subprocess.run(
+            [
+                cli_path,
+                "bundle",
+                "deploy",
+                f'--var="library_ref={library_ref}"',
+                f'--var="demo_catalog={catalog}"',
+                f'--var="demo_schema={schema}"',
+                "--auto-approve",
+            ],
+            check=True,
+            capture_output=True,
+            cwd=path,
+        )
+        subprocess.run([cli_path, "bundle", "run", "dqx_demo_job"], check=True, capture_output=True, cwd=path)
+    except subprocess.CalledProcessError as ex:
+        raise AssertionError(ex.stderr) from ex
+
+
 def validate_run_status(run: Run, client: WorkspaceClient) -> None:
     """
     Validates that a job task run completed successfully.