From 2ff3d18b2d06e9370184334d586d6a1c204664d8 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 14 Jul 2025 09:27:45 -0400 Subject: [PATCH 01/91] Run demo notebooks in automated tests --- Makefile | 5 ++++- pyproject.toml | 1 + tests/end_to_end/__init__.py | 0 tests/end_to_end/test_run_demos.py | 35 ++++++++++++++++++++++++++++++ 4 files changed, 40 insertions(+), 1 deletion(-) create mode 100644 tests/end_to_end/__init__.py create mode 100644 tests/end_to_end/test_run_demos.py diff --git a/Makefile b/Makefile index 0af9db073..a3c66fb3b 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -all: clean dev lint fmt test integration coverage +all: clean dev lint fmt test integration coverage end-to-end clean: rm -fr .venv clean htmlcov .mypy_cache .pytest_cache .ruff_cache .coverage coverage.xml @@ -24,6 +24,9 @@ test: integration: hatch run integration +end-to-end: + hatch run end_to_end + coverage: hatch run coverage; open htmlcov/index.html diff --git a/pyproject.toml b/pyproject.toml index a90cfac2f..25aecaf65 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -84,6 +84,7 @@ path = ".venv" test = "pytest tests/unit/ -n 10 --cov --cov-report=xml:coverage-unit.xml --timeout 30 --durations 20" coverage = "pytest tests/ -n 10 --cov --cov-report=html --timeout 600 --durations 20" integration = "pytest tests/integration/ -n 10 --cov --cov-report=xml --timeout 600 --durations 20" +end_to_end = "pytest tests/end_to_end/ -n 10 --cov --cov-report=xml --timeout 600 --durations 20" fmt = ["black .", "ruff check . --fix", "mypy .", diff --git a/tests/end_to_end/__init__.py b/tests/end_to_end/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/end_to_end/test_run_demos.py b/tests/end_to_end/test_run_demos.py new file mode 100644 index 000000000..51203fff4 --- /dev/null +++ b/tests/end_to_end/test_run_demos.py @@ -0,0 +1,35 @@ +import logging +import os + +from pathlib import Path +from databricks.sdk import WorkspaceClient +from databricks.sdk.service.jobs import NotebookTask, SubmitTask, RunLifecycleStateV2State, TerminationTypeType + + +logging.getLogger("tests").setLevel("DEBUG") +logging.getLogger("databricks.labs.dqx").setLevel("DEBUG") +logger = logging.getLogger(__name__) + + +def test_run_all_demo_notebooks_succeed(make_notebook): + demo_folder = f"{Path(__file__).parent.parent.parent}/demos" + notebook_paths = [path for path in os.listdir(demo_folder) if path.endswith(".py")] + ws = WorkspaceClient() + + run_ids = [] + for notebook_path in notebook_paths: + path = f"{demo_folder}/{notebook_path}" + with open(path, "rb") as f: + notebook = make_notebook(content=f) + wait_for_run = ws.jobs.submit( + tasks=[SubmitTask(task_key="demo_run", notebook_task=NotebookTask(notebook_path=notebook.as_fuse()))] + ) + run_ids.append(wait_for_run.response.run_id) + + for run_id in run_ids: + run = ws.jobs.get_run(run_id) + if not run.status.state == RunLifecycleStateV2State.TERMINATED: + run_details = run.status.termination_details + assert ( + run_details.type == TerminationTypeType.SUCCESS + ), f"Run ended with status {run_details.type.value}: {run_details.message}" From 634c58add1f4277505de2701d955f55680a2d3c1 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 14 Jul 2025 12:16:15 -0400 Subject: [PATCH 02/91] Refactor --- Makefile | 4 ++-- pyproject.toml | 2 +- tests/{end_to_end => e2e}/__init__.py | 0 tests/{end_to_end => e2e}/test_run_demos.py | 6 +++--- 4 files changed, 6 insertions(+), 6 deletions(-) rename tests/{end_to_end => e2e}/__init__.py (100%) rename tests/{end_to_end => e2e}/test_run_demos.py (86%) diff --git a/Makefile b/Makefile index a3c66fb3b..48f62df3b 100644 --- a/Makefile +++ b/Makefile @@ -24,8 +24,8 @@ test: integration: hatch run integration -end-to-end: - hatch run end_to_end +e2e: + hatch run e2e coverage: hatch run coverage; open htmlcov/index.html diff --git a/pyproject.toml b/pyproject.toml index 25aecaf65..00e078d6e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -84,7 +84,7 @@ path = ".venv" test = "pytest tests/unit/ -n 10 --cov --cov-report=xml:coverage-unit.xml --timeout 30 --durations 20" coverage = "pytest tests/ -n 10 --cov --cov-report=html --timeout 600 --durations 20" integration = "pytest tests/integration/ -n 10 --cov --cov-report=xml --timeout 600 --durations 20" -end_to_end = "pytest tests/end_to_end/ -n 10 --cov --cov-report=xml --timeout 600 --durations 20" +e2e = "pytest tests/e2e/ -n 10 --cov --cov-report=xml --timeout 600 --durations 20" fmt = ["black .", "ruff check . --fix", "mypy .", diff --git a/tests/end_to_end/__init__.py b/tests/e2e/__init__.py similarity index 100% rename from tests/end_to_end/__init__.py rename to tests/e2e/__init__.py diff --git a/tests/end_to_end/test_run_demos.py b/tests/e2e/test_run_demos.py similarity index 86% rename from tests/end_to_end/test_run_demos.py rename to tests/e2e/test_run_demos.py index 51203fff4..65cc91a85 100644 --- a/tests/end_to_end/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -12,13 +12,13 @@ def test_run_all_demo_notebooks_succeed(make_notebook): - demo_folder = f"{Path(__file__).parent.parent.parent}/demos" + demo_folder = Path(__file__).parent.parent.parent / "demos" notebook_paths = [path for path in os.listdir(demo_folder) if path.endswith(".py")] ws = WorkspaceClient() run_ids = [] for notebook_path in notebook_paths: - path = f"{demo_folder}/{notebook_path}" + path = demo_folder / notebook_path with open(path, "rb") as f: notebook = make_notebook(content=f) wait_for_run = ws.jobs.submit( @@ -28,7 +28,7 @@ def test_run_all_demo_notebooks_succeed(make_notebook): for run_id in run_ids: run = ws.jobs.get_run(run_id) - if not run.status.state == RunLifecycleStateV2State.TERMINATED: + if run.status.state == RunLifecycleStateV2State.TERMINATED: run_details = run.status.termination_details assert ( run_details.type == TerminationTypeType.SUCCESS From ba612cdf4ac45779528f02d128592fc84f00727c Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Tue, 15 Jul 2025 10:00:51 +0200 Subject: [PATCH 03/91] added e2e tests to github actions --- .github/workflows/acceptance.yml | 35 ++++++++++++++++++++++++++++++++ tests/e2e/.codegen.json | 12 +++++++++++ 2 files changed, 47 insertions(+) create mode 100644 tests/e2e/.codegen.json diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index c346a937a..465c37512 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -19,6 +19,7 @@ concurrency: cancel-in-progress: true jobs: + integration: # Only run this job for PRs from branches on the main repository and not from forks. # Workflows triggered by PRs from forks don't have access to the tool environment. @@ -103,3 +104,37 @@ jobs: ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }} DATABRICKS_SERVERLESS_COMPUTE_ID: ${{ env.DATABRICKS_SERVERLESS_COMPUTE_ID }} + + e2e: + # Only run this job for PRs from branches on the main repository and not from forks. + # Workflows triggered by PRs from forks don't have access to the tool environment. + # PRs from forks to be tested by the reviewer(s) / maintainer(s) before merging. + if: github.event_name == 'pull_request' && !github.event.pull_request.draft && !github.event.pull_request.head.repo.fork + environment: tool + runs-on: larger + steps: + - name: Checkout Code + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Install Python + uses: actions/setup-python@v5 + with: + cache: 'pip' + cache-dependency-path: '**/pyproject.toml' + python-version: '3.10' + + - name: Install hatch + run: pip install hatch==1.9.4 + + - name: Run e2e tests + uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.3 + with: + vault_uri: ${{ secrets.VAULT_URI }} + timeout: 2h + directory: tests/e2e + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} + ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }} diff --git a/tests/e2e/.codegen.json b/tests/e2e/.codegen.json new file mode 100644 index 000000000..fd49fc564 --- /dev/null +++ b/tests/e2e/.codegen.json @@ -0,0 +1,12 @@ +{ + "version": { + "src/databricks/labs/dqx/__about__.py": "__version__ = \"$VERSION\"" + }, + "toolchain": { + "required": ["python3", "hatch"], + "pre_setup": ["hatch env create"], + "prepend_path": ".venv/bin", + "acceptance_path": "", + "test": [] + } +} \ No newline at end of file From b2328d8bd1c6cc335e3b30477c59d39a9d47aefc Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Tue, 15 Jul 2025 12:28:14 +0200 Subject: [PATCH 04/91] test --- .codegen.json | 4 +--- .github/workflows/acceptance.yml | 1 - 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/.codegen.json b/.codegen.json index 4cebef620..74b9ad344 100644 --- a/.codegen.json +++ b/.codegen.json @@ -5,8 +5,6 @@ "toolchain": { "required": ["python3", "hatch"], "pre_setup": ["hatch env create"], - "prepend_path": ".venv/bin", - "acceptance_path": "tests/integration", - "test": [] + "prepend_path": ".venv/bin" } } \ No newline at end of file diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index 465c37512..0cd22acce 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -133,7 +133,6 @@ jobs: with: vault_uri: ${{ secrets.VAULT_URI }} timeout: 2h - directory: tests/e2e env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} From 96de2d569650619b72646ead7d80822e6954c62c Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Tue, 15 Jul 2025 13:12:43 +0200 Subject: [PATCH 05/91] test --- .codegen.json | 3 ++- .github/workflows/acceptance.yml | 2 ++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/.codegen.json b/.codegen.json index 74b9ad344..bd11afde0 100644 --- a/.codegen.json +++ b/.codegen.json @@ -5,6 +5,7 @@ "toolchain": { "required": ["python3", "hatch"], "pre_setup": ["hatch env create"], - "prepend_path": ".venv/bin" + "prepend_path": ".venv/bin", + "acceptance_path": "tests/integration" } } \ No newline at end of file diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index 0cd22acce..20ea48263 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -128,11 +128,13 @@ jobs: - name: Install hatch run: pip install hatch==1.9.4 + # TODO: update after this PR is released: https://github.com/databrickslabs/sandbox/pull/502 - name: Run e2e tests uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.3 with: vault_uri: ${{ secrets.VAULT_URI }} timeout: 2h + codegen_path: tests/e2e/.codegen.json env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} From ec92f9185e4cfc18b0dc6029e17098e72650b9cb Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 08:42:42 -0400 Subject: [PATCH 06/91] Refactor demo tests --- tests/e2e/test_run_demos.py | 36 ++++++++++++++++++++---------------- 1 file changed, 20 insertions(+), 16 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 65cc91a85..659236516 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -1,8 +1,10 @@ +import asyncio import logging import os from pathlib import Path from databricks.sdk import WorkspaceClient +from databricks.sdk.service.workspace import ImportFormat from databricks.sdk.service.jobs import NotebookTask, SubmitTask, RunLifecycleStateV2State, TerminationTypeType @@ -13,23 +15,25 @@ def test_run_all_demo_notebooks_succeed(make_notebook): demo_folder = Path(__file__).parent.parent.parent / "demos" - notebook_paths = [path for path in os.listdir(demo_folder) if path.endswith(".py")] + notebook_paths = [path for path in os.listdir(demo_folder) if path.endswith((".py", ".ipynb", ".dbc"))] ws = WorkspaceClient() - run_ids = [] - for notebook_path in notebook_paths: + async def get_submit_job_result(notebook_path): path = demo_folder / notebook_path with open(path, "rb") as f: - notebook = make_notebook(content=f) - wait_for_run = ws.jobs.submit( - tasks=[SubmitTask(task_key="demo_run", notebook_task=NotebookTask(notebook_path=notebook.as_fuse()))] - ) - run_ids.append(wait_for_run.response.run_id) - - for run_id in run_ids: - run = ws.jobs.get_run(run_id) - if run.status.state == RunLifecycleStateV2State.TERMINATED: - run_details = run.status.termination_details - assert ( - run_details.type == TerminationTypeType.SUCCESS - ), f"Run ended with status {run_details.type.value}: {run_details.message}" + notebook = make_notebook(content=f, format=ImportFormat.AUTO) + return await ws.jobs.submit( + tasks=[SubmitTask(task_key="demo_run", notebook_task=NotebookTask(notebook_path=notebook.as_fuse()))] + ) + + async def validate_submit_job_runs(runs): + for completion in asyncio.as_completed(runs): + run = await completion + if run.status.state == RunLifecycleStateV2State.TERMINATED: + run_details = run.status.termination_details + assert ( + run_details.type == TerminationTypeType.SUCCESS + ), f"Run ended with status {run_details.type.value}: {run_details.message}" + + runs = [get_submit_job_result(notebook_path) for notebook_path in notebook_paths] + asyncio.run(validate_submit_job_runs(runs)) From d8448e7d367e6ef9aa06b67afbc9d49cb02ec97b Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Tue, 15 Jul 2025 15:00:19 +0200 Subject: [PATCH 07/91] test --- .github/workflows/acceptance.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index 20ea48263..9cc774898 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -130,7 +130,7 @@ jobs: # TODO: update after this PR is released: https://github.com/databrickslabs/sandbox/pull/502 - name: Run e2e tests - uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.3 + uses: databrickslabs/sandbox/acceptance@codegen with: vault_uri: ${{ secrets.VAULT_URI }} timeout: 2h From 471b6138b5880cf054a7480c417472ef38db7bf6 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Tue, 15 Jul 2025 16:18:21 +0200 Subject: [PATCH 08/91] updated acceptance tests --- .codegen.json | 3 +-- .github/workflows/acceptance.yml | 46 +++++++++++++++++++++++++++----- tests/e2e/.codegen.json | 3 +-- tests/integration/.codegen.json | 11 ++++++++ 4 files changed, 52 insertions(+), 11 deletions(-) create mode 100644 tests/integration/.codegen.json diff --git a/.codegen.json b/.codegen.json index bd11afde0..74b9ad344 100644 --- a/.codegen.json +++ b/.codegen.json @@ -5,7 +5,6 @@ "toolchain": { "required": ["python3", "hatch"], "pre_setup": ["hatch env create"], - "prepend_path": ".venv/bin", - "acceptance_path": "tests/integration" + "prepend_path": ".venv/bin" } } \ No newline at end of file diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index 9cc774898..00559770d 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -54,10 +54,11 @@ jobs: # Run tests from `tests/integration` as defined in .codegen.json # and generate code coverage for modules defined in .coveragerc - name: Run integration tests and generate test coverage report - uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.3 + uses: databrickslabs/sandbox/acceptance@acceptance/main with: vault_uri: ${{ secrets.VAULT_URI }} timeout: 2h + codegen_path: tests/integration/.codegen.json env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} @@ -95,10 +96,11 @@ jobs: run: pip install hatch==1.9.4 - name: Run integration tests on serverless cluster - uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.3 + uses: databrickslabs/sandbox/acceptance@acceptance/main with: vault_uri: ${{ secrets.VAULT_URI }} timeout: 2h + codegen_path: tests/integration/.codegen.json env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} @@ -106,9 +108,6 @@ jobs: DATABRICKS_SERVERLESS_COMPUTE_ID: ${{ env.DATABRICKS_SERVERLESS_COMPUTE_ID }} e2e: - # Only run this job for PRs from branches on the main repository and not from forks. - # Workflows triggered by PRs from forks don't have access to the tool environment. - # PRs from forks to be tested by the reviewer(s) / maintainer(s) before merging. if: github.event_name == 'pull_request' && !github.event.pull_request.draft && !github.event.pull_request.head.repo.fork environment: tool runs-on: larger @@ -128,9 +127,41 @@ jobs: - name: Install hatch run: pip install hatch==1.9.4 - # TODO: update after this PR is released: https://github.com/databrickslabs/sandbox/pull/502 - name: Run e2e tests - uses: databrickslabs/sandbox/acceptance@codegen + uses: databrickslabs/sandbox/acceptance@acceptance/main + with: + vault_uri: ${{ secrets.VAULT_URI }} + timeout: 2h + codegen_path: tests/e2e/.codegen.json + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} + ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }} + + e2e_serverless: + if: github.event_name == 'pull_request' && !github.event.pull_request.draft && !github.event.pull_request.head.repo.fork + environment: tool + runs-on: larger + env: + DATABRICKS_SERVERLESS_COMPUTE_ID: auto + steps: + - name: Checkout Code + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Install Python + uses: actions/setup-python@v5 + with: + cache: 'pip' + cache-dependency-path: '**/pyproject.toml' + python-version: '3.10' + + - name: Install hatch + run: pip install hatch==1.9.4 + + - name: Run e2e tests on serverless cluster + uses: databrickslabs/sandbox/acceptance@acceptance/main with: vault_uri: ${{ secrets.VAULT_URI }} timeout: 2h @@ -139,3 +170,4 @@ jobs: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }} + DATABRICKS_SERVERLESS_COMPUTE_ID: ${{ env.DATABRICKS_SERVERLESS_COMPUTE_ID }} diff --git a/tests/e2e/.codegen.json b/tests/e2e/.codegen.json index fd49fc564..a76875481 100644 --- a/tests/e2e/.codegen.json +++ b/tests/e2e/.codegen.json @@ -6,7 +6,6 @@ "required": ["python3", "hatch"], "pre_setup": ["hatch env create"], "prepend_path": ".venv/bin", - "acceptance_path": "", - "test": [] + "acceptance_path": "tests/e2e" } } \ No newline at end of file diff --git a/tests/integration/.codegen.json b/tests/integration/.codegen.json new file mode 100644 index 000000000..bd11afde0 --- /dev/null +++ b/tests/integration/.codegen.json @@ -0,0 +1,11 @@ +{ + "version": { + "src/databricks/labs/dqx/__about__.py": "__version__ = \"$VERSION\"" + }, + "toolchain": { + "required": ["python3", "hatch"], + "pre_setup": ["hatch env create"], + "prepend_path": ".venv/bin", + "acceptance_path": "tests/integration" + } +} \ No newline at end of file From a548a82435d1becc73be5a91c1740a764da1ae68 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Tue, 15 Jul 2025 16:52:59 +0200 Subject: [PATCH 09/91] updated acceptance action version updated nightly tests --- .github/workflows/acceptance.yml | 8 +-- .github/workflows/nightly.yml | 104 ++++++++++++++++++++++++++++++- 2 files changed, 107 insertions(+), 5 deletions(-) diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index 00559770d..d7ec6f9b4 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -54,7 +54,7 @@ jobs: # Run tests from `tests/integration` as defined in .codegen.json # and generate code coverage for modules defined in .coveragerc - name: Run integration tests and generate test coverage report - uses: databrickslabs/sandbox/acceptance@acceptance/main + uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.4 with: vault_uri: ${{ secrets.VAULT_URI }} timeout: 2h @@ -96,7 +96,7 @@ jobs: run: pip install hatch==1.9.4 - name: Run integration tests on serverless cluster - uses: databrickslabs/sandbox/acceptance@acceptance/main + uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.4 with: vault_uri: ${{ secrets.VAULT_URI }} timeout: 2h @@ -128,7 +128,7 @@ jobs: run: pip install hatch==1.9.4 - name: Run e2e tests - uses: databrickslabs/sandbox/acceptance@acceptance/main + uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.4 with: vault_uri: ${{ secrets.VAULT_URI }} timeout: 2h @@ -161,7 +161,7 @@ jobs: run: pip install hatch==1.9.4 - name: Run e2e tests on serverless cluster - uses: databrickslabs/sandbox/acceptance@acceptance/main + uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.4 with: vault_uri: ${{ secrets.VAULT_URI }} timeout: 2h diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 416039ec0..3f9719e16 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -45,11 +45,12 @@ jobs: # Run tests from `tests/integration` as defined in .codegen.json # and generate code coverage for modules defined in .coveragerc - name: Run integration tests and generate test coverage report - uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.3 + uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.4 with: vault_uri: ${{ secrets.VAULT_URI }} timeout: 2h create_issues: true + codegen_path: tests/integration/.codegen.json env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} @@ -60,3 +61,104 @@ jobs: uses: codecov/codecov-action@v5 with: use_oidc: true + + serverless_integration: + environment: tool + runs-on: larger + env: + DATABRICKS_SERVERLESS_COMPUTE_ID: auto + steps: + - name: Checkout Code + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Install Python + uses: actions/setup-python@v5 + with: + cache: 'pip' + cache-dependency-path: '**/pyproject.toml' + python-version: '3.10' + + - name: Install hatch + run: pip install hatch==1.9.4 + + - name: Run integration tests on serverless cluster + uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.4 + with: + vault_uri: ${{ secrets.VAULT_URI }} + timeout: 2h + create_issues: true + codegen_path: tests/integration/.codegen.json + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} + ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }} + DATABRICKS_SERVERLESS_COMPUTE_ID: ${{ env.DATABRICKS_SERVERLESS_COMPUTE_ID }} + + e2e: + if: github.event_name == 'pull_request' && !github.event.pull_request.draft && !github.event.pull_request.head.repo.fork + environment: tool + runs-on: larger + steps: + - name: Checkout Code + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Install Python + uses: actions/setup-python@v5 + with: + cache: 'pip' + cache-dependency-path: '**/pyproject.toml' + python-version: '3.10' + + - name: Install hatch + run: pip install hatch==1.9.4 + + - name: Run e2e tests + uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.4 + with: + vault_uri: ${{ secrets.VAULT_URI }} + timeout: 2h + create_issues: true + codegen_path: tests/e2e/.codegen.json + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} + ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }} + + e2e_serverless: + if: github.event_name == 'pull_request' && !github.event.pull_request.draft && !github.event.pull_request.head.repo.fork + environment: tool + runs-on: larger + env: + DATABRICKS_SERVERLESS_COMPUTE_ID: auto + steps: + - name: Checkout Code + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Install Python + uses: actions/setup-python@v5 + with: + cache: 'pip' + cache-dependency-path: '**/pyproject.toml' + python-version: '3.10' + + - name: Install hatch + run: pip install hatch==1.9.4 + + - name: Run e2e tests on serverless cluster + uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.4 + with: + vault_uri: ${{ secrets.VAULT_URI }} + timeout: 2h + create_issues: true + codegen_path: tests/e2e/.codegen.json + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} + ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }} + DATABRICKS_SERVERLESS_COMPUTE_ID: ${{ env.DATABRICKS_SERVERLESS_COMPUTE_ID }} From 87de2f6a9dd31c5551cc28d1010b01043ed08d38 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Tue, 15 Jul 2025 17:16:58 +0200 Subject: [PATCH 10/91] refactor --- .github/workflows/acceptance.yml | 2 +- .github/workflows/nightly.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index d7ec6f9b4..8a9a051c1 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -70,7 +70,7 @@ jobs: with: use_oidc: true - serverless_integration: + integration_serverless: # Only run this job for PRs from branches on the main repository and not from forks. # Workflows triggered by PRs from forks don't have access to the tool environment. # PRs from forks to be tested by the reviewer(s) / maintainer(s) before merging. diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 3f9719e16..0d1b8c58a 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -62,7 +62,7 @@ jobs: with: use_oidc: true - serverless_integration: + integration_serverless: environment: tool runs-on: larger env: From 75d49ae16fc9db13f74f51f4cf7d61629eb27969 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 11:32:22 -0400 Subject: [PATCH 11/91] Send path as string --- tests/e2e/test_run_demos.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 659236516..c1992a9f1 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -23,7 +23,9 @@ async def get_submit_job_result(notebook_path): with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.AUTO) return await ws.jobs.submit( - tasks=[SubmitTask(task_key="demo_run", notebook_task=NotebookTask(notebook_path=notebook.as_fuse()))] + tasks=[ + SubmitTask(task_key="demo_run", notebook_task=NotebookTask(notebook_path=notebook.as_fuse().as_posix())) + ] ) async def validate_submit_job_runs(runs): From d377992c03885eb6da15c3b0369765470ae7a6a9 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 12:04:40 -0400 Subject: [PATCH 12/91] Refactor --- Makefile | 2 +- tests/e2e/test_run_demos.py | 22 +++++++++++++++------- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/Makefile b/Makefile index 48f62df3b..f86616d7b 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -all: clean dev lint fmt test integration coverage end-to-end +all: clean dev lint fmt test integration coverage e2e clean: rm -fr .venv clean htmlcov .mypy_cache .pytest_cache .ruff_cache .coverage coverage.xml diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index c1992a9f1..82c339181 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -18,24 +18,32 @@ def test_run_all_demo_notebooks_succeed(make_notebook): notebook_paths = [path for path in os.listdir(demo_folder) if path.endswith((".py", ".ipynb", ".dbc"))] ws = WorkspaceClient() - async def get_submit_job_result(notebook_path): + run_ids = [] + for notebook_path in notebook_paths: path = demo_folder / notebook_path with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.AUTO) - return await ws.jobs.submit( + job_run = ws.jobs.submit( tasks=[ SubmitTask(task_key="demo_run", notebook_task=NotebookTask(notebook_path=notebook.as_fuse().as_posix())) ] ) + run_ids.append([job_run.run_id]) - async def validate_submit_job_runs(runs): - for completion in asyncio.as_completed(runs): - run = await completion + async def wait_for_completion(run_id, poll_interval=10): + while True: + run = ws.jobs.get_run(run_id) if run.status.state == RunLifecycleStateV2State.TERMINATED: run_details = run.status.termination_details assert ( run_details.type == TerminationTypeType.SUCCESS ), f"Run ended with status {run_details.type.value}: {run_details.message}" + return + await asyncio.sleep(poll_interval) + + async def validate_submit_job_runs(runs): + for completion in asyncio.as_completed(runs): + await completion - runs = [get_submit_job_result(notebook_path) for notebook_path in notebook_paths] - asyncio.run(validate_submit_job_runs(runs)) + demo_runs = [wait_for_completion(run_id) for run_id in run_ids] + asyncio.run(validate_submit_job_runs(demo_runs)) From 44bf5ef5409ea254a805feba60f744e7737fb32c Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 12:28:21 -0400 Subject: [PATCH 13/91] Fix import format --- tests/e2e/test_run_demos.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 82c339181..9baaaea53 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -20,9 +20,15 @@ def test_run_all_demo_notebooks_succeed(make_notebook): run_ids = [] for notebook_path in notebook_paths: + logger.info(f"Running demo notebook '{notebook_path}'") path = demo_folder / notebook_path + import_format = ( + ImportFormat.JUPYTER + if notebook_path.endswith(".ipynb") + else ImportFormat.DBC if notebook_path.endswith(".dbc") else ImportFormat.SOURCE + ) with open(path, "rb") as f: - notebook = make_notebook(content=f, format=ImportFormat.AUTO) + notebook = make_notebook(content=f, format=import_format) job_run = ws.jobs.submit( tasks=[ SubmitTask(task_key="demo_run", notebook_task=NotebookTask(notebook_path=notebook.as_fuse().as_posix())) @@ -42,7 +48,7 @@ async def wait_for_completion(run_id, poll_interval=10): await asyncio.sleep(poll_interval) async def validate_submit_job_runs(runs): - for completion in asyncio.as_completed(runs): + for completion in asyncio.as_completed(runs, timeout=1800): await completion demo_runs = [wait_for_completion(run_id) for run_id in run_ids] From 083f3ff9fb15037f825acfae7cd97b3bfe0a4995 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 12:43:22 -0400 Subject: [PATCH 14/91] Clarify error messages --- tests/e2e/test_run_demos.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 9baaaea53..429af37f8 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -31,19 +31,21 @@ def test_run_all_demo_notebooks_succeed(make_notebook): notebook = make_notebook(content=f, format=import_format) job_run = ws.jobs.submit( tasks=[ - SubmitTask(task_key="demo_run", notebook_task=NotebookTask(notebook_path=notebook.as_fuse().as_posix())) + SubmitTask( + task_key=notebook_path, notebook_task=NotebookTask(notebook_path=notebook.as_fuse().as_posix()) + ) ] ) run_ids.append([job_run.run_id]) - async def wait_for_completion(run_id, poll_interval=10): + async def wait_for_completion(run_id, poll_interval=30): while True: run = ws.jobs.get_run(run_id) if run.status.state == RunLifecycleStateV2State.TERMINATED: run_details = run.status.termination_details assert ( run_details.type == TerminationTypeType.SUCCESS - ), f"Run ended with status {run_details.type.value}: {run_details.message}" + ), f"Run of '{run.tasks[0].task_key}' failed with output: {run.tasks[0].state.state_message}" return await asyncio.sleep(poll_interval) From b642196d3cd772d42f094bb504564723abb5ef0e Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 12:52:18 -0400 Subject: [PATCH 15/91] Fix task key characters --- tests/e2e/test_run_demos.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 429af37f8..c1a334ce2 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -32,7 +32,8 @@ def test_run_all_demo_notebooks_succeed(make_notebook): job_run = ws.jobs.submit( tasks=[ SubmitTask( - task_key=notebook_path, notebook_task=NotebookTask(notebook_path=notebook.as_fuse().as_posix()) + task_key=notebook_path.replace(".", "_"), + notebook_task=NotebookTask(notebook_path=notebook.as_fuse().as_posix()), ) ] ) From e306c10896d3baf9c5239a0020c9c2eacf7f679b Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 14:21:41 -0400 Subject: [PATCH 16/91] Sample test a few notebooks --- tests/e2e/test_run_demos.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index c1a334ce2..bbb8549e8 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -12,10 +12,17 @@ logging.getLogger("databricks.labs.dqx").setLevel("DEBUG") logger = logging.getLogger(__name__) +INCLUDE_PATHS = [ + "dqx_demo_library.py", + "dqx_demo_pii_detection.py", + "dqx_demo_manufacturing_demo.py", + "dqx_quick_start_demo.py", + ] + def test_run_all_demo_notebooks_succeed(make_notebook): demo_folder = Path(__file__).parent.parent.parent / "demos" - notebook_paths = [path for path in os.listdir(demo_folder) if path.endswith((".py", ".ipynb", ".dbc"))] + notebook_paths = [path for path in os.listdir(demo_folder) if path in INCLUDE_PATHS] ws = WorkspaceClient() run_ids = [] From cd4998bd75c272c84b097575d87fabf291efb7fa Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 15:15:30 -0400 Subject: [PATCH 17/91] Refactor --- tests/e2e/test_run_demos.py | 100 +++++++++++++++++------------------- 1 file changed, 48 insertions(+), 52 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index bbb8549e8..87c7ac676 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -1,65 +1,61 @@ -import asyncio import logging -import os from pathlib import Path from databricks.sdk import WorkspaceClient from databricks.sdk.service.workspace import ImportFormat -from databricks.sdk.service.jobs import NotebookTask, SubmitTask, RunLifecycleStateV2State, TerminationTypeType +from databricks.sdk.service.jobs import NotebookTask, SubmitTask, TerminationTypeType logging.getLogger("tests").setLevel("DEBUG") logging.getLogger("databricks.labs.dqx").setLevel("DEBUG") logger = logging.getLogger(__name__) -INCLUDE_PATHS = [ - "dqx_demo_library.py", - "dqx_demo_pii_detection.py", - "dqx_demo_manufacturing_demo.py", - "dqx_quick_start_demo.py", - ] - -def test_run_all_demo_notebooks_succeed(make_notebook): - demo_folder = Path(__file__).parent.parent.parent / "demos" - notebook_paths = [path for path in os.listdir(demo_folder) if path in INCLUDE_PATHS] +def test_run_dqx_demo_library(make_notebook): + path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_library.py" ws = WorkspaceClient() - - run_ids = [] - for notebook_path in notebook_paths: - logger.info(f"Running demo notebook '{notebook_path}'") - path = demo_folder / notebook_path - import_format = ( - ImportFormat.JUPYTER - if notebook_path.endswith(".ipynb") - else ImportFormat.DBC if notebook_path.endswith(".dbc") else ImportFormat.SOURCE - ) - with open(path, "rb") as f: - notebook = make_notebook(content=f, format=import_format) - job_run = ws.jobs.submit( - tasks=[ - SubmitTask( - task_key=notebook_path.replace(".", "_"), - notebook_task=NotebookTask(notebook_path=notebook.as_fuse().as_posix()), - ) - ] - ) - run_ids.append([job_run.run_id]) - - async def wait_for_completion(run_id, poll_interval=30): - while True: - run = ws.jobs.get_run(run_id) - if run.status.state == RunLifecycleStateV2State.TERMINATED: - run_details = run.status.termination_details - assert ( - run_details.type == TerminationTypeType.SUCCESS - ), f"Run of '{run.tasks[0].task_key}' failed with output: {run.tasks[0].state.state_message}" - return - await asyncio.sleep(poll_interval) - - async def validate_submit_job_runs(runs): - for completion in asyncio.as_completed(runs, timeout=1800): - await completion - - demo_runs = [wait_for_completion(run_id) for run_id in run_ids] - asyncio.run(validate_submit_job_runs(demo_runs)) + with open(path, "rb") as f: + notebook = make_notebook(content=f, format=ImportFormat.SOURCE) + notebook_path = notebook.as_fuse().as_posix() + run = ws.jobs.submit_and_wait( + tasks=[SubmitTask(task_key="dqx_demo_library", notebook_task=NotebookTask(notebook_path=notebook_path))] + ) + run_details = run.status.termination_details + task = run.tasks[0] + assert ( + run_details.type == TerminationTypeType.SUCCESS + ), f"Run of '{task.task_key}' failed with output: {task.status.termination_details.message}" + + +def test_run_dqx_manufacturing_demo(make_notebook): + path = Path(__file__).parent.parent.parent / "demos" / "dqx_manufacturing_demo.py" + ws = WorkspaceClient() + with open(path, "rb") as f: + notebook = make_notebook(content=f, format=ImportFormat.SOURCE) + notebook_path = notebook.as_fuse().as_posix() + run = ws.jobs.submit_and_wait( + tasks=[SubmitTask(task_key="dqx_manufacturing_demo", notebook_task=NotebookTask(notebook_path=notebook_path))] + ) + run_details = run.status.termination_details + task = run.tasks[0] + assert ( + run_details.type == TerminationTypeType.SUCCESS + ), f"Run of '{task.task_key}' failed with output: {task.status.termination_details.message}" + + +def test_run_dqx_quick_start_demo_library(make_notebook): + path = Path(__file__).parent.parent.parent / "demos" / "dqx_quick_start_demo_library.py" + ws = WorkspaceClient() + with open(path, "rb") as f: + notebook = make_notebook(content=f, format=ImportFormat.SOURCE) + notebook_path = notebook.as_fuse().as_posix() + run = ws.jobs.submit_and_wait( + tasks=[ + SubmitTask(task_key="dqx_quick_start_demo_library", notebook_task=NotebookTask(notebook_path=notebook_path)) + ] + ) + run_details = run.status.termination_details + task = run.tasks[0] + assert ( + run_details.type == TerminationTypeType.SUCCESS + ), f"Run of '{task.task_key}' failed with output: {task.status.termination_details.message}" From 701aa9a0402b63af97d72891e3c66fe0146d7774 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 17:07:00 -0400 Subject: [PATCH 18/91] Parameterize schemas and catalogs --- demos/dqx_demo_library.py | 38 +++++++++++++++++++++++++---------- tests/e2e/test_run_demos.py | 40 +++++++++++++++++++++++++------------ 2 files changed, 55 insertions(+), 23 deletions(-) diff --git a/demos/dqx_demo_library.py b/demos/dqx_demo_library.py index a2fcc726d..5c752e628 100644 --- a/demos/dqx_demo_library.py +++ b/demos/dqx_demo_library.py @@ -14,6 +14,24 @@ # COMMAND ---------- +# Set Catalog and Schema for Demo Dataset +default_database = "main" +default_schema_name = "default" + +dbutils.widgets.text("demo_database", default_database, "Catalog Name") +dbutils.widgets.text("demo_schema", default_schema_name, "Schema Name") + +database = dbutils.widgets.get("demo_database") +schema = dbutils.widgets.get("demo_schema") + +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"CREATE SCHEMA IF NOT EXISTS {schema}") + +# COMMAND ---------- + # MAGIC %md # MAGIC ## Generation of quality rule/check candidates using Profiler # MAGIC Data profiling is typically performed as a one-time action for the input dataset to discover the initial set of quality rule candidates. @@ -65,7 +83,7 @@ dq_engine.save_checks_in_workspace_file(checks=checks, workspace_path=checks_file) # save generated checks in a Delta table -dq_engine.save_checks_in_table(checks=checks, table_name="main.default.dqx_checks_table", mode="overwrite") +dq_engine.save_checks_in_table(checks=checks, table_name=f"{database}.{schema}.dqx_checks_table", mode="overwrite") # COMMAND ---------- @@ -106,7 +124,7 @@ # load checks from a Delta table dq_engine = DQEngine(WorkspaceClient()) -checks = dq_engine.load_checks_from_table(table_name="main.default.dqx_checks_table") +checks = dq_engine.load_checks_from_table(table_name=f"{database}.{schema}.dqx_checks_table") # Option 1: apply quality rules and quarantine invalid records valid_df, quarantine_df = dq_engine.apply_checks_by_metadata_and_split(input_df, checks) @@ -429,17 +447,17 @@ dq_engine.save_results_in_table( output_df=silver_df, quarantine_df=quarantine_df, - output_config=OutputConfig("main.default.dqx_output", mode="overwrite"), - quarantine_config=OutputConfig("main.default.dqx_quarantine", mode="overwrite") + output_config=OutputConfig(f"{database}.{schema}.dqx_output", mode="overwrite"), + quarantine_config=OutputConfig(f"{database}.{schema}.dqx_quarantine", mode="overwrite") ) # COMMAND ---------- -display(spark.table("main.default.dqx_output")) +display(spark.table(f"{database}.{schema}.dqx_output")) # COMMAND ---------- -display(spark.table("main.default.dqx_quarantine")) +display(spark.table(f"{database}.{schema}.dqx_quarantine")) # COMMAND ---------- @@ -452,13 +470,13 @@ dq_engine.apply_checks_by_metadata_and_save_in_table( input_config=InputConfig("/databricks-datasets/delta-sharing/samples/nyctaxi_2019"), checks=checks, - output_config=OutputConfig("main.default.dqx_e2e_output", mode="overwrite"), - quarantine_config=OutputConfig("main.default.dqx_e2e_quarantine", mode="overwrite") + output_config=OutputConfig(f"{database}.{schema}.dqx_e2e_output", mode="overwrite"), + quarantine_config=OutputConfig(f"{database}.{schema}.dqx_e2e_quarantine", mode="overwrite") ) # display the results saved to output and quarantine tables -display(spark.table("main.default.dqx_e2e_output")) -display(spark.table("main.default.dqx_e2e_quarantine")) +display(spark.table(f"{database}.{schema}.dqx_e2e_output")) +display(spark.table(f"{database}.{schema}.dqx_e2e_quarantine")) # COMMAND ---------- diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 87c7ac676..f5abd9e36 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -11,36 +11,52 @@ logger = logging.getLogger(__name__) -def test_run_dqx_demo_library(make_notebook): +def test_run_dqx_demo_library(make_notebook, make_schema): path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_library.py" ws = WorkspaceClient() with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) + + catalog = "main" + schema = make_schema(catalog_name=catalog).name notebook_path = notebook.as_fuse().as_posix() run = ws.jobs.submit_and_wait( - tasks=[SubmitTask(task_key="dqx_demo_library", notebook_task=NotebookTask(notebook_path=notebook_path))] + tasks=[ + SubmitTask( + task_key="dqx_demo_library", + notebook_task=NotebookTask( + notebook_path=notebook_path, base_parameters={"demo_database": catalog, "demo_schema": schema} + ), + ) + ] ) run_details = run.status.termination_details task = run.tasks[0] - assert ( - run_details.type == TerminationTypeType.SUCCESS - ), f"Run of '{task.task_key}' failed with output: {task.status.termination_details.message}" + assert run_details.type == TerminationTypeType.SUCCESS, f"Run of '{task.task_key}' failed" -def test_run_dqx_manufacturing_demo(make_notebook): +def test_run_dqx_manufacturing_demo(make_notebook, make_schema): path = Path(__file__).parent.parent.parent / "demos" / "dqx_manufacturing_demo.py" ws = WorkspaceClient() with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) + + catalog = "main" + schema = make_schema(catalog_name=catalog).name notebook_path = notebook.as_fuse().as_posix() run = ws.jobs.submit_and_wait( - tasks=[SubmitTask(task_key="dqx_manufacturing_demo", notebook_task=NotebookTask(notebook_path=notebook_path))] + tasks=[ + SubmitTask( + task_key="dqx_manufacturing_demo", + notebook_task=NotebookTask( + notebook_path=notebook_path, base_parameters={"demo_database": catalog, "demo_schema": schema} + ), + ) + ] ) run_details = run.status.termination_details task = run.tasks[0] - assert ( - run_details.type == TerminationTypeType.SUCCESS - ), f"Run of '{task.task_key}' failed with output: {task.status.termination_details.message}" + assert run_details.type == TerminationTypeType.SUCCESS, f"Run of '{task.task_key}' failed" def test_run_dqx_quick_start_demo_library(make_notebook): @@ -56,6 +72,4 @@ def test_run_dqx_quick_start_demo_library(make_notebook): ) run_details = run.status.termination_details task = run.tasks[0] - assert ( - run_details.type == TerminationTypeType.SUCCESS - ), f"Run of '{task.task_key}' failed with output: {task.status.termination_details.message}" + assert run_details.type == TerminationTypeType.SUCCESS, f"Run of '{task.task_key}' failed" From 2554020d1c15233f216ba6785abd17b79a2ad4c6 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 18:08:33 -0400 Subject: [PATCH 19/91] Refactor --- tests/e2e/test_run_demos.py | 82 ++++++++++++++++++++++--------------- 1 file changed, 49 insertions(+), 33 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index f5abd9e36..5c2f238c0 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -1,9 +1,10 @@ import logging +import time from pathlib import Path from databricks.sdk import WorkspaceClient from databricks.sdk.service.workspace import ImportFormat -from databricks.sdk.service.jobs import NotebookTask, SubmitTask, TerminationTypeType +from databricks.sdk.service.jobs import NotebookTask, Task, RunLifecycleStateV2State, TerminationTypeType logging.getLogger("tests").setLevel("DEBUG") @@ -11,7 +12,10 @@ logger = logging.getLogger(__name__) -def test_run_dqx_demo_library(make_notebook, make_schema): +RETRY_INTERVAL_SECONDS = 30 + + +def test_run_dqx_demo_library(make_notebook, make_schema, make_job): path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_library.py" ws = WorkspaceClient() with open(path, "rb") as f: @@ -20,22 +24,25 @@ def test_run_dqx_demo_library(make_notebook, make_schema): catalog = "main" schema = make_schema(catalog_name=catalog).name notebook_path = notebook.as_fuse().as_posix() - run = ws.jobs.submit_and_wait( - tasks=[ - SubmitTask( - task_key="dqx_demo_library", - notebook_task=NotebookTask( - notebook_path=notebook_path, base_parameters={"demo_database": catalog, "demo_schema": schema} - ), - ) - ] + notebook_task = NotebookTask( + notebook_path=notebook_path, + base_parameters={"demo_database": catalog, "demo_schema": schema} ) - run_details = run.status.termination_details + job = make_job(tasks=[Task(task_key="dqx_demo_library", notebook_task=notebook_task)]) + run = ws.jobs.run_now(job.job_id) + + while True: + run_details = ws.jobs.get_run(run.run_id) + if run_details.status.state == RunLifecycleStateV2State.TERMINATED: + break + time.sleep(RETRY_INTERVAL_SECONDS) + task = run.tasks[0] - assert run_details.type == TerminationTypeType.SUCCESS, f"Run of '{task.task_key}' failed" + termination_details = run_details.status.termination_details + assert termination_details.type == TerminationTypeType.SUCCESS, f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" -def test_run_dqx_manufacturing_demo(make_notebook, make_schema): +def test_run_dqx_manufacturing_demo(make_notebook, make_schema, make_job): path = Path(__file__).parent.parent.parent / "demos" / "dqx_manufacturing_demo.py" ws = WorkspaceClient() with open(path, "rb") as f: @@ -44,32 +51,41 @@ def test_run_dqx_manufacturing_demo(make_notebook, make_schema): catalog = "main" schema = make_schema(catalog_name=catalog).name notebook_path = notebook.as_fuse().as_posix() - run = ws.jobs.submit_and_wait( - tasks=[ - SubmitTask( - task_key="dqx_manufacturing_demo", - notebook_task=NotebookTask( - notebook_path=notebook_path, base_parameters={"demo_database": catalog, "demo_schema": schema} - ), - ) - ] + notebook_task = NotebookTask( + notebook_path=notebook_path, + base_parameters={"demo_database": catalog, "demo_schema": schema} ) - run_details = run.status.termination_details + job = make_job(tasks=[Task(task_key="dqx_manufacturing_demo", notebook_task=notebook_task)]) + run = ws.jobs.run_now(job.job_id) + + while True: + run_details = ws.jobs.get_run(run.run_id) + if run_details.status.state == RunLifecycleStateV2State.TERMINATED: + break + time.sleep(RETRY_INTERVAL_SECONDS) + task = run.tasks[0] - assert run_details.type == TerminationTypeType.SUCCESS, f"Run of '{task.task_key}' failed" + termination_details = run_details.status.termination_details + assert termination_details.type == TerminationTypeType.SUCCESS, f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" -def test_run_dqx_quick_start_demo_library(make_notebook): +def test_run_dqx_quick_start_demo_library(make_notebook, make_job): path = Path(__file__).parent.parent.parent / "demos" / "dqx_quick_start_demo_library.py" ws = WorkspaceClient() with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) + notebook_path = notebook.as_fuse().as_posix() - run = ws.jobs.submit_and_wait( - tasks=[ - SubmitTask(task_key="dqx_quick_start_demo_library", notebook_task=NotebookTask(notebook_path=notebook_path)) - ] - ) - run_details = run.status.termination_details + notebook_task = NotebookTask(notebook_path=notebook_path) + job = make_job(tasks=[Task(task_key="dqx_quick_start_demo_library", notebook_task=notebook_task)]) + run = ws.jobs.run_now(job.job_id) + + while True: + run_details = ws.jobs.get_run(run.run_id) + if run_details.status.state == RunLifecycleStateV2State.TERMINATED: + break + time.sleep(RETRY_INTERVAL_SECONDS) + task = run.tasks[0] - assert run_details.type == TerminationTypeType.SUCCESS, f"Run of '{task.task_key}' failed" + termination_details = run_details.status.termination_details + assert termination_details.type == TerminationTypeType.SUCCESS, f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" From 68ffa9b0d17cb5c45b2f243cd920115e74576316 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 18:08:59 -0400 Subject: [PATCH 20/91] Format --- tests/e2e/test_run_demos.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 5c2f238c0..bd4d452e2 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -25,8 +25,7 @@ def test_run_dqx_demo_library(make_notebook, make_schema, make_job): schema = make_schema(catalog_name=catalog).name notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask( - notebook_path=notebook_path, - base_parameters={"demo_database": catalog, "demo_schema": schema} + notebook_path=notebook_path, base_parameters={"demo_database": catalog, "demo_schema": schema} ) job = make_job(tasks=[Task(task_key="dqx_demo_library", notebook_task=notebook_task)]) run = ws.jobs.run_now(job.job_id) @@ -39,7 +38,9 @@ def test_run_dqx_demo_library(make_notebook, make_schema, make_job): task = run.tasks[0] termination_details = run_details.status.termination_details - assert termination_details.type == TerminationTypeType.SUCCESS, f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" + assert ( + termination_details.type == TerminationTypeType.SUCCESS + ), f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" def test_run_dqx_manufacturing_demo(make_notebook, make_schema, make_job): @@ -52,8 +53,7 @@ def test_run_dqx_manufacturing_demo(make_notebook, make_schema, make_job): schema = make_schema(catalog_name=catalog).name notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask( - notebook_path=notebook_path, - base_parameters={"demo_database": catalog, "demo_schema": schema} + notebook_path=notebook_path, base_parameters={"demo_database": catalog, "demo_schema": schema} ) job = make_job(tasks=[Task(task_key="dqx_manufacturing_demo", notebook_task=notebook_task)]) run = ws.jobs.run_now(job.job_id) @@ -66,7 +66,9 @@ def test_run_dqx_manufacturing_demo(make_notebook, make_schema, make_job): task = run.tasks[0] termination_details = run_details.status.termination_details - assert termination_details.type == TerminationTypeType.SUCCESS, f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" + assert ( + termination_details.type == TerminationTypeType.SUCCESS + ), f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" def test_run_dqx_quick_start_demo_library(make_notebook, make_job): @@ -88,4 +90,6 @@ def test_run_dqx_quick_start_demo_library(make_notebook, make_job): task = run.tasks[0] termination_details = run_details.status.termination_details - assert termination_details.type == TerminationTypeType.SUCCESS, f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" + assert ( + termination_details.type == TerminationTypeType.SUCCESS + ), f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" From 9607bb3f357182b9c542ec3d3ad3263b3259e918 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 18:15:24 -0400 Subject: [PATCH 21/91] Fix test --- tests/e2e/test_run_demos.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index bd4d452e2..245ed3b07 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -36,7 +36,7 @@ def test_run_dqx_demo_library(make_notebook, make_schema, make_job): break time.sleep(RETRY_INTERVAL_SECONDS) - task = run.tasks[0] + task = run_details.tasks[0] termination_details = run_details.status.termination_details assert ( termination_details.type == TerminationTypeType.SUCCESS @@ -64,7 +64,7 @@ def test_run_dqx_manufacturing_demo(make_notebook, make_schema, make_job): break time.sleep(RETRY_INTERVAL_SECONDS) - task = run.tasks[0] + task = run_details.tasks[0] termination_details = run_details.status.termination_details assert ( termination_details.type == TerminationTypeType.SUCCESS @@ -88,7 +88,7 @@ def test_run_dqx_quick_start_demo_library(make_notebook, make_job): break time.sleep(RETRY_INTERVAL_SECONDS) - task = run.tasks[0] + task = run_details.tasks[0] termination_details = run_details.status.termination_details assert ( termination_details.type == TerminationTypeType.SUCCESS From 634bc22d389ad61c592288dc61f0dc77f7b9c9c2 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 20:02:51 -0400 Subject: [PATCH 22/91] Use make_catalog --- tests/e2e/test_run_demos.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 245ed3b07..048ad0e76 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -13,16 +13,17 @@ RETRY_INTERVAL_SECONDS = 30 +CATALOG_NAME = "main" -def test_run_dqx_demo_library(make_notebook, make_schema, make_job): +def test_run_dqx_demo_library(make_notebook, make_catalog, make_schema, make_job): path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_library.py" ws = WorkspaceClient() with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) - catalog = "main" - schema = make_schema(catalog_name=catalog).name + catalog = make_catalog(name=CATALOG_NAME).name + schema = make_schema(catalog_name=CATALOG_NAME).name notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask( notebook_path=notebook_path, base_parameters={"demo_database": catalog, "demo_schema": schema} @@ -43,14 +44,14 @@ def test_run_dqx_demo_library(make_notebook, make_schema, make_job): ), f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" -def test_run_dqx_manufacturing_demo(make_notebook, make_schema, make_job): +def test_run_dqx_manufacturing_demo(make_notebook, make_catalog, make_schema, make_job): path = Path(__file__).parent.parent.parent / "demos" / "dqx_manufacturing_demo.py" ws = WorkspaceClient() with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) - catalog = "main" - schema = make_schema(catalog_name=catalog).name + catalog = make_catalog(name=CATALOG_NAME).name + schema = make_schema(catalog_name=CATALOG_NAME).name notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask( notebook_path=notebook_path, base_parameters={"demo_database": catalog, "demo_schema": schema} From a7bb0cd090b994e2054ab8d6df31c6b6399c5edc Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 20:18:27 -0400 Subject: [PATCH 23/91] Update demos and tests --- demos/dqx_demo_library.py | 3 --- demos/dqx_demo_pii_detection.py | 2 +- demos/dqx_manufacturing_demo.py | 5 ----- tests/e2e/test_run_demos.py | 33 ++++++++++++++++++++++++++++----- 4 files changed, 29 insertions(+), 14 deletions(-) diff --git a/demos/dqx_demo_library.py b/demos/dqx_demo_library.py index 5c752e628..8560a82f6 100644 --- a/demos/dqx_demo_library.py +++ b/demos/dqx_demo_library.py @@ -27,9 +27,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"CREATE SCHEMA IF NOT EXISTS {schema}") - # COMMAND ---------- # MAGIC %md diff --git a/demos/dqx_demo_pii_detection.py b/demos/dqx_demo_pii_detection.py index 0426d02c9..df27081a5 100644 --- a/demos/dqx_demo_pii_detection.py +++ b/demos/dqx_demo_pii_detection.py @@ -17,7 +17,7 @@ # COMMAND ---------- -# MAGIC %pip install databricks-labs-dqx presidio-analyzer==2.2.358 numpy==1.26 --quiet +# MAGIC %pip install databricks-labs-dqx==0.5.0 presidio-analyzer==2.2.358 numpy==1.26 --quiet # COMMAND ---------- diff --git a/demos/dqx_manufacturing_demo.py b/demos/dqx_manufacturing_demo.py index 18291ee20..dc2df4d57 100644 --- a/demos/dqx_manufacturing_demo.py +++ b/demos/dqx_manufacturing_demo.py @@ -92,11 +92,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" diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 048ad0e76..beb175b40 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -13,7 +13,6 @@ RETRY_INTERVAL_SECONDS = 30 -CATALOG_NAME = "main" def test_run_dqx_demo_library(make_notebook, make_catalog, make_schema, make_job): @@ -22,8 +21,8 @@ def test_run_dqx_demo_library(make_notebook, make_catalog, make_schema, make_job with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) - catalog = make_catalog(name=CATALOG_NAME).name - schema = make_schema(catalog_name=CATALOG_NAME).name + catalog = "main" + schema = make_schema(catalog_name=catalog).name notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask( notebook_path=notebook_path, base_parameters={"demo_database": catalog, "demo_schema": schema} @@ -50,8 +49,8 @@ def test_run_dqx_manufacturing_demo(make_notebook, make_catalog, make_schema, ma with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) - catalog = make_catalog(name=CATALOG_NAME).name - schema = make_schema(catalog_name=CATALOG_NAME).name + catalog = "main" + schema = make_schema(catalog_name=catalog).name notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask( notebook_path=notebook_path, base_parameters={"demo_database": catalog, "demo_schema": schema} @@ -94,3 +93,27 @@ def test_run_dqx_quick_start_demo_library(make_notebook, make_job): assert ( termination_details.type == TerminationTypeType.SUCCESS ), f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" + + +def test_run_dqx_demo_pii_detection(make_notebook, make_job): + path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_pii_detection.py" + ws = WorkspaceClient() + with open(path, "rb") as f: + notebook = make_notebook(content=f, format=ImportFormat.SOURCE) + + notebook_path = notebook.as_fuse().as_posix() + notebook_task = NotebookTask(notebook_path=notebook_path) + job = make_job(tasks=[Task(task_key="dqx_demo_pii_detection", notebook_task=notebook_task)]) + run = ws.jobs.run_now(job.job_id) + + while True: + run_details = ws.jobs.get_run(run.run_id) + if run_details.status.state == RunLifecycleStateV2State.TERMINATED: + break + time.sleep(RETRY_INTERVAL_SECONDS) + + task = run_details.tasks[0] + termination_details = run_details.status.termination_details + assert ( + termination_details.type == TerminationTypeType.SUCCESS + ), f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" From 967fde4559dd1ec95ae2e041763defd9e8e4c573 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 20:53:25 -0400 Subject: [PATCH 24/91] Fix tests --- tests/e2e/test_run_demos.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index beb175b40..4b7c4f9a7 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -43,11 +43,12 @@ def test_run_dqx_demo_library(make_notebook, make_catalog, make_schema, make_job ), f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" -def test_run_dqx_manufacturing_demo(make_notebook, make_catalog, make_schema, make_job): +def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema, make_job): path = Path(__file__).parent.parent.parent / "demos" / "dqx_manufacturing_demo.py" ws = WorkspaceClient() with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) + make_directory("quality_rules") catalog = "main" schema = make_schema(catalog_name=catalog).name @@ -95,7 +96,7 @@ def test_run_dqx_quick_start_demo_library(make_notebook, make_job): ), f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" -def test_run_dqx_demo_pii_detection(make_notebook, make_job): +def test_run_dqx_demo_pii_detection(make_notebook, make_cluster, make_job): path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_pii_detection.py" ws = WorkspaceClient() with open(path, "rb") as f: @@ -103,7 +104,8 @@ def test_run_dqx_demo_pii_detection(make_notebook, make_job): notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask(notebook_path=notebook_path) - job = make_job(tasks=[Task(task_key="dqx_demo_pii_detection", notebook_task=notebook_task)]) + cluster = make_cluster(spark_version="15.4.x-scala2.12") + job = make_job(tasks=[Task(task_key="dqx_demo_pii_detection", notebook_task=notebook_task, existing_cluster_id=cluster.cluster_id)]) run = ws.jobs.run_now(job.job_id) while True: From 278c90bd902cec6b881e4cadc8696d733a909da2 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 21:01:11 -0400 Subject: [PATCH 25/91] Fix tests --- tests/e2e/test_run_demos.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 4b7c4f9a7..8e6e8f11a 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -48,7 +48,7 @@ def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema, ws = WorkspaceClient() with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) - make_directory("quality_rules") + make_directory(path="quality_rules") catalog = "main" schema = make_schema(catalog_name=catalog).name @@ -104,7 +104,7 @@ def test_run_dqx_demo_pii_detection(make_notebook, make_cluster, make_job): notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask(notebook_path=notebook_path) - cluster = make_cluster(spark_version="15.4.x-scala2.12") + cluster = make_cluster(single_node=True, spark_version="15.4.x-scala2.12") job = make_job(tasks=[Task(task_key="dqx_demo_pii_detection", notebook_task=notebook_task, existing_cluster_id=cluster.cluster_id)]) run = ws.jobs.run_now(job.job_id) From 2b251737630a4557657b028eca25836eadeba2c6 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 15 Jul 2025 21:04:37 -0400 Subject: [PATCH 26/91] Fix tests --- tests/e2e/test_run_demos.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 8e6e8f11a..a66d1bf04 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -48,7 +48,7 @@ def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema, ws = WorkspaceClient() with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) - make_directory(path="quality_rules") + make_directory(path="/quality_rules") catalog = "main" schema = make_schema(catalog_name=catalog).name From 205370ad12bda5c6860e1d97c28381b37788bf80 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Thu, 24 Jul 2025 23:10:46 -0400 Subject: [PATCH 27/91] Fix tests --- tests/e2e/test_run_demos.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index a66d1bf04..7ab13c096 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -48,7 +48,8 @@ def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema, ws = WorkspaceClient() with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) - make_directory(path="/quality_rules") + folder = notebook.as_fuse().parent / "quality_rules" + make_directory(path=folder) catalog = "main" schema = make_schema(catalog_name=catalog).name From f43ec33f82cd4686e5fce8553eac4a4cda0da229 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Fri, 25 Jul 2025 16:17:14 -0400 Subject: [PATCH 28/91] Fix demos & tests --- demos/dqx_demo_library.py | 408 ++++++++++++++++---------------- demos/dqx_demo_pii_detection.py | 69 +++--- demos/dqx_manufacturing_demo.py | 16 +- tests/e2e/test_run_demos.py | 8 +- 4 files changed, 251 insertions(+), 250 deletions(-) diff --git a/demos/dqx_demo_library.py b/demos/dqx_demo_library.py index 3e591d0bc..f373bd141 100644 --- a/demos/dqx_demo_library.py +++ b/demos/dqx_demo_library.py @@ -14,21 +14,6 @@ # COMMAND ---------- -# Set Catalog and Schema for Demo Dataset -default_database = "main" -default_schema_name = "default" - -dbutils.widgets.text("demo_database", default_database, "Catalog Name") -dbutils.widgets.text("demo_schema", default_schema_name, "Schema Name") - -database = dbutils.widgets.get("demo_database") -schema = dbutils.widgets.get("demo_schema") - -print(f"Selected Catalog for Demo Dataset: {database}") -print(f"Selected Schema for Demo Dataset: {schema}") - -# COMMAND ---------- - # MAGIC %md # MAGIC ## Generation of quality rule/check candidates using Profiler # MAGIC Data profiling is typically performed as a one-time action for the input dataset to discover the initial set of quality rule candidates. @@ -41,8 +26,21 @@ from databricks.labs.dqx.profiler.dlt_generator import DQDltGenerator from databricks.labs.dqx.engine import DQEngine from databricks.sdk import WorkspaceClient +import os import yaml +default_file_directory = os.getcwd() +default_catalog = "main" +default_schema = "default" + +dbutils.widgets.text("demo_file_directory", default_file_directory, "File Directory") +dbutils.widgets.text("demo_catalog_name", default_catalog, "Catalog Name") +dbutils.widgets.text("demo_schema_name", default_schema, "Schema Name") + +demo_file_directory = dbutils.widgets.get("demo_file_directory") +demo_catalog_name = dbutils.widgets.get("demo_catalog_name") +demo_schema_name = dbutils.widgets.get("demo_schema_name") + schema = "col1: int, col2: int, col3: int, col4 int" input_df = spark.createDataFrame([[1, 3, 3, 1], [2, None, 4, 1]], schema) @@ -75,12 +73,12 @@ # save generated checks in a workspace file user_name = spark.sql("select current_user() as user").collect()[0]["user"] -checks_file = f"/Workspace/Users/{user_name}/dqx_demo_checks.yml" +checks_file = f"{demo_file_directory}/dqx_demo_checks.yml" dq_engine = DQEngine(ws) dq_engine.save_checks_in_workspace_file(checks=checks, workspace_path=checks_file) # save generated checks in a Delta table -dq_engine.save_checks_in_table(checks=checks, table_name=f"{database}.{schema}.dqx_checks_table", mode="overwrite") +dq_engine.save_checks_in_table(checks=checks, table_name=f"{demo_catalog_name}.{demo_schema_name}.dqx_checks_table", mode="overwrite") # COMMAND ---------- @@ -121,7 +119,7 @@ # load checks from a Delta table dq_engine = DQEngine(WorkspaceClient()) -checks = dq_engine.load_checks_from_table(table_name=f"{database}.{schema}.dqx_checks_table") +checks = dq_engine.load_checks_from_table(table_name=f"{demo_catalog_name}.{demo_schema_name}.dqx_checks_table") # Option 1: apply quality rules and quarantine invalid records valid_df, quarantine_df = dq_engine.apply_checks_by_metadata_and_split(input_df, checks) @@ -286,24 +284,24 @@ import pyspark.sql.functions as F checks = [ - DQRowRule( # check for a single column + DQRowRule( # check for a single column name="col3_is_null_or_empty", criticality="warn", check_func=check_funcs.is_not_null_and_not_empty, column="col3" - ), - *DQForEachColRule( # check for multiple columns - columns=["col1", "col2"], - criticality="error", - check_func=check_funcs.is_not_null).get_rules(), - DQRowRule( # check with a filter + ), + *DQForEachColRule( # check for multiple columns + columns=["col1", "col2"], + criticality="error", + check_func=check_funcs.is_not_null).get_rules(), + DQRowRule( # check with a filter name="col_4_is_null_or_empty", criticality="warn", filter="col1 < 3", check_func=check_funcs.is_not_null_and_not_empty, column="col4" - ), - DQRowRule( + ), + DQRowRule( criticality="warn", check_func=check_funcs.is_not_null_and_not_empty, column='col3', @@ -311,51 +309,51 @@ "check_type": "completeness", "responsible_data_steward": "someone@email.com" } - ), - DQRowRule( # provide check func arguments using positional arguments - criticality="warn", - check_func=check_funcs.is_in_list, - column="col1", - check_func_args=[[1, 2]] - ), - DQRowRule( # provide check func arguments using keyword arguments - criticality="warn", - check_func=check_funcs.is_in_list, - column="col2", - check_func_kwargs={"allowed": [1, 2]} - ), - DQRowRule( # check for a struct field - # "error" criticality used if not provided - check_func=check_funcs.is_not_null, - column="col7.field1" - ), - DQRowRule( # check for a map element - criticality="error", - check_func=check_funcs.is_not_null, - column=F.try_element_at("col5", F.lit("key1")) - ), - DQRowRule( # check for an array element - criticality="error", - check_func=check_funcs.is_not_null, - column=F.try_element_at("col6", F.lit(1)) - ), - DQDatasetRule( # check uniqueness of composite key, multi-column rule - criticality="error", - check_func=check_funcs.is_unique, - columns=["col1", "col2"] - ), - DQDatasetRule( - criticality="error", - check_func=check_funcs.is_aggr_not_greater_than, - column="col1", - check_func_kwargs={"aggr_type": "count", "limit": 10}, - ), - DQDatasetRule( - criticality="error", - check_func=check_funcs.is_aggr_not_less_than, - column="col1", - check_func_kwargs={"aggr_type": "avg", "limit": 1.2}, - ), + ), + DQRowRule( # provide check func arguments using positional arguments + criticality="warn", + check_func=check_funcs.is_in_list, + column="col1", + check_func_args=[[1, 2]] + ), + DQRowRule( # provide check func arguments using keyword arguments + criticality="warn", + check_func=check_funcs.is_in_list, + column="col2", + check_func_kwargs={"allowed": [1, 2]} + ), + DQRowRule( # check for a struct field + # "error" criticality used if not provided + check_func=check_funcs.is_not_null, + column="col7.field1" + ), + DQRowRule( # check for a map element + criticality="error", + check_func=check_funcs.is_not_null, + column=F.try_element_at("col5", F.lit("key1")) + ), + DQRowRule( # check for an array element + criticality="error", + check_func=check_funcs.is_not_null, + column=F.try_element_at("col6", F.lit(1)) + ), + DQDatasetRule( # check uniqueness of composite key, multi-column rule + criticality="error", + check_func=check_funcs.is_unique, + columns=["col1", "col2"] + ), + DQDatasetRule( + criticality="error", + check_func=check_funcs.is_aggr_not_greater_than, + column="col1", + check_func_kwargs={"aggr_type": "count", "limit": 10}, + ), + DQDatasetRule( + criticality="error", + check_func=check_funcs.is_aggr_not_less_than, + column="col1", + check_func_kwargs={"aggr_type": "avg", "limit": 1.2}, + ), ] schema = "col1: int, col2: int, col3: int, col4 int, col5: map, col6: array, col7: struct" @@ -444,17 +442,17 @@ dq_engine.save_results_in_table( output_df=silver_df, quarantine_df=quarantine_df, - output_config=OutputConfig(f"{database}.{schema}.dqx_output", mode="overwrite"), - quarantine_config=OutputConfig(f"{database}.{schema}.dqx_quarantine", mode="overwrite") + output_config=OutputConfig(f"{demo_catalog_name}.{demo_schema_name}.dqx_output", mode="overwrite"), + quarantine_config=OutputConfig(f"{demo_catalog_name}.{demo_schema_name}.dqx_quarantine", mode="overwrite") ) # COMMAND ---------- -display(spark.table(f"{database}.{schema}.dqx_output")) +display(spark.table(f"{demo_catalog_name}.{demo_schema_name}.dqx_output")) # COMMAND ---------- -display(spark.table(f"{database}.{schema}.dqx_quarantine")) +display(spark.table(f"{demo_catalog_name}.{demo_schema_name}.dqx_quarantine")) # COMMAND ---------- @@ -467,13 +465,13 @@ dq_engine.apply_checks_by_metadata_and_save_in_table( input_config=InputConfig("/databricks-datasets/delta-sharing/samples/nyctaxi_2019"), checks=checks, - output_config=OutputConfig(f"{database}.{schema}.dqx_e2e_output", mode="overwrite"), - quarantine_config=OutputConfig(f"{database}.{schema}.dqx_e2e_quarantine", mode="overwrite") + output_config=OutputConfig(f"{demo_catalog_name}.{demo_schema_name}.dqx_e2e_output", mode="overwrite"), + quarantine_config=OutputConfig(f"{demo_catalog_name}.{demo_schema_name}.dqx_e2e_quarantine", mode="overwrite") ) # display the results saved to output and quarantine tables -display(spark.table(f"{database}.{schema}.dqx_e2e_output")) -display(spark.table(f"{database}.{schema}.dqx_e2e_quarantine")) +display(spark.table(f"{demo_catalog_name}.{demo_schema_name}.dqx_e2e_output")) +display(spark.table(f"{demo_catalog_name}.{demo_schema_name}.dqx_e2e_quarantine")) # COMMAND ---------- @@ -486,7 +484,6 @@ from databricks.labs.dqx.engine import DQEngine from databricks.sdk import WorkspaceClient - checks = [ DQDatasetRule( criticality="error", @@ -497,7 +494,7 @@ # either provide reference DataFrame name "ref_df_name": "ref_df_key", # or provide name of the reference table - #"ref_table": "catalog1.schema1.ref_table", + # "ref_table": "catalog1.schema1.ref_table", }, ), DQDatasetRule( @@ -531,35 +528,34 @@ from databricks.labs.dqx.engine import DQEngine from databricks.sdk import WorkspaceClient - checks = yaml.safe_load( - """ - - criticality: error - check: - function: foreign_key - arguments: - columns: - - col1 - ref_columns: - - ref_col1 - # either provide reference DataFrame name - ref_df_name: ref_df_key - # or provide name of the reference table - #ref_table: catalog1.schema1.ref_table - - - criticality: warn - name: foreign_key_check_on_composite_key - check: - function: foreign_key - arguments: - columns: - - col1 - - col2 - ref_columns: - - ref_col1 - - ref_col2 - ref_df_name: ref_df_key - """) + """ + - criticality: error + check: + function: foreign_key + arguments: + columns: + - col1 + ref_columns: + - ref_col1 + # either provide reference DataFrame name + ref_df_name: ref_df_key + # or provide name of the reference table + #ref_table: catalog1.schema1.ref_table + + - criticality: warn + name: foreign_key_check_on_composite_key + check: + function: foreign_key + arguments: + columns: + - col1 + - col2 + ref_columns: + - ref_col1 + - ref_col2 + ref_df_name: ref_df_key + """) input_df = spark.createDataFrame([[1, 1], [2, 2], [None, None]], "col1: int, col2: int") reference_df = spark.createDataFrame([[1, 1]], "ref_col1: int, ref_col2: int") @@ -594,7 +590,9 @@ @register_rule("row") def not_ends_with(column: str, suffix: str) -> Column: col_expr = F.col(column) - return make_condition(col_expr.endswith(suffix), f"Column {column} ends with {suffix}", f"{column}_ends_with_{suffix}") + return make_condition(col_expr.endswith(suffix), f"Column {column} ends with {suffix}", + f"{column}_ends_with_{suffix}") + # COMMAND ---------- @@ -608,15 +606,14 @@ def not_ends_with(column: str, suffix: str) -> Column: from databricks.labs.dqx.rule import DQRowRule from databricks.labs.dqx.check_funcs import is_not_null_and_not_empty, sql_expression - checks = [ # custom check DQRowRule(criticality="warn", check_func=not_ends_with, column="col1", check_func_kwargs={"suffix": "foo"}), # sql expression check DQRowRule(criticality="warn", check_func=sql_expression, check_func_kwargs={ - "expression": "col1 like 'str%'", "msg": "col1 not starting with 'str'" - } - ), + "expression": "col1 like 'str%'", "msg": "col1 not starting with 'str'" + } + ), # built-in check DQRowRule(criticality="error", check_func=is_not_null_and_not_empty, column="col1"), ] @@ -640,30 +637,29 @@ def not_ends_with(column: str, suffix: str) -> Column: from databricks.labs.dqx.engine import DQEngine from databricks.sdk import WorkspaceClient - checks = yaml.safe_load( -""" -# custom python check -- criticality: warn - check: - function: not_ends_with - arguments: - column: col1 - suffix: foo -# sql expression check -- criticality: warn - check: - function: sql_expression - arguments: - expression: col1 like 'str%' - msg: col1 not starting with 'str' -# built-in check -- criticality: error - check: - function: is_not_null_and_not_empty - arguments: - column: col1 -""" + """ + # custom python check + - criticality: warn + check: + function: not_ends_with + arguments: + column: col1 + suffix: foo + # sql expression check + - criticality: warn + check: + function: sql_expression + arguments: + expression: col1 like 'str%' + msg: col1 not starting with 'str' + # built-in check + - criticality: error + check: + function: is_not_null_and_not_empty + arguments: + column: col1 + """ ) schema = "col1: string, col2: string" @@ -673,7 +669,7 @@ def not_ends_with(column: str, suffix: str) -> Column: custom_check_functions = {"not_ends_with": not_ends_with} # alternatively, you can also use globals to include all available functions -#custom_check_functions = globals() +# custom_check_functions = globals() status = dq_engine.validate_checks(checks, custom_check_functions) assert not status.has_errors @@ -704,7 +700,6 @@ def not_ends_with(column: str, suffix: str) -> Column: [2, 2, 110] ], "measurement_id: int, sensor_id: int, reading_value: int") - # reference specs sensor_specs_df = spark.createDataFrame([ [1, 5], @@ -765,33 +760,33 @@ def not_ends_with(column: str, suffix: str) -> Column: # using YAML declarative approach checks = yaml.safe_load( - """ - - criticality: error - check: - function: sql_query - arguments: - merge_columns: - - sensor_id - condition_column: condition - msg: one of the sensor reading is greater than limit - name: sensor_reading_check - negate: false - input_placeholder: sensor - query: | - WITH joined AS ( - SELECT - sensor.*, - COALESCE(specs.min_threshold, 100) AS effective_threshold - FROM {{ sensor }} sensor - LEFT JOIN {{ sensor_specs }} specs - ON sensor.sensor_id = specs.sensor_id - ) - SELECT - sensor_id, - MAX(CASE WHEN reading_value > effective_threshold THEN 1 ELSE 0 END) = 1 AS condition - FROM joined - GROUP BY sensor_id - """ + """ + - criticality: error + check: + function: sql_query + arguments: + merge_columns: + - sensor_id + condition_column: condition + msg: one of the sensor reading is greater than limit + name: sensor_reading_check + negate: false + input_placeholder: sensor + query: | + WITH joined AS ( + SELECT + sensor.*, + COALESCE(specs.min_threshold, 100) AS effective_threshold + FROM {{ sensor }} sensor + LEFT JOIN {{ sensor_specs }} specs + ON sensor.sensor_id = specs.sensor_id + ) + SELECT + sensor_id, + MAX(CASE WHEN reading_value > effective_threshold THEN 1 ELSE 0 END) = 1 AS condition + FROM joined + GROUP BY sensor_id + """ ) ref_dfs = {"sensor_specs": sensor_specs_df} @@ -816,7 +811,6 @@ def not_ends_with(column: str, suffix: str) -> Column: @register_rule("dataset") # must be registered as dataset-level check def sensor_reading_less_than(default_limit: int) -> tuple[Column, Callable]: - # make sure any column added to the dataframe is unique condition_col = "condition" + uuid.uuid4().hex @@ -837,7 +831,8 @@ def apply(df: DataFrame, ref_dfs: dict[str, DataFrame]) -> DataFrame: # Join readings with specs sensor_df = df.join(sensor_specs_df, on="sensor_id", how="left") - sensor_df = sensor_df.withColumn("effective_threshold", F.coalesce(F.col("min_threshold"), F.lit(default_limit))) + sensor_df = sensor_df.withColumn("effective_threshold", + F.coalesce(F.col("min_threshold"), F.lit(default_limit))) # Check if any reading is below the spec-defined min_threshold per sensor aggr_df = ( @@ -878,23 +873,21 @@ def apply(df: DataFrame, ref_dfs: dict[str, DataFrame]) -> DataFrame: @register_rule("dataset") # must be registered as dataset-level check def sensor_reading_less_than2(default_limit: int) -> tuple[Column, Callable]: + # make sure any column added to the dataframe and registered temp views are unique + # in case the check / function is applied multiple times + unique_str = uuid.uuid4().hex + condition_col = "condition_" + unique_str - # make sure any column added to the dataframe and registered temp views are unique - # in case the check / function is applied multiple times - unique_str = uuid.uuid4().hex - condition_col = "condition_" + unique_str - - def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> DataFrame: + def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> DataFrame: + # Register the main and reference DataFrames as temporary views + sensor_view_unique = "sensor_" + unique_str + df.createOrReplaceTempView(sensor_view_unique) - # Register the main and reference DataFrames as temporary views - sensor_view_unique = "sensor_" + unique_str - df.createOrReplaceTempView(sensor_view_unique) + sensor_specs_view_unique = "sensor_specs_" + unique_str + ref_dfs["sensor_specs"].createOrReplaceTempView(sensor_specs_view_unique) - sensor_specs_view_unique = "sensor_specs_" + unique_str - ref_dfs["sensor_specs"].createOrReplaceTempView(sensor_specs_view_unique) - - # Perform the check - query = f""" + # Perform the check + query = f""" WITH joined AS ( SELECT sensor.*, @@ -920,16 +913,17 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> ON sensor.sensor_id = aggr.sensor_id """ - return spark.sql(query) + return spark.sql(query) + + return ( + make_condition( + condition=F.col(condition_col), # check if condition column is True + message=f"one of the sensor reading is greater than limit", + alias="sensor_reading_check", + ), + apply + ) - return ( - make_condition( - condition=F.col(condition_col), # check if condition column is True - message=f"one of the sensor reading is greater than limit", - alias="sensor_reading_check", - ), - apply - ) # COMMAND ---------- @@ -942,14 +936,13 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> from databricks.sdk import WorkspaceClient from databricks.labs.dqx.rule import DQDatasetRule - checks = [ - DQDatasetRule(criticality="error", check_func=sensor_reading_less_than, name="sensor_reading_exceeded", - check_func_kwargs={ - "default_limit": 100 - } - ), - # any other checks ... + DQDatasetRule(criticality="error", check_func=sensor_reading_less_than, name="sensor_reading_exceeded", + check_func_kwargs={ + "default_limit": 100 + } + ), + # any other checks ... ] # Pass reference DataFrame with sensor specifications @@ -982,7 +975,7 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> custom_check_functions = {"sensor_reading_less_than": sensor_reading_less_than} # list of custom check functions # or include all functions with globals() for simplicity -#custom_check_functions=globals() +# custom_check_functions=globals() # Pass reference DataFrame with sensor specifications ref_dfs = {"sensor_specs": sensor_specs_df} @@ -1005,7 +998,6 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> from databricks.labs.dqx.engine import DQEngine from databricks.sdk import WorkspaceClient - sensor_df.createOrReplaceTempView("sensor") sensor_specs_df.createOrReplaceTempView("sensor_specs") @@ -1068,11 +1060,10 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> from databricks.labs.dqx.engine import DQEngine from databricks.sdk import WorkspaceClient - schema = "id: int, id2: int, val: string" df = spark.createDataFrame( [ - + [1, 1, "val1"], [4, 5, "val1"], [6, 6, "val1"], @@ -1091,16 +1082,17 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> ) checks = [ - DQDatasetRule( + DQDatasetRule( criticality="error", check_func=check_funcs.compare_datasets, columns=["id", "id2"], # pass df.columns if you don't have primary key and want to match against all columns check_func_kwargs={ - "ref_columns": ["id", "id2"], # pass df2.columns if you don't have primary key and want to match against all columns - "ref_df_name": "ref_df", - "check_missing_records": True # if wanting to get info about missing records + "ref_columns": ["id", "id2"], + # pass df2.columns if you don't have primary key and want to match against all columns + "ref_df_name": "ref_df", + "check_missing_records": True # if wanting to get info about missing records }, - ) + ) ] ref_dfs = {"ref_df": df2} diff --git a/demos/dqx_demo_pii_detection.py b/demos/dqx_demo_pii_detection.py index df27081a5..280990e99 100644 --- a/demos/dqx_demo_pii_detection.py +++ b/demos/dqx_demo_pii_detection.py @@ -1,11 +1,11 @@ # Databricks notebook source # MAGIC %md # MAGIC # Using DQX for PII Detection -# MAGIC Increased regulation makes Databricks customers responsible for any Personally Identifiable Information (PII) stored in Unity Catalog. While [Lakehouse Monitoring](https://docs.databricks.com/aws/en/lakehouse-monitoring/data-classification#discover-sensitive-data) can identify sensitive data in-place, many customers need to proactively quarantine or anonymize PII before persisting the data. +# MAGIC Increased regulation makes Databricks customers responsible for any personally-identifying information (PII) stored in Unity Catalog. While [Lakehouse Monitoring](https://docs.databricks.com/aws/en/lakehouse-monitoring/data-classification#discover-sensitive-data) can identify sensitive data in-place, many customers need to proactively quarantine or anonymize PII before writing the data to Delta. # MAGIC # MAGIC [Databricks Labs' DQX project](https://databrickslabs.github.io/dqx/) provides in-flight data quality monitoring for Spark `DataFrames`. Customers can apply checks, get row-level metadata, and quarantine failing records. Workloads can use DQX's built-in checks or custom user-defined functions. # MAGIC -# MAGIC In this notebook, we'll use DQX with a custom check function to detect PII in JSON strings. +# MAGIC In this notebook, we'll use DQX with a custom function to detect PII in JSON strings. # COMMAND ---------- @@ -17,24 +17,25 @@ # COMMAND ---------- -# MAGIC %pip install databricks-labs-dqx==0.5.0 presidio-analyzer==2.2.358 numpy==1.26 --quiet +# MAGIC %pip install databricks-labs-dqx==0.4.0 presidio-analyzer numpy==1.23.5 --quiet # COMMAND ---------- -# MAGIC %restart_python +dbutils.library.restartPython() # COMMAND ---------- import json import pandas as pd +from typing import Iterator from pyspark.sql.functions import concat_ws, col, lit, pandas_udf from pyspark.sql import Column from presidio_analyzer import AnalyzerEngine from databricks.sdk import WorkspaceClient from databricks.labs.dqx.engine import DQEngine -from databricks.labs.dqx.rule import DQRowRule -from databricks.labs.dqx.check_funcs import make_condition +from databricks.labs.dqx.rule import DQColRule +from databricks.labs.dqx.row_checks import make_condition # COMMAND ---------- @@ -52,22 +53,30 @@ # Create a wrapper function to generate the entity mapping results: def get_entity_mapping(data: str) -> str | None: - # Run the Presidio analyzer to detect PII in the string: - results = analyzer.analyze( - text=data, - entities=["PERSON", "EMAIL_ADDRESS"], - language='en', - score_threshold=0.5, - ) - - # Validate and return the results: - results = [ - result.to_dict() - for result in results.entity_mapping() - if result.score >= 0.5 - ] - if results != []: - return json.dumps(results) + if data: + # Run the Presidio analyzer to detect PII in the string: + results = analyzer.analyze( + text=data, + entities=["PERSON", "EMAIL_ADDRESS"], + language='en', + score_threshold=0.5, + ) + if results != []: + output = [] + # Validate and return the results: + for result in results: + # Ignore if the result is low confidence: + if result.score < 0.6: + continue + # Append the result to the output: + output.append({ + "entity_type": result.entity_type, + "start": int(result.start), + "end": int(result.end), + "score": float(result.score), + }) + # Return the results as JSON: + return json.dumps(output) return None # COMMAND ---------- @@ -92,25 +101,25 @@ def contains_pii(batch: pd.Series) -> pd.Series: # COMMAND ---------- -def does_not_contain_pii(column: str) -> Column: +def does_not_contain_pii(col_name: str) -> Column: # Define a PII detection expression calling the pandas UDF: - pii_info = contains_pii(col(column)) + pii_info = contains_pii(col(col_name)) # Return the DQX condition that uses the PII detection expression: return make_condition( pii_info.isNotNull(), concat_ws( ' ', - lit(column), + lit(col_name), lit('contains pii with the following info:'), pii_info ), - f'{column}_contains_pii' + f'{col_name}_contains_pii' ) # Define the DQX rule: checks = [ - DQRowRule(criticality='error', check_func=does_not_contain_pii, column='val') + DQColRule(criticality='error', check_func=does_not_contain_pii, col_name='val') ] # Initialize the DQX engine: @@ -118,9 +127,9 @@ def does_not_contain_pii(column: str) -> Column: # Create some sample data: data = [ - ['{"key1": 1, "key2": "greg h", "key3": "222-32-1031"}'], - ['{"key1": 2, "key2": "Mickey mouse", "key3": "blue"}'], - ['{"key1": 3, "key2": "Prius", "key3": "Red"}'] + ['My name is John Smith'], + ['The sky is blue, road runner'], + ['Jane Smith sent an email to sara@info.com'] ] df = spark.createDataFrame(data, 'val string') diff --git a/demos/dqx_manufacturing_demo.py b/demos/dqx_manufacturing_demo.py index dc2df4d57..7d3f499e5 100644 --- a/demos/dqx_manufacturing_demo.py +++ b/demos/dqx_manufacturing_demo.py @@ -80,14 +80,14 @@ # COMMAND ---------- # DBTITLE 1,Set Catalog and Schema for Demo Dataset -default_database = "main" +default_database_name = "main" default_schema_name = "default" -dbutils.widgets.text("demo_database", default_database, "Catalog Name") -dbutils.widgets.text("demo_schema", default_schema_name, "Schema Name") +dbutils.widgets.text("demo_database_name", default_database_name, "Catalog Name") +dbutils.widgets.text("demo_schema_name", default_schema_name, "Schema Name") -database = dbutils.widgets.get("demo_database") -schema = dbutils.widgets.get("demo_schema") +database = dbutils.widgets.get("demo_database_name") +schema = dbutils.widgets.get("demo_schema_name") print(f"Selected Catalog for Demo Dataset: {database}") print(f"Selected Schema for Demo Dataset: {schema}") @@ -799,6 +799,8 @@ print(status) assert not status.has_errors +# COMMAND ---------- + # save checks in a workspace location sensor_dq_rules_yaml = f"{quality_rules_path}/sensor_dq_rules.yml" dq_engine.save_checks_in_workspace_file(sensor_dq_checks, workspace_path=sensor_dq_rules_yaml) @@ -862,10 +864,6 @@ def firmware_version_start_with_v(column: str) -> col: # COMMAND ---------- # DBTITLE 1,Add custom DQ rule in YAML -# Open existing Sensor DQ Rules YAML -with open(sensor_dq_rules_yaml, 'r') as f: - sensor_dq_checks = yaml.safe_load(f) - # Define Custom Check in YAML byor_quality_rule = { 'criticality': 'error', diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 7ab13c096..3a26b3327 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -15,17 +15,19 @@ RETRY_INTERVAL_SECONDS = 30 -def test_run_dqx_demo_library(make_notebook, make_catalog, make_schema, make_job): +def test_run_dqx_demo_library(make_notebook, make_schema, make_job): path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_library.py" ws = WorkspaceClient() with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) + directory = notebook.as_fuse().parent catalog = "main" schema = make_schema(catalog_name=catalog).name notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask( - notebook_path=notebook_path, base_parameters={"demo_database": catalog, "demo_schema": schema} + notebook_path=notebook_path, + base_parameters={"demo_database_name": catalog, "demo_schema_name": schema, "demo_file_directory": directory}, ) job = make_job(tasks=[Task(task_key="dqx_demo_library", notebook_task=notebook_task)]) run = ws.jobs.run_now(job.job_id) @@ -55,7 +57,7 @@ def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema, schema = make_schema(catalog_name=catalog).name notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask( - notebook_path=notebook_path, base_parameters={"demo_database": catalog, "demo_schema": schema} + notebook_path=notebook_path, base_parameters={"demo_database": catalog, "demo_schema": schema}, ) job = make_job(tasks=[Task(task_key="dqx_manufacturing_demo", notebook_task=notebook_task)]) run = ws.jobs.run_now(job.job_id) From 5a65502adbf339bd547c81a9c083845e7d7a96f0 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Fri, 25 Jul 2025 16:35:01 -0400 Subject: [PATCH 29/91] Add test for DLT pipeline demo --- tests/e2e/test_run_demos.py | 31 ++++++++++++++++++++++++++++--- 1 file changed, 28 insertions(+), 3 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 3a26b3327..5992e3e61 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -4,6 +4,7 @@ from pathlib import Path from databricks.sdk import WorkspaceClient from databricks.sdk.service.workspace import ImportFormat +from databricks.sdk.service.pipelines import UpdateInfoState from databricks.sdk.service.jobs import NotebookTask, Task, RunLifecycleStateV2State, TerminationTypeType @@ -20,7 +21,7 @@ def test_run_dqx_demo_library(make_notebook, make_schema, make_job): ws = WorkspaceClient() with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) - directory = notebook.as_fuse().parent + directory = notebook.as_fuse().parent.as_posix() catalog = "main" schema = make_schema(catalog_name=catalog).name @@ -57,7 +58,8 @@ def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema, schema = make_schema(catalog_name=catalog).name notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask( - notebook_path=notebook_path, base_parameters={"demo_database": catalog, "demo_schema": schema}, + notebook_path=notebook_path, + base_parameters={"demo_database": catalog, "demo_schema": schema}, ) job = make_job(tasks=[Task(task_key="dqx_manufacturing_demo", notebook_task=notebook_task)]) run = ws.jobs.run_now(job.job_id) @@ -108,7 +110,11 @@ def test_run_dqx_demo_pii_detection(make_notebook, make_cluster, make_job): notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask(notebook_path=notebook_path) cluster = make_cluster(single_node=True, spark_version="15.4.x-scala2.12") - job = make_job(tasks=[Task(task_key="dqx_demo_pii_detection", notebook_task=notebook_task, existing_cluster_id=cluster.cluster_id)]) + job = make_job( + tasks=[ + Task(task_key="dqx_demo_pii_detection", notebook_task=notebook_task, existing_cluster_id=cluster.cluster_id) + ] + ) run = ws.jobs.run_now(job.job_id) while True: @@ -122,3 +128,22 @@ def test_run_dqx_demo_pii_detection(make_notebook, make_cluster, make_job): assert ( termination_details.type == TerminationTypeType.SUCCESS ), f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" + + +def test_run_dqx_dlt_demo(make_notebook, make_pipeline): + path = Path(__file__).parent.parent.parent / "demos" / "dqx_dlt_demo.py" + ws = WorkspaceClient() + with open(path, "rb") as f: + notebook = make_notebook(content=f, format=ImportFormat.SOURCE) + + notebook_path = notebook.as_fuse().as_posix() + pipeline = make_pipeline(libraries=[notebook_path]) + update = ws.pipelines.start_update(pipeline.pipeline_id) + + while True: + update_details = ws.pipelines.get_update(pipeline_id=pipeline.pipeline_id, update_id=update.update_id) + if update_details.update.state in [UpdateInfoState.CANCELED, UpdateInfoState.COMPLETED, UpdateInfoState.FAILED]: + break + time.sleep(RETRY_INTERVAL_SECONDS) + + assert update_details.update.state == UpdateInfoState.COMPLETED, f"Run of pipeline '{pipeline.pipeline_id}' failed" From 73fd1040a233d9a1070b7141a462248266150e10 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Fri, 25 Jul 2025 16:57:47 -0400 Subject: [PATCH 30/91] Fix DLT pipeline test --- tests/e2e/test_run_demos.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 5992e3e61..b13609285 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -137,7 +137,7 @@ def test_run_dqx_dlt_demo(make_notebook, make_pipeline): notebook = make_notebook(content=f, format=ImportFormat.SOURCE) notebook_path = notebook.as_fuse().as_posix() - pipeline = make_pipeline(libraries=[notebook_path]) + pipeline = make_pipeline({"libraries": [{"notebook": {"path": notebook_path}}]}) update = ws.pipelines.start_update(pipeline.pipeline_id) while True: From dd5e5a1c0e9a9000a81cfeea39d5ffe26368d200 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Fri, 25 Jul 2025 17:08:42 -0400 Subject: [PATCH 31/91] Fix DLT pipeline test --- tests/e2e/test_run_demos.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index b13609285..5693dfd10 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -4,7 +4,7 @@ from pathlib import Path from databricks.sdk import WorkspaceClient from databricks.sdk.service.workspace import ImportFormat -from databricks.sdk.service.pipelines import UpdateInfoState +from databricks.sdk.service.pipelines import NotebookLibrary, PipelineLibrary, UpdateInfoState from databricks.sdk.service.jobs import NotebookTask, Task, RunLifecycleStateV2State, TerminationTypeType @@ -137,7 +137,7 @@ def test_run_dqx_dlt_demo(make_notebook, make_pipeline): notebook = make_notebook(content=f, format=ImportFormat.SOURCE) notebook_path = notebook.as_fuse().as_posix() - pipeline = make_pipeline({"libraries": [{"notebook": {"path": notebook_path}}]}) + pipeline = make_pipeline(libraries=[PipelineLibrary(notebook=NotebookLibrary(notebook_path))]) update = ws.pipelines.start_update(pipeline.pipeline_id) while True: From d340dd5a9d3670223fc32c3f9777d4cd30ab9666 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Fri, 25 Jul 2025 18:17:15 -0400 Subject: [PATCH 32/91] Add test for dqx_demo_tool --- tests/e2e/test_run_demos.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 5693dfd10..656e081f7 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -2,6 +2,7 @@ import time from pathlib import Path +from tests.integration.conftest import installation_ctx from databricks.sdk import WorkspaceClient from databricks.sdk.service.workspace import ImportFormat from databricks.sdk.service.pipelines import NotebookLibrary, PipelineLibrary, UpdateInfoState @@ -147,3 +148,34 @@ def test_run_dqx_dlt_demo(make_notebook, make_pipeline): time.sleep(RETRY_INTERVAL_SECONDS) assert update_details.update.state == UpdateInfoState.COMPLETED, f"Run of pipeline '{pipeline.pipeline_id}' failed" + + +def test_run_dqx_demo_tool(installation_ctx, make_notebook, make_job): + installation_ctx.prompts.extend({ + r'Provide location for the input data *': '/databricks-datasets/delta-sharing/samples/nyctaxi_2019', + r'Provide format for the input data (e.g. delta, parquet, csv, json) *': 'delta', + r'Provide output table in the format `catalog.schema.table` or `schema.table`': 'main.dqx.output_table', + r'Provide quarantined table in the format `catalog.schema.table` or `schema.table` ': 'main.dqx.quarantine_table', + }) + installation_ctx.workspace_installer.run(installation_ctx.config) + path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_tool.py" + ws = WorkspaceClient() + with open(path, "rb") as f: + notebook = make_notebook(content=f, format=ImportFormat.SOURCE) + + notebook_path = notebook.as_fuse().as_posix() + notebook_task = NotebookTask(notebook_path=notebook_path) + job = make_job(tasks=[Task(task_key="dqx_demo_tool", notebook_task=notebook_task)]) + run = ws.jobs.run_now(job.job_id) + + while True: + run_details = ws.jobs.get_run(run.run_id) + if run_details.status.state == RunLifecycleStateV2State.TERMINATED: + break + time.sleep(RETRY_INTERVAL_SECONDS) + + task = run_details.tasks[0] + termination_details = run_details.status.termination_details + assert ( + termination_details.type == TerminationTypeType.SUCCESS + ), f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" From ef0ed8d71c179517dd8ade2b9dd77607dd66e287 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sat, 26 Jul 2025 07:55:52 -0400 Subject: [PATCH 33/91] Handle whl folder as parameter --- demos/dqx_demo_tool.py | 8 ++++++-- tests/e2e/test_run_demos.py | 22 ++++++++++++++-------- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/demos/dqx_demo_tool.py b/demos/dqx_demo_tool.py index 7519513fe..687118898 100644 --- a/demos/dqx_demo_tool.py +++ b/demos/dqx_demo_tool.py @@ -52,7 +52,11 @@ import os user_name = spark.sql("select current_user() as user").collect()[0]["user"] -dqx_wheel_files = glob.glob(f"/Workspace/Users/{user_name}/.dqx/wheels/databricks_labs_dqx-*.whl") +default_dqx_wheel_files = glob.glob(f"/Workspace/Users/{user_name}/.dqx/wheels/databricks_labs_dqx-*.whl") + +dbutils.widgets.text("dqx_wheel_files", default_dqx_wheel_files, "DQX Wheel Files Folder") +dqx_wheel_files = dbutils.widgets.get("dqx_wheel_files") + dqx_latest_wheel = max(dqx_wheel_files, key=os.path.getctime) %pip install {dqx_latest_wheel} %restart_python @@ -149,7 +153,7 @@ function: sql_expression arguments: expression: pickup_datetime <= dropoff_datetime - msg: pickup time must not be greater than dropff time + msg: pickup time must not be greater than dropoff time name: pickup_datetime_greater_than_dropoff_datetime criticality: error - check: diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 656e081f7..8297e42d6 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -2,7 +2,6 @@ import time from pathlib import Path -from tests.integration.conftest import installation_ctx from databricks.sdk import WorkspaceClient from databricks.sdk.service.workspace import ImportFormat from databricks.sdk.service.pipelines import NotebookLibrary, PipelineLibrary, UpdateInfoState @@ -151,20 +150,27 @@ def test_run_dqx_dlt_demo(make_notebook, make_pipeline): def test_run_dqx_demo_tool(installation_ctx, make_notebook, make_job): - installation_ctx.prompts.extend({ - r'Provide location for the input data *': '/databricks-datasets/delta-sharing/samples/nyctaxi_2019', - r'Provide format for the input data (e.g. delta, parquet, csv, json) *': 'delta', - r'Provide output table in the format `catalog.schema.table` or `schema.table`': 'main.dqx.output_table', - r'Provide quarantined table in the format `catalog.schema.table` or `schema.table` ': 'main.dqx.quarantine_table', - }) + installation_ctx.prompts.extend( + { + r'Provide location for the input data *': '/databricks-datasets/delta-sharing/samples/nyctaxi_2019', + r'Provide format for the input data (e.g. delta, parquet, csv, json) *': 'delta', + r'Provide output table in the format `catalog.schema.table` or `schema.table`': 'main.dqx.output_table', + r'Provide quarantined table in the format `catalog.schema.table` or `schema.table` ': 'main.dqx.quarantine_table', + } + ) installation_ctx.workspace_installer.run(installation_ctx.config) + path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_tool.py" ws = WorkspaceClient() with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) + install_path = installation_ctx.installation.install_folder() notebook_path = notebook.as_fuse().as_posix() - notebook_task = NotebookTask(notebook_path=notebook_path) + notebook_task = NotebookTask( + notebook_path=notebook_path, + base_parameters={"dqx_wheel_files": f"{install_path}/wheels/databricks_labs_dqx-*.whl"}, + ) job = make_job(tasks=[Task(task_key="dqx_demo_tool", notebook_task=notebook_task)]) run = ws.jobs.run_now(job.job_id) From 2eb6510f895ccc7a892c1cf2c9cd12c81fa5eaf3 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sat, 26 Jul 2025 08:28:56 -0400 Subject: [PATCH 34/91] Fix installation test --- tests/e2e/test_run_demos.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 8297e42d6..decf127b0 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -1,6 +1,7 @@ import logging import time +from tests.integration.conftest import installation_ctx from pathlib import Path from databricks.sdk import WorkspaceClient from databricks.sdk.service.workspace import ImportFormat From f184f0493c3215015d8f130e7b0a38503b2d6c59 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sat, 26 Jul 2025 15:35:40 -0400 Subject: [PATCH 35/91] Fix demo tests --- demos/dqx_demo_tool.py | 8 +- tests/conftest.py | 161 ++++++++++++++++++++++++++++++++- tests/e2e/test_run_demos.py | 2 - tests/integration/conftest.py | 164 +--------------------------------- 4 files changed, 167 insertions(+), 168 deletions(-) diff --git a/demos/dqx_demo_tool.py b/demos/dqx_demo_tool.py index 687118898..82a3fffd4 100644 --- a/demos/dqx_demo_tool.py +++ b/demos/dqx_demo_tool.py @@ -53,9 +53,11 @@ user_name = spark.sql("select current_user() as user").collect()[0]["user"] default_dqx_wheel_files = glob.glob(f"/Workspace/Users/{user_name}/.dqx/wheels/databricks_labs_dqx-*.whl") +default_dqx_wheel_files_path = f"/Workspace/Users/{user_name}/.dqx/wheels/databricks_labs_dqx-*.whl" -dbutils.widgets.text("dqx_wheel_files", default_dqx_wheel_files, "DQX Wheel Files Folder") -dqx_wheel_files = dbutils.widgets.get("dqx_wheel_files") +dbutils.widgets.text("dqx_wheel_files_path", default_dqx_wheel_files_path, "DQX Wheel Files Folder") +dqx_wheel_files_path = dbutils.widgets.get("dqx_wheel_files_path") +dqx_wheel_files = glob.glob(dqx_wheel_files_path) dqx_latest_wheel = max(dqx_wheel_files, key=os.path.getctime) %pip install {dqx_latest_wheel} @@ -250,4 +252,4 @@ ctx = WorkspaceContext(WorkspaceClient()) dashboards_folder_link = f"{ctx.installation.workspace_link('')}dashboards/" print(f"Open a dashboard from the following folder and refresh it:") -print(dashboards_folder_link) \ No newline at end of file +print(dashboards_folder_link) diff --git a/tests/conftest.py b/tests/conftest.py index c2e3de5f2..658b6f5ba 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,10 +1,169 @@ import os -import pytest +from collections.abc import Callable, Generator +from dataclasses import replace +from functools import cached_property +import pytest +from databricks.labs.blueprint.installation import Installation, MockInstallation +from databricks.labs.blueprint.tui import MockPrompts +from databricks.labs.blueprint.wheels import ProductInfo +from databricks.labs.dqx.__about__ import __version__ +from databricks.labs.dqx.config import WorkspaceConfig, RunConfig +from databricks.labs.dqx.contexts.workflows import RuntimeContext +from databricks.labs.dqx.installer.install import WorkspaceInstaller, WorkspaceInstallation +from databricks.labs.dqx.installer.workflows_installer import WorkflowsDeployment +from databricks.labs.dqx.installer.workflow_task import Task +from databricks.labs.dqx.runtime import Workflows from databricks.labs.pytester.fixtures.baseline import factory +from databricks.sdk import WorkspaceClient from databricks.sdk.service.workspace import ImportFormat +@pytest.fixture +def debug_env_name(): + return "ws" # Specify the name of the debug environment from ~/.databricks/debug-env.json + + +@pytest.fixture +def product_info(): + return "dqx", __version__ + + +@pytest.fixture +def set_utc_timezone(): + """ + Set the timezone to UTC for the duration of the test to make sure spark timestamps are handled the same way regardless of the environment. + """ + os.environ["TZ"] = "UTC" + yield + os.environ.pop("TZ") + + +class CommonUtils: + def __init__(self, env_or_skip_fixture, ws): + self._env_or_skip = env_or_skip_fixture + self._ws = ws + + @cached_property + def installation(self): + return MockInstallation() + + @cached_property + def workspace_client(self) -> WorkspaceClient: + return self._ws + + +class MockRuntimeContext(CommonUtils, RuntimeContext): + def __init__(self, env_or_skip_fixture, ws_fixture) -> None: + super().__init__( + env_or_skip_fixture, + ws_fixture, + ) + self._env_or_skip = env_or_skip_fixture + + @cached_property + def config(self) -> WorkspaceConfig: + return WorkspaceConfig( + run_configs=[RunConfig()], + connect=self.workspace_client.config, + ) + + +class MockInstallationContext(MockRuntimeContext): + __test__ = False + + def __init__(self, env_or_skip_fixture, ws, check_file): + super().__init__(env_or_skip_fixture, ws) + self.check_file = check_file + + @cached_property + def installation(self): + return Installation(self.workspace_client, self.product_info.product_name()) + + @cached_property + def environ(self) -> dict[str, str]: + return {**os.environ} + + @cached_property + def workspace_installer(self): + return WorkspaceInstaller( + self.workspace_client, + self.environ, + ).replace(prompts=self.prompts, installation=self.installation, product_info=self.product_info) + + @cached_property + def config_transform(self) -> Callable[[WorkspaceConfig], WorkspaceConfig]: + return lambda wc: wc + + @cached_property + def config(self) -> WorkspaceConfig: + workspace_config = self.workspace_installer.configure() + + for i, run_config in enumerate(workspace_config.run_configs): + workspace_config.run_configs[i] = replace(run_config, checks_file=self.check_file) + + workspace_config = self.config_transform(workspace_config) + self.installation.save(workspace_config) + return workspace_config + + @cached_property + def product_info(self): + return ProductInfo.for_testing(WorkspaceConfig) + + @cached_property + def tasks(self) -> list[Task]: + return Workflows.all(self.config).tasks() + + @cached_property + def workflows_deployment(self) -> WorkflowsDeployment: + return WorkflowsDeployment( + self.config, + self.config.get_run_config().name, + self.installation, + self.install_state, + self.workspace_client, + self.product_info.wheels(self.workspace_client), + self.product_info, + self.tasks, + ) + + @cached_property + def prompts(self): + return MockPrompts( + { + r'Provide location for the input data (path or a table)': 'main.dqx_test.input_table', + r'Provide output table in the format `catalog.schema.table` or `schema.table`': 'main.dqx_test.output_table', + r'Do you want to uninstall DQX.*': 'yes', + r".*PRO or SERVERLESS SQL warehouse.*": "1", + r".*": "", + } + | (self.extend_prompts or {}) + ) + + @cached_property + def extend_prompts(self): + return {} + + @cached_property + def workspace_installation(self) -> WorkspaceInstallation: + return WorkspaceInstallation( + self.config, + self.installation, + self.install_state, + self.workspace_client, + self.workflows_deployment, + self.prompts, + self.product_info, + ) + + +@pytest.fixture +def installation_ctx(ws, env_or_skip, check_file="checks.yml") -> Generator[MockInstallationContext, None, None]: + ctx = MockInstallationContext(env_or_skip, ws, check_file) + yield ctx.replace(workspace_client=ws) + ctx.workspace_installation.uninstall() + + @pytest.fixture def checks_yaml_content(): return """- criticality: error diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index decf127b0..c0ee7973a 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -1,14 +1,12 @@ import logging import time -from tests.integration.conftest import installation_ctx from pathlib import Path from databricks.sdk import WorkspaceClient from databricks.sdk.service.workspace import ImportFormat from databricks.sdk.service.pipelines import NotebookLibrary, PipelineLibrary, UpdateInfoState from databricks.sdk.service.jobs import NotebookTask, Task, RunLifecycleStateV2State, TerminationTypeType - logging.getLogger("tests").setLevel("DEBUG") logging.getLogger("databricks.labs.dqx").setLevel("DEBUG") logger = logging.getLogger(__name__) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 1064b1197..105dbfcb5 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -1,23 +1,9 @@ -import os import logging -from collections.abc import Callable, Generator -from functools import cached_property -from dataclasses import replace from unittest.mock import patch import pytest -from databricks.labs.dqx.contexts.workflows import RuntimeContext -from databricks.labs.dqx.__about__ import __version__ -from databricks.sdk import WorkspaceClient -from databricks.labs.blueprint.wheels import ProductInfo -from databricks.labs.dqx.config import WorkspaceConfig, RunConfig, InputConfig, OutputConfig -from databricks.labs.blueprint.installation import Installation, MockInstallation -from databricks.labs.dqx.installer.install import WorkspaceInstaller, WorkspaceInstallation -from databricks.labs.blueprint.tui import MockPrompts -from databricks.labs.dqx.runtime import Workflows -from databricks.labs.dqx.installer.workflow_task import Task -from databricks.labs.dqx.installer.workflows_installer import WorkflowsDeployment +from databricks.labs.dqx.config import InputConfig, OutputConfig logging.getLogger("tests").setLevel("DEBUG") @@ -26,152 +12,6 @@ logger = logging.getLogger(__name__) -@pytest.fixture -def debug_env_name(): - return "ws" # Specify the name of the debug environment from ~/.databricks/debug-env.json - - -@pytest.fixture -def product_info(): - return "dqx", __version__ - - -@pytest.fixture -def set_utc_timezone(): - """ - Set the timezone to UTC for the duration of the test to make sure spark timestamps - are handled the same way regardless of the environment. - """ - os.environ["TZ"] = "UTC" - yield - os.environ.pop("TZ") - - -class CommonUtils: - def __init__(self, env_or_skip_fixture, ws): - self._env_or_skip = env_or_skip_fixture - self._ws = ws - - @cached_property - def installation(self): - return MockInstallation() - - @cached_property - def workspace_client(self) -> WorkspaceClient: - return self._ws - - -class MockRuntimeContext(CommonUtils, RuntimeContext): - def __init__(self, env_or_skip_fixture, ws_fixture) -> None: - super().__init__( - env_or_skip_fixture, - ws_fixture, - ) - self._env_or_skip = env_or_skip_fixture - - @cached_property - def config(self) -> WorkspaceConfig: - return WorkspaceConfig( - run_configs=[RunConfig()], - connect=self.workspace_client.config, - ) - - -class MockInstallationContext(MockRuntimeContext): - __test__ = False - - def __init__(self, env_or_skip_fixture, ws, check_file): - super().__init__(env_or_skip_fixture, ws) - self.check_file = check_file - - @cached_property - def installation(self): - return Installation(self.workspace_client, self.product_info.product_name()) - - @cached_property - def environ(self) -> dict[str, str]: - return {**os.environ} - - @cached_property - def workspace_installer(self): - return WorkspaceInstaller( - self.workspace_client, - self.environ, - ).replace(prompts=self.prompts, installation=self.installation, product_info=self.product_info) - - @cached_property - def config_transform(self) -> Callable[[WorkspaceConfig], WorkspaceConfig]: - return lambda wc: wc - - @cached_property - def config(self) -> WorkspaceConfig: - workspace_config = self.workspace_installer.configure() - - for i, run_config in enumerate(workspace_config.run_configs): - workspace_config.run_configs[i] = replace(run_config, checks_file=self.check_file) - - workspace_config = self.config_transform(workspace_config) - self.installation.save(workspace_config) - return workspace_config - - @cached_property - def product_info(self): - return ProductInfo.for_testing(WorkspaceConfig) - - @cached_property - def tasks(self) -> list[Task]: - return Workflows.all(self.config).tasks() - - @cached_property - def workflows_deployment(self) -> WorkflowsDeployment: - return WorkflowsDeployment( - self.config, - self.config.get_run_config().name, - self.installation, - self.install_state, - self.workspace_client, - self.product_info.wheels(self.workspace_client), - self.product_info, - self.tasks, - ) - - @cached_property - def prompts(self): - return MockPrompts( - { - r'Provide location for the input data (path or a table)': 'main.dqx_test.input_table', - r'Provide output table in the format `catalog.schema.table` or `schema.table`': 'main.dqx_test.output_table', - r'Do you want to uninstall DQX.*': 'yes', - r".*PRO or SERVERLESS SQL warehouse.*": "1", - r".*": "", - } - | (self.extend_prompts or {}) - ) - - @cached_property - def extend_prompts(self): - return {} - - @cached_property - def workspace_installation(self) -> WorkspaceInstallation: - return WorkspaceInstallation( - self.config, - self.installation, - self.install_state, - self.workspace_client, - self.workflows_deployment, - self.prompts, - self.product_info, - ) - - -@pytest.fixture -def installation_ctx(ws, env_or_skip, check_file="checks.yml") -> Generator[MockInstallationContext, None, None]: - ctx = MockInstallationContext(env_or_skip, ws, check_file) - yield ctx.replace(workspace_client=ws) - ctx.workspace_installation.uninstall() - - @pytest.fixture def webbrowser_open(): with patch("webbrowser.open") as mock_open: @@ -179,7 +19,7 @@ def webbrowser_open(): @pytest.fixture -def setup_workflows(installation_ctx: MockInstallationContext, make_schema, make_table): +def setup_workflows(installation_ctx, make_schema, make_table): """ Set up the workflows for the tests Existing cluster can be used by adding: From 75de17f7934d5ea8bd6259e870937d0e6835d4a5 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sat, 26 Jul 2025 15:42:55 -0400 Subject: [PATCH 36/91] Fix test --- tests/e2e/test_run_demos.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index c0ee7973a..c253eda1a 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -168,7 +168,7 @@ def test_run_dqx_demo_tool(installation_ctx, make_notebook, make_job): notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask( notebook_path=notebook_path, - base_parameters={"dqx_wheel_files": f"{install_path}/wheels/databricks_labs_dqx-*.whl"}, + base_parameters={"dqx_wheel_files_path": f"{install_path}/wheels/databricks_labs_dqx-*.whl"}, ) job = make_job(tasks=[Task(task_key="dqx_demo_tool", notebook_task=notebook_task)]) run = ws.jobs.run_now(job.job_id) From 67b6737a449edd7b42311934e751d8c6abb98c57 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sat, 26 Jul 2025 16:19:51 -0400 Subject: [PATCH 37/91] Test commit --- demos/dqx_demo_tool.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/demos/dqx_demo_tool.py b/demos/dqx_demo_tool.py index 82a3fffd4..29e17b3ec 100644 --- a/demos/dqx_demo_tool.py +++ b/demos/dqx_demo_tool.py @@ -52,14 +52,16 @@ import os user_name = spark.sql("select current_user() as user").collect()[0]["user"] -default_dqx_wheel_files = glob.glob(f"/Workspace/Users/{user_name}/.dqx/wheels/databricks_labs_dqx-*.whl") default_dqx_wheel_files_path = f"/Workspace/Users/{user_name}/.dqx/wheels/databricks_labs_dqx-*.whl" dbutils.widgets.text("dqx_wheel_files_path", default_dqx_wheel_files_path, "DQX Wheel Files Folder") dqx_wheel_files_path = dbutils.widgets.get("dqx_wheel_files_path") dqx_wheel_files = glob.glob(dqx_wheel_files_path) +try: + dqx_latest_wheel = max(dqx_wheel_files, key=os.path.getctime) +except: + raise ValueError(f"No files in path: {dqx_wheel_files_path}") -dqx_latest_wheel = max(dqx_wheel_files, key=os.path.getctime) %pip install {dqx_latest_wheel} %restart_python From e1139fd0d3788fb514955de78f5cb722b71543f9 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sat, 26 Jul 2025 16:47:57 -0400 Subject: [PATCH 38/91] Fix tool demo path parameter --- tests/e2e/test_run_demos.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index c253eda1a..639c04640 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -168,7 +168,7 @@ def test_run_dqx_demo_tool(installation_ctx, make_notebook, make_job): notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask( notebook_path=notebook_path, - base_parameters={"dqx_wheel_files_path": f"{install_path}/wheels/databricks_labs_dqx-*.whl"}, + base_parameters={"dqx_wheel_files_path": f"/Workspace{install_path}/wheels/databricks_labs_dqx-*.whl"}, ) job = make_job(tasks=[Task(task_key="dqx_demo_tool", notebook_task=notebook_task)]) run = ws.jobs.run_now(job.job_id) From 8fc496f7495697246dfead99a2c189bd4eb06b66 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sat, 26 Jul 2025 17:21:01 -0400 Subject: [PATCH 39/91] Add product info --- demos/dqx_demo_tool.py | 10 +++++++--- tests/e2e/test_run_demos.py | 6 +++++- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/demos/dqx_demo_tool.py b/demos/dqx_demo_tool.py index 29e17b3ec..6477c65ad 100644 --- a/demos/dqx_demo_tool.py +++ b/demos/dqx_demo_tool.py @@ -53,8 +53,12 @@ user_name = spark.sql("select current_user() as user").collect()[0]["user"] default_dqx_wheel_files_path = f"/Workspace/Users/{user_name}/.dqx/wheels/databricks_labs_dqx-*.whl" +default_dqx_product_name = "dqx" dbutils.widgets.text("dqx_wheel_files_path", default_dqx_wheel_files_path, "DQX Wheel Files Folder") +dbutils.widgets.text("dqx_product_name", default_product_name, "DQX Product Name") + +dqx_product_name = dbutils.widgets.get("product_name") dqx_wheel_files_path = dbutils.widgets.get("dqx_wheel_files_path") dqx_wheel_files = glob.glob(dqx_wheel_files_path) try: @@ -175,7 +179,7 @@ dq_engine = DQEngine(WorkspaceClient()) # save checks to location specified in the default run configuration inside workspace installation folder -dq_engine.save_checks_in_installation(checks, run_config_name="default") +dq_engine.save_checks_in_installation(checks, run_config_name="default", product_name=dqx_product_name) # or save it to an arbitrary workspace location #dq_engine.save_checks_in_workspace_file(checks, workspace_path="/Shared/App1/checks.yml") @@ -190,7 +194,7 @@ from databricks.labs.dqx.utils import read_input_data from databricks.sdk import WorkspaceClient -run_config = dq_engine.load_run_config(run_config_name="default", assume_user=True) +run_config = dq_engine.load_run_config(run_config_name="default", assume_user=True, product_name=dqx_product_name) # read the data, limit to 1000 rows for demo purpose bronze_df = read_input_data(spark, run_config.input_config).limit(1000) @@ -201,7 +205,7 @@ dq_engine = DQEngine(WorkspaceClient()) # load checks from location defined in the run configuration -checks = dq_engine.load_checks_from_installation(assume_user=True, run_config_name="default") +checks = dq_engine.load_checks_from_installation(assume_user=True, run_config_name="default", product_name=dqx_product_name) # or load checks from arbitrary workspace file # checks = dq_engine.load_checks_from_workspace_file(workspace_path="/Shared/App1/checks.yml") print(checks) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 639c04640..8a5202c14 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -158,6 +158,7 @@ def test_run_dqx_demo_tool(installation_ctx, make_notebook, make_job): } ) installation_ctx.workspace_installer.run(installation_ctx.config) + product_name = installation_ctx.product_info.product_name() path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_tool.py" ws = WorkspaceClient() @@ -168,7 +169,10 @@ def test_run_dqx_demo_tool(installation_ctx, make_notebook, make_job): notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask( notebook_path=notebook_path, - base_parameters={"dqx_wheel_files_path": f"/Workspace{install_path}/wheels/databricks_labs_dqx-*.whl"}, + base_parameters={ + "dqx_wheel_files_path": f"/Workspace{install_path}/wheels/databricks_labs_dqx-*.whl", + "dqx_product_name": product_name, + }, ) job = make_job(tasks=[Task(task_key="dqx_demo_tool", notebook_task=notebook_task)]) run = ws.jobs.run_now(job.job_id) From 3a190ff2452c1b80b0aebdd080ea42604f03b4f6 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sat, 26 Jul 2025 17:26:00 -0400 Subject: [PATCH 40/91] Add product info --- demos/dqx_demo_tool.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/demos/dqx_demo_tool.py b/demos/dqx_demo_tool.py index 6477c65ad..ed1e2450b 100644 --- a/demos/dqx_demo_tool.py +++ b/demos/dqx_demo_tool.py @@ -56,7 +56,7 @@ default_dqx_product_name = "dqx" dbutils.widgets.text("dqx_wheel_files_path", default_dqx_wheel_files_path, "DQX Wheel Files Folder") -dbutils.widgets.text("dqx_product_name", default_product_name, "DQX Product Name") +dbutils.widgets.text("dqx_product_name", default_dqx_product_name, "DQX Product Name") dqx_product_name = dbutils.widgets.get("product_name") dqx_wheel_files_path = dbutils.widgets.get("dqx_wheel_files_path") From 0c281b83b72890359e924e7c73284b9f17dc9b98 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sat, 26 Jul 2025 17:31:18 -0400 Subject: [PATCH 41/91] Add product info --- demos/dqx_demo_tool.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/demos/dqx_demo_tool.py b/demos/dqx_demo_tool.py index ed1e2450b..9b4397701 100644 --- a/demos/dqx_demo_tool.py +++ b/demos/dqx_demo_tool.py @@ -58,7 +58,7 @@ dbutils.widgets.text("dqx_wheel_files_path", default_dqx_wheel_files_path, "DQX Wheel Files Folder") dbutils.widgets.text("dqx_product_name", default_dqx_product_name, "DQX Product Name") -dqx_product_name = dbutils.widgets.get("product_name") +dqx_product_name = dbutils.widgets.get("dqx_product_name") dqx_wheel_files_path = dbutils.widgets.get("dqx_wheel_files_path") dqx_wheel_files = glob.glob(dqx_wheel_files_path) try: From eac0b40c3f89adbb37e3628eac6366294040fbf8 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sat, 26 Jul 2025 17:46:41 -0400 Subject: [PATCH 42/91] Add product info --- demos/dqx_demo_tool.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/demos/dqx_demo_tool.py b/demos/dqx_demo_tool.py index 9b4397701..75ebb22d4 100644 --- a/demos/dqx_demo_tool.py +++ b/demos/dqx_demo_tool.py @@ -100,7 +100,7 @@ ws = WorkspaceClient() dq_engine = DQEngine(ws) -run_config = dq_engine.load_run_config(run_config_name="default", assume_user=True) +run_config = dq_engine.load_run_config(run_config_name="default", assume_user=True, product_name=dqx_product_name) # read the input data, limit to 1000 rows for demo purpose input_df = read_input_data(spark, run_config.input_config).limit(1000) @@ -118,7 +118,7 @@ print(yaml.safe_dump(checks)) # save generated checks to location specified in the default run configuration inside workspace installation folder -dq_engine.save_checks_in_installation(checks, run_config_name="default") +dq_engine.save_checks_in_installation(checks, run_config_name="default", product_name=dqx_product_name) # or save it to an arbitrary workspace location #dq_engine.save_checks_in_workspace_file(checks, workspace_path="/Shared/App1/checks.yml") @@ -236,6 +236,7 @@ quarantine_df=quarantine_df, output_config=run_config.output_config, quarantine_config=run_config.quarantine_config, + product_name=dqx_product_name ) display(spark.sql(f"SELECT * FROM {run_config.output_config.location}")) From f4ab58d945c4d2c9169dba18d48f534647a27479 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sat, 26 Jul 2025 18:00:35 -0400 Subject: [PATCH 43/91] Add product info --- demos/dqx_demo_tool.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/demos/dqx_demo_tool.py b/demos/dqx_demo_tool.py index 75ebb22d4..2ade575a0 100644 --- a/demos/dqx_demo_tool.py +++ b/demos/dqx_demo_tool.py @@ -58,7 +58,6 @@ dbutils.widgets.text("dqx_wheel_files_path", default_dqx_wheel_files_path, "DQX Wheel Files Folder") dbutils.widgets.text("dqx_product_name", default_dqx_product_name, "DQX Product Name") -dqx_product_name = dbutils.widgets.get("dqx_product_name") dqx_wheel_files_path = dbutils.widgets.get("dqx_wheel_files_path") dqx_wheel_files = glob.glob(dqx_wheel_files_path) try: @@ -98,6 +97,9 @@ from databricks.labs.dqx.utils import read_input_data from databricks.sdk import WorkspaceClient +dqx_product_name = dbutils.widgets.get("dqx_product_name") + +# setup the DQEngine ws = WorkspaceClient() dq_engine = DQEngine(ws) run_config = dq_engine.load_run_config(run_config_name="default", assume_user=True, product_name=dqx_product_name) From 220d9ff56845a2e87ee156dead5af413c2c02c6d Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sat, 26 Jul 2025 18:26:03 -0400 Subject: [PATCH 44/91] Add product info --- demos/dqx_demo_tool.py | 8 +++++++- tests/e2e/test_run_demos.py | 8 +++++--- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/demos/dqx_demo_tool.py b/demos/dqx_demo_tool.py index 2ade575a0..1d9f5eb9b 100644 --- a/demos/dqx_demo_tool.py +++ b/demos/dqx_demo_tool.py @@ -228,7 +228,13 @@ # COMMAND ---------- -quarantine_catalog, quarantine_schema, _ = run_config.quarantine_config.location.split(".") +try: + quarantine_catalog, quarantine_schema, _ = run_config.quarantine_config.location.split(".") +except: + try: + raise ValueError(f"Quarantine config: {run_config.quarantine_config.__dict__}") + except: + raise ValueError(f"Run config: {run_config.__dict__}") spark.sql(f"CREATE CATALOG IF NOT EXISTS {quarantine_catalog}") spark.sql(f"CREATE SCHEMA IF NOT EXISTS {quarantine_catalog}.{quarantine_schema}") diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 8a5202c14..bacdd8d7f 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -148,13 +148,15 @@ def test_run_dqx_dlt_demo(make_notebook, make_pipeline): assert update_details.update.state == UpdateInfoState.COMPLETED, f"Run of pipeline '{pipeline.pipeline_id}' failed" -def test_run_dqx_demo_tool(installation_ctx, make_notebook, make_job): +def test_run_dqx_demo_tool(installation_ctx, make_schema, make_notebook, make_job): + catalog = "main" + schema = make_schema(catalog_name=catalog).name installation_ctx.prompts.extend( { r'Provide location for the input data *': '/databricks-datasets/delta-sharing/samples/nyctaxi_2019', r'Provide format for the input data (e.g. delta, parquet, csv, json) *': 'delta', - r'Provide output table in the format `catalog.schema.table` or `schema.table`': 'main.dqx.output_table', - r'Provide quarantined table in the format `catalog.schema.table` or `schema.table` ': 'main.dqx.quarantine_table', + r'Provide output table in the format `catalog.schema.table` or `schema.table`': f'{catalog}.{schema}.output_table', + r'Provide quarantined table in the format `catalog.schema.table` or `schema.table`': f'{catalog}.{schema}.quarantine_table', } ) installation_ctx.workspace_installer.run(installation_ctx.config) From 396f44d453f16e97d34c38c1d31c9841c3e20f0a Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sun, 27 Jul 2025 08:31:52 -0400 Subject: [PATCH 45/91] Update prompts extension --- tests/e2e/test_run_demos.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index bacdd8d7f..386f64d1d 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -151,13 +151,13 @@ def test_run_dqx_dlt_demo(make_notebook, make_pipeline): def test_run_dqx_demo_tool(installation_ctx, make_schema, make_notebook, make_job): catalog = "main" schema = make_schema(catalog_name=catalog).name - installation_ctx.prompts.extend( - { + installation_ctx.replace( + extend_prompts={ r'Provide location for the input data *': '/databricks-datasets/delta-sharing/samples/nyctaxi_2019', r'Provide format for the input data (e.g. delta, parquet, csv, json) *': 'delta', r'Provide output table in the format `catalog.schema.table` or `schema.table`': f'{catalog}.{schema}.output_table', r'Provide quarantined table in the format `catalog.schema.table` or `schema.table`': f'{catalog}.{schema}.quarantine_table', - } + }, ) installation_ctx.workspace_installer.run(installation_ctx.config) product_name = installation_ctx.product_info.product_name() From ebae8dc0864ea574002bd60ba5cd82001e9e6a41 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sun, 27 Jul 2025 09:19:18 -0400 Subject: [PATCH 46/91] Fix demo to run without catalog and schema creation --- demos/dqx_demo_tool.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/demos/dqx_demo_tool.py b/demos/dqx_demo_tool.py index 1d9f5eb9b..3fbfd82f4 100644 --- a/demos/dqx_demo_tool.py +++ b/demos/dqx_demo_tool.py @@ -228,17 +228,6 @@ # COMMAND ---------- -try: - quarantine_catalog, quarantine_schema, _ = run_config.quarantine_config.location.split(".") -except: - try: - raise ValueError(f"Quarantine config: {run_config.quarantine_config.__dict__}") - except: - raise ValueError(f"Run config: {run_config.__dict__}") - -spark.sql(f"CREATE CATALOG IF NOT EXISTS {quarantine_catalog}") -spark.sql(f"CREATE SCHEMA IF NOT EXISTS {quarantine_catalog}.{quarantine_schema}") - dq_engine.save_results_in_table( output_df=silver_df, quarantine_df=quarantine_df, From 4094d402b23c4c528a5e770c14b5962b255c8a9b Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sun, 27 Jul 2025 09:50:26 -0400 Subject: [PATCH 47/91] Update dashboard link --- demos/dqx_demo_tool.py | 11 ++++------- tests/e2e/test_run_demos.py | 2 +- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/demos/dqx_demo_tool.py b/demos/dqx_demo_tool.py index 3fbfd82f4..5abbb30b7 100644 --- a/demos/dqx_demo_tool.py +++ b/demos/dqx_demo_tool.py @@ -52,13 +52,13 @@ import os user_name = spark.sql("select current_user() as user").collect()[0]["user"] -default_dqx_wheel_files_path = f"/Workspace/Users/{user_name}/.dqx/wheels/databricks_labs_dqx-*.whl" +default_dqx_installation_path = f"/Workspace/Users/{user_name}/.dqx" default_dqx_product_name = "dqx" -dbutils.widgets.text("dqx_wheel_files_path", default_dqx_wheel_files_path, "DQX Wheel Files Folder") +dbutils.widgets.text("dqx_installation_path", default_dqx_installation_path, "DQX Installation Folder") dbutils.widgets.text("dqx_product_name", default_dqx_product_name, "DQX Product Name") -dqx_wheel_files_path = dbutils.widgets.get("dqx_wheel_files_path") +dqx_wheel_files_path = f"{dbutils.widgets.get('dqx_installation_path')}/wheels/databricks_labs_dqx-*.whl" dqx_wheel_files = glob.glob(dqx_wheel_files_path) try: dqx_latest_wheel = max(dqx_wheel_files, key=os.path.getctime) @@ -251,9 +251,6 @@ # COMMAND ---------- -from databricks.labs.dqx.contexts.workspace import WorkspaceContext - -ctx = WorkspaceContext(WorkspaceClient()) -dashboards_folder_link = f"{ctx.installation.workspace_link('')}dashboards/" +dashboards_folder_link = f"{dbutils.widgets.get('dqx_installation_path')}/dashboards/" print(f"Open a dashboard from the following folder and refresh it:") print(dashboards_folder_link) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 386f64d1d..82404c0c1 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -172,7 +172,7 @@ def test_run_dqx_demo_tool(installation_ctx, make_schema, make_notebook, make_jo notebook_task = NotebookTask( notebook_path=notebook_path, base_parameters={ - "dqx_wheel_files_path": f"/Workspace{install_path}/wheels/databricks_labs_dqx-*.whl", + "dqx_installation_path": f"/Workspace{install_path}", "dqx_product_name": product_name, }, ) From 87b46db791036b68ec7a6030be13e2c5a18e9c8a Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Sun, 27 Jul 2025 10:09:06 -0400 Subject: [PATCH 48/91] Update PR template --- .github/pull_request_template.md | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 75a9f2a48..ec30869d4 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -12,3 +12,4 @@ Resolves #.. - [ ] manually tested - [ ] added unit tests - [ ] added integration tests +- [ ] added end-to-end tests From d12b7b1331aa29bdefb0c06cee789cbc7a7d40aa Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 28 Jul 2025 16:01:10 -0400 Subject: [PATCH 49/91] Run demos from the open branch --- .github/workflows/acceptance.yml | 2 + demos/dqx_demo_library.py | 13 +++- demos/dqx_demo_pii_detection.py | 18 +++-- demos/dqx_manufacturing_demo.py | 11 ++- demos/dqx_quick_start_demo_library.py | 10 ++- tests/e2e/test_run_demos.py | 106 +++++++++++--------------- 6 files changed, 84 insertions(+), 76 deletions(-) diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index 8a9a051c1..246d45189 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -134,6 +134,7 @@ jobs: timeout: 2h codegen_path: tests/e2e/.codegen.json env: + REF_NAME: ${{ github.ref_name }} GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }} @@ -167,6 +168,7 @@ jobs: timeout: 2h codegen_path: tests/e2e/.codegen.json env: + REF_NAME: ${{ github.ref_name }} GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }} diff --git a/demos/dqx_demo_library.py b/demos/dqx_demo_library.py index 2b9ad7c22..3e8c592a4 100644 --- a/demos/dqx_demo_library.py +++ b/demos/dqx_demo_library.py @@ -9,8 +9,14 @@ # COMMAND ---------- -# MAGIC %pip install databricks-labs-dqx -# MAGIC %restart_python +dbutils.widgets.text("test_library_ref", "", "Test Library Ref") + +if dbutils.widgets.get("test_library_ref") != "": + %pip install '{dbutils.widgets.get("test_library_ref")}' +else: + %pip install databricks-labs-dqx + +%restart_python # COMMAND ---------- @@ -612,8 +618,7 @@ def not_ends_with(column: str, suffix: str) -> Column: # sql expression check DQRowRule(criticality="warn", check_func=sql_expression, check_func_kwargs={ "expression": "col1 like 'str%'", "msg": "col1 not starting with 'str'" - } - ), + }), # built-in check DQRowRule(criticality="error", check_func=is_not_null_and_not_empty, column="col1"), ] diff --git a/demos/dqx_demo_pii_detection.py b/demos/dqx_demo_pii_detection.py index 280990e99..679f8e366 100644 --- a/demos/dqx_demo_pii_detection.py +++ b/demos/dqx_demo_pii_detection.py @@ -1,7 +1,7 @@ # Databricks notebook source # MAGIC %md # MAGIC # Using DQX for PII Detection -# MAGIC Increased regulation makes Databricks customers responsible for any personally-identifying information (PII) stored in Unity Catalog. While [Lakehouse Monitoring](https://docs.databricks.com/aws/en/lakehouse-monitoring/data-classification#discover-sensitive-data) can identify sensitive data in-place, many customers need to proactively quarantine or anonymize PII before writing the data to Delta. +# MAGIC Increased regulation makes Databricks customers responsible for any Personally Identifiable Information (PII) stored in Unity Catalog. While [Lakehouse Monitoring](https://docs.databricks.com/aws/en/lakehouse-monitoring/data-classification#discover-sensitive-data) can identify sensitive data in-place, many customers need to proactively quarantine or anonymize PII before writing the data to Delta. # MAGIC # MAGIC [Databricks Labs' DQX project](https://databrickslabs.github.io/dqx/) provides in-flight data quality monitoring for Spark `DataFrames`. Customers can apply checks, get row-level metadata, and quarantine failing records. Workloads can use DQX's built-in checks or custom user-defined functions. # MAGIC @@ -17,7 +17,16 @@ # COMMAND ---------- -# MAGIC %pip install databricks-labs-dqx==0.4.0 presidio-analyzer numpy==1.23.5 --quiet +dbutils.widgets.text("test_library_ref", "", "Test Library Ref") + +if dbutils.widgets.get("test_library_ref") != "": + %pip install '{dbutils.widgets.get("test_library_ref")}' +else: + %pip install databricks-labs-dqx + +%pip install presidio_analyzer numpy==1.23.5 + +%restart_python # COMMAND ---------- @@ -28,13 +37,12 @@ import json import pandas as pd -from typing import Iterator from pyspark.sql.functions import concat_ws, col, lit, pandas_udf from pyspark.sql import Column from presidio_analyzer import AnalyzerEngine from databricks.sdk import WorkspaceClient from databricks.labs.dqx.engine import DQEngine -from databricks.labs.dqx.rule import DQColRule +from databricks.labs.dqx.rule import DQRowRule from databricks.labs.dqx.row_checks import make_condition # COMMAND ---------- @@ -119,7 +127,7 @@ def does_not_contain_pii(col_name: str) -> Column: # Define the DQX rule: checks = [ - DQColRule(criticality='error', check_func=does_not_contain_pii, col_name='val') + DQRowRule(criticality='error', check_func=does_not_contain_pii, column='val') ] # Initialize the DQX engine: diff --git a/demos/dqx_manufacturing_demo.py b/demos/dqx_manufacturing_demo.py index 7d3f499e5..55cb3d8b2 100644 --- a/demos/dqx_manufacturing_demo.py +++ b/demos/dqx_manufacturing_demo.py @@ -52,8 +52,15 @@ # COMMAND ---------- # DBTITLE 1,Install DQX Library -# MAGIC %pip install databricks-labs-dqx==0.5.0 -# MAGIC %restart_python + +dbutils.widgets.text("test_library_ref", "", "Test Library Ref") + +if dbutils.widgets.get("test_library_ref") != "": + %pip install '{dbutils.widgets.get("test_library_ref")}' +else: + %pip install databricks-labs-dqx + +%restart_python # COMMAND ---------- diff --git a/demos/dqx_quick_start_demo_library.py b/demos/dqx_quick_start_demo_library.py index f2ff6e91e..f16e5f5d9 100644 --- a/demos/dqx_quick_start_demo_library.py +++ b/demos/dqx_quick_start_demo_library.py @@ -14,8 +14,14 @@ # COMMAND ---------- -# MAGIC %pip install databricks-labs-dqx -# MAGIC %restart_python +dbutils.widgets.text("test_library_ref", "", "Test Library Ref") + +if dbutils.widgets.get("test_library_ref") != "": + %pip install '{dbutils.widgets.get("test_library_ref")}' +else: + %pip install databricks-labs-dqx + +%restart_python # COMMAND ---------- diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 82404c0c1..d58780380 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -1,11 +1,13 @@ import logging +import os import time +from datetime import timedelta from pathlib import Path from databricks.sdk import WorkspaceClient from databricks.sdk.service.workspace import ImportFormat from databricks.sdk.service.pipelines import NotebookLibrary, PipelineLibrary, UpdateInfoState -from databricks.sdk.service.jobs import NotebookTask, Task, RunLifecycleStateV2State, TerminationTypeType +from databricks.sdk.service.jobs import NotebookTask, Run, Task, TerminationTypeType logging.getLogger("tests").setLevel("DEBUG") logging.getLogger("databricks.labs.dqx").setLevel("DEBUG") @@ -13,6 +15,7 @@ RETRY_INTERVAL_SECONDS = 30 +TEST_LIBRARY_REF = f"git+https://github.com/databrickslabs/dqx.git@{os.getenv('REF_NAME')}" def test_run_dqx_demo_library(make_notebook, make_schema, make_job): @@ -27,22 +30,18 @@ def test_run_dqx_demo_library(make_notebook, make_schema, make_job): notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask( notebook_path=notebook_path, - base_parameters={"demo_database_name": catalog, "demo_schema_name": schema, "demo_file_directory": directory}, + base_parameters={ + "demo_database_name": catalog, + "demo_schema_name": schema, + "demo_file_directory": directory, + "test_library_ref": TEST_LIBRARY_REF, + }, ) job = make_job(tasks=[Task(task_key="dqx_demo_library", notebook_task=notebook_task)]) - run = ws.jobs.run_now(job.job_id) - while True: - run_details = ws.jobs.get_run(run.run_id) - if run_details.status.state == RunLifecycleStateV2State.TERMINATED: - break - time.sleep(RETRY_INTERVAL_SECONDS) - - task = run_details.tasks[0] - termination_details = run_details.status.termination_details - assert ( - termination_details.type == TerminationTypeType.SUCCESS - ), f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" + waiter = ws.jobs.run_now(job.job_id) + run = waiter.result(timeout=timedelta(minutes=30), callback=lambda r: validate_demo_run_status(r, ws)) + logging.info(f"Job run {run.run_id} completed successfully for dqx_manufacturing_demo") def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema, make_job): @@ -58,22 +57,13 @@ def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema, notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask( notebook_path=notebook_path, - base_parameters={"demo_database": catalog, "demo_schema": schema}, + base_parameters={"demo_database": catalog, "demo_schema": schema, "test_library_ref": TEST_LIBRARY_REF}, ) job = make_job(tasks=[Task(task_key="dqx_manufacturing_demo", notebook_task=notebook_task)]) - run = ws.jobs.run_now(job.job_id) - while True: - run_details = ws.jobs.get_run(run.run_id) - if run_details.status.state == RunLifecycleStateV2State.TERMINATED: - break - time.sleep(RETRY_INTERVAL_SECONDS) - - task = run_details.tasks[0] - termination_details = run_details.status.termination_details - assert ( - termination_details.type == TerminationTypeType.SUCCESS - ), f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" + waiter = ws.jobs.run_now(job.job_id) + run = waiter.result(timeout=timedelta(minutes=30), callback=lambda r: validate_demo_run_status(r, ws)) + logging.info(f"Job run {run.run_id} completed successfully for dqx_manufacturing_demo") def test_run_dqx_quick_start_demo_library(make_notebook, make_job): @@ -83,21 +73,12 @@ def test_run_dqx_quick_start_demo_library(make_notebook, make_job): notebook = make_notebook(content=f, format=ImportFormat.SOURCE) notebook_path = notebook.as_fuse().as_posix() - notebook_task = NotebookTask(notebook_path=notebook_path) + notebook_task = NotebookTask(notebook_path=notebook_path, base_parameters={"test_library_ref": TEST_LIBRARY_REF}) job = make_job(tasks=[Task(task_key="dqx_quick_start_demo_library", notebook_task=notebook_task)]) - run = ws.jobs.run_now(job.job_id) - - while True: - run_details = ws.jobs.get_run(run.run_id) - if run_details.status.state == RunLifecycleStateV2State.TERMINATED: - break - time.sleep(RETRY_INTERVAL_SECONDS) - task = run_details.tasks[0] - termination_details = run_details.status.termination_details - assert ( - termination_details.type == TerminationTypeType.SUCCESS - ), f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" + waiter = ws.jobs.run_now(job.job_id) + run = waiter.result(timeout=timedelta(minutes=30), callback=lambda r: validate_demo_run_status(r, ws)) + logging.info(f"Job run {run.run_id} completed successfully for dqx_quick_start_demo_library") def test_run_dqx_demo_pii_detection(make_notebook, make_cluster, make_job): @@ -107,26 +88,17 @@ def test_run_dqx_demo_pii_detection(make_notebook, make_cluster, make_job): notebook = make_notebook(content=f, format=ImportFormat.SOURCE) notebook_path = notebook.as_fuse().as_posix() - notebook_task = NotebookTask(notebook_path=notebook_path) + notebook_task = NotebookTask(notebook_path=notebook_path, base_parameters={"test_library_ref": TEST_LIBRARY_REF}) cluster = make_cluster(single_node=True, spark_version="15.4.x-scala2.12") job = make_job( tasks=[ Task(task_key="dqx_demo_pii_detection", notebook_task=notebook_task, existing_cluster_id=cluster.cluster_id) ] ) - run = ws.jobs.run_now(job.job_id) - while True: - run_details = ws.jobs.get_run(run.run_id) - if run_details.status.state == RunLifecycleStateV2State.TERMINATED: - break - time.sleep(RETRY_INTERVAL_SECONDS) - - task = run_details.tasks[0] - termination_details = run_details.status.termination_details - assert ( - termination_details.type == TerminationTypeType.SUCCESS - ), f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" + waiter = ws.jobs.run_now(job.job_id) + run = waiter.result(timeout=timedelta(minutes=30), callback=lambda r: validate_demo_run_status(r, ws)) + logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_pii_detection") def test_run_dqx_dlt_demo(make_notebook, make_pipeline): @@ -161,13 +133,13 @@ def test_run_dqx_demo_tool(installation_ctx, make_schema, make_notebook, make_jo ) installation_ctx.workspace_installer.run(installation_ctx.config) product_name = installation_ctx.product_info.product_name() + install_path = installation_ctx.installation.install_folder() path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_tool.py" ws = WorkspaceClient() with open(path, "rb") as f: notebook = make_notebook(content=f, format=ImportFormat.SOURCE) - install_path = installation_ctx.installation.install_folder() notebook_path = notebook.as_fuse().as_posix() notebook_task = NotebookTask( notebook_path=notebook_path, @@ -177,16 +149,24 @@ def test_run_dqx_demo_tool(installation_ctx, make_schema, make_notebook, make_jo }, ) job = make_job(tasks=[Task(task_key="dqx_demo_tool", notebook_task=notebook_task)]) - run = ws.jobs.run_now(job.job_id) - while True: - run_details = ws.jobs.get_run(run.run_id) - if run_details.status.state == RunLifecycleStateV2State.TERMINATED: - break - time.sleep(RETRY_INTERVAL_SECONDS) + waiter = ws.jobs.run_now(job.job_id) + run = waiter.result(timeout=timedelta(minutes=30), callback=lambda r: validate_demo_run_status(r, ws)) + logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_tool") + + +def validate_demo_run_status(run: Run, client: WorkspaceClient | None) -> None: + """ + Validates that a demo run completed successfully. + :param run: `Run` object returned from a `WorkspaceClient.jobs.submit(...)` command + :param client: `WorkspaceClient` used for getting the run-level errors + """ + if not client: + client = WorkspaceClient() + + task = run.tasks[0] + termination_details = run.status.termination_details - task = run_details.tasks[0] - termination_details = run_details.status.termination_details assert ( termination_details.type == TerminationTypeType.SUCCESS - ), f"Run of '{task.task_key}' failed with message: {ws.jobs.get_run_output(task.run_id).error}" + ), f"Run of '{task.task_key}' failed with message: {client.jobs.get_run_output(task.run_id).error}" From 639eb105219dbd69b432e1ecefeda342329a9b47 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 28 Jul 2025 17:42:48 -0400 Subject: [PATCH 50/91] Fix tests --- tests/e2e/test_run_demos.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index d58780380..84f787aec 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -151,19 +151,16 @@ def test_run_dqx_demo_tool(installation_ctx, make_schema, make_notebook, make_jo job = make_job(tasks=[Task(task_key="dqx_demo_tool", notebook_task=notebook_task)]) waiter = ws.jobs.run_now(job.job_id) - run = waiter.result(timeout=timedelta(minutes=30), callback=lambda r: validate_demo_run_status(r, ws)) + run = waiter.result(timeout=timedelta(minutes=30), callback=validate_demo_run_status) logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_tool") -def validate_demo_run_status(run: Run, client: WorkspaceClient | None) -> None: +def validate_demo_run_status(run: Run) -> None: """ Validates that a demo run completed successfully. :param run: `Run` object returned from a `WorkspaceClient.jobs.submit(...)` command - :param client: `WorkspaceClient` used for getting the run-level errors """ - if not client: - client = WorkspaceClient() - + client = WorkspaceClient() task = run.tasks[0] termination_details = run.status.termination_details From 1cd8b25951e67798335af9cd3f25d377f970875e Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 28 Jul 2025 17:47:44 -0400 Subject: [PATCH 51/91] Fix tests --- tests/e2e/test_run_demos.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 84f787aec..290ceaabf 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -40,7 +40,7 @@ def test_run_dqx_demo_library(make_notebook, make_schema, make_job): job = make_job(tasks=[Task(task_key="dqx_demo_library", notebook_task=notebook_task)]) waiter = ws.jobs.run_now(job.job_id) - run = waiter.result(timeout=timedelta(minutes=30), callback=lambda r: validate_demo_run_status(r, ws)) + run = waiter.result(timeout=timedelta(minutes=30), callback=validate_demo_run_status) logging.info(f"Job run {run.run_id} completed successfully for dqx_manufacturing_demo") @@ -62,7 +62,7 @@ def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema, job = make_job(tasks=[Task(task_key="dqx_manufacturing_demo", notebook_task=notebook_task)]) waiter = ws.jobs.run_now(job.job_id) - run = waiter.result(timeout=timedelta(minutes=30), callback=lambda r: validate_demo_run_status(r, ws)) + run = waiter.result(timeout=timedelta(minutes=30), callback=validate_demo_run_status) logging.info(f"Job run {run.run_id} completed successfully for dqx_manufacturing_demo") @@ -77,7 +77,7 @@ def test_run_dqx_quick_start_demo_library(make_notebook, make_job): job = make_job(tasks=[Task(task_key="dqx_quick_start_demo_library", notebook_task=notebook_task)]) waiter = ws.jobs.run_now(job.job_id) - run = waiter.result(timeout=timedelta(minutes=30), callback=lambda r: validate_demo_run_status(r, ws)) + run = waiter.result(timeout=timedelta(minutes=30), callback=validate_demo_run_status) logging.info(f"Job run {run.run_id} completed successfully for dqx_quick_start_demo_library") @@ -97,7 +97,7 @@ def test_run_dqx_demo_pii_detection(make_notebook, make_cluster, make_job): ) waiter = ws.jobs.run_now(job.job_id) - run = waiter.result(timeout=timedelta(minutes=30), callback=lambda r: validate_demo_run_status(r, ws)) + run = waiter.result(timeout=timedelta(minutes=30), callback=validate_demo_run_status) logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_pii_detection") From 56e39c91a6f3b819082304cfdeb6eda3542cbf86 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 28 Jul 2025 18:07:49 -0400 Subject: [PATCH 52/91] Fix tests --- tests/e2e/test_run_demos.py | 49 +++++++++++++++++++++++++++---------- 1 file changed, 36 insertions(+), 13 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 290ceaabf..2945c4c11 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -39,9 +39,13 @@ def test_run_dqx_demo_library(make_notebook, make_schema, make_job): ) job = make_job(tasks=[Task(task_key="dqx_demo_library", notebook_task=notebook_task)]) - waiter = ws.jobs.run_now(job.job_id) - run = waiter.result(timeout=timedelta(minutes=30), callback=validate_demo_run_status) - logging.info(f"Job run {run.run_id} completed successfully for dqx_manufacturing_demo") + waiter = ws.jobs.run_now_and_wait(job.job_id) + run = ws.jobs.wait_get_run_job_terminated_or_skipped( + run_id=waiter.run_id, + timeout=timedelta(minutes=30), + callback=lambda r: validate_demo_run_status(r, client=ws), + ) + logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_library") def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema, make_job): @@ -61,8 +65,12 @@ def test_run_dqx_manufacturing_demo(make_notebook, make_directory, make_schema, ) job = make_job(tasks=[Task(task_key="dqx_manufacturing_demo", notebook_task=notebook_task)]) - waiter = ws.jobs.run_now(job.job_id) - run = waiter.result(timeout=timedelta(minutes=30), callback=validate_demo_run_status) + waiter = ws.jobs.run_now_and_wait(job.job_id) + run = ws.jobs.wait_get_run_job_terminated_or_skipped( + run_id=waiter.run_id, + timeout=timedelta(minutes=30), + callback=lambda r: validate_demo_run_status(r, client=ws), + ) logging.info(f"Job run {run.run_id} completed successfully for dqx_manufacturing_demo") @@ -76,8 +84,12 @@ def test_run_dqx_quick_start_demo_library(make_notebook, make_job): notebook_task = NotebookTask(notebook_path=notebook_path, base_parameters={"test_library_ref": TEST_LIBRARY_REF}) job = make_job(tasks=[Task(task_key="dqx_quick_start_demo_library", notebook_task=notebook_task)]) - waiter = ws.jobs.run_now(job.job_id) - run = waiter.result(timeout=timedelta(minutes=30), callback=validate_demo_run_status) + waiter = ws.jobs.run_now_and_wait(job.job_id) + run = ws.jobs.wait_get_run_job_terminated_or_skipped( + run_id=waiter.run_id, + timeout=timedelta(minutes=30), + callback=lambda r: validate_demo_run_status(r, client=ws), + ) logging.info(f"Job run {run.run_id} completed successfully for dqx_quick_start_demo_library") @@ -96,8 +108,12 @@ def test_run_dqx_demo_pii_detection(make_notebook, make_cluster, make_job): ] ) - waiter = ws.jobs.run_now(job.job_id) - run = waiter.result(timeout=timedelta(minutes=30), callback=validate_demo_run_status) + waiter = ws.jobs.run_now_and_wait(job.job_id) + run = ws.jobs.wait_get_run_job_terminated_or_skipped( + run_id=waiter.run_id, + timeout=timedelta(minutes=30), + callback=lambda r: validate_demo_run_status(r, client=ws), + ) logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_pii_detection") @@ -150,17 +166,24 @@ def test_run_dqx_demo_tool(installation_ctx, make_schema, make_notebook, make_jo ) job = make_job(tasks=[Task(task_key="dqx_demo_tool", notebook_task=notebook_task)]) - waiter = ws.jobs.run_now(job.job_id) - run = waiter.result(timeout=timedelta(minutes=30), callback=validate_demo_run_status) + waiter = ws.jobs.run_now_and_wait(job.job_id) + run = ws.jobs.wait_get_run_job_terminated_or_skipped( + run_id=waiter.run_id, + timeout=timedelta(minutes=30), + callback=lambda r: validate_demo_run_status(r, client=ws), + ) logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_tool") -def validate_demo_run_status(run: Run) -> None: +def validate_demo_run_status(run: Run, client: WorkspaceClient | None) -> None: """ Validates that a demo run completed successfully. :param run: `Run` object returned from a `WorkspaceClient.jobs.submit(...)` command + :param client: Optional `WorkspaceClient` object for getting task output """ - client = WorkspaceClient() + if client is None: + client = WorkspaceClient() + task = run.tasks[0] termination_details = run.status.termination_details From 22269ae10618393241d74a285ca5d1dc0bd1e265 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 28 Jul 2025 18:20:43 -0400 Subject: [PATCH 53/91] Fix tests --- .github/workflows/acceptance.yml | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index 246d45189..ae459b996 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -127,6 +127,12 @@ jobs: - name: Install hatch run: pip install hatch==1.9.4 + - name: Print env var + run: echo Github ref = $REF_NAME + env: + REF_NAME: ${{ github.ref_name }} + + - name: Run e2e tests uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.4 with: @@ -158,6 +164,11 @@ jobs: cache-dependency-path: '**/pyproject.toml' python-version: '3.10' + - name: Print env var + run: echo Github ref = $REF_NAME + env: + REF_NAME: ${{ github.ref_name }} + - name: Install hatch run: pip install hatch==1.9.4 From 778543d78165028e3fe63503cf4c6efebe20f0d9 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 28 Jul 2025 18:29:56 -0400 Subject: [PATCH 54/91] Fix tests --- .github/workflows/acceptance.yml | 11 ----------- tests/e2e/test_run_demos.py | 4 ++-- 2 files changed, 2 insertions(+), 13 deletions(-) diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index ae459b996..246d45189 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -127,12 +127,6 @@ jobs: - name: Install hatch run: pip install hatch==1.9.4 - - name: Print env var - run: echo Github ref = $REF_NAME - env: - REF_NAME: ${{ github.ref_name }} - - - name: Run e2e tests uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.4 with: @@ -164,11 +158,6 @@ jobs: cache-dependency-path: '**/pyproject.toml' python-version: '3.10' - - name: Print env var - run: echo Github ref = $REF_NAME - env: - REF_NAME: ${{ github.ref_name }} - - name: Install hatch run: pip install hatch==1.9.4 diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 2945c4c11..077e674da 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -13,9 +13,9 @@ logging.getLogger("databricks.labs.dqx").setLevel("DEBUG") logger = logging.getLogger(__name__) - RETRY_INTERVAL_SECONDS = 30 -TEST_LIBRARY_REF = f"git+https://github.com/databrickslabs/dqx.git@{os.getenv('REF_NAME')}" +TEST_LIBRARY_REF = f"git+https://github.com/databrickslabs/dqx.git@refs/pull/{os.getenv('REF_NAME')}" +logger.info(f"Running demo tests from {TEST_LIBRARY_REF}") def test_run_dqx_demo_library(make_notebook, make_schema, make_job): From 9d659c7c0bd6f07f107a766eb75c48de3ec7c345 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 28 Jul 2025 18:53:45 -0400 Subject: [PATCH 55/91] Fix PII detection demo --- demos/dqx_demo_pii_detection.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/demos/dqx_demo_pii_detection.py b/demos/dqx_demo_pii_detection.py index 679f8e366..ef5728429 100644 --- a/demos/dqx_demo_pii_detection.py +++ b/demos/dqx_demo_pii_detection.py @@ -74,7 +74,7 @@ def get_entity_mapping(data: str) -> str | None: # Validate and return the results: for result in results: # Ignore if the result is low confidence: - if result.score < 0.6: + if result.score < 0.5: continue # Append the result to the output: output.append({ @@ -83,8 +83,9 @@ def get_entity_mapping(data: str) -> str | None: "end": int(result.end), "score": float(result.score), }) - # Return the results as JSON: - return json.dumps(output) + if output != []: + # Return the results as JSON: + return json.dumps(output) return None # COMMAND ---------- From ca6f386b89cebbf39416389d6a9067fb3b174395 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 28 Jul 2025 19:00:14 -0400 Subject: [PATCH 56/91] Fix PII detection demo --- demos/dqx_demo_pii_detection.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/demos/dqx_demo_pii_detection.py b/demos/dqx_demo_pii_detection.py index ef5728429..4b976ac78 100644 --- a/demos/dqx_demo_pii_detection.py +++ b/demos/dqx_demo_pii_detection.py @@ -26,8 +26,6 @@ %pip install presidio_analyzer numpy==1.23.5 -%restart_python - # COMMAND ---------- dbutils.library.restartPython() From 74305be66b4d799bec9d34b13f2c82757f9e51bd Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 28 Jul 2025 19:39:44 -0400 Subject: [PATCH 57/91] Fix PII detection demo --- demos/dqx_demo_pii_detection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/demos/dqx_demo_pii_detection.py b/demos/dqx_demo_pii_detection.py index 4b976ac78..4d59a6c96 100644 --- a/demos/dqx_demo_pii_detection.py +++ b/demos/dqx_demo_pii_detection.py @@ -41,7 +41,7 @@ from databricks.sdk import WorkspaceClient from databricks.labs.dqx.engine import DQEngine from databricks.labs.dqx.rule import DQRowRule -from databricks.labs.dqx.row_checks import make_condition +from databricks.labs.dqx.check_funcs import make_condition # COMMAND ---------- From f95423f33475815010c029a56ef6be0f7e23a10b Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 28 Jul 2025 19:58:20 -0400 Subject: [PATCH 58/91] Fix PII detection demo --- demos/dqx_demo_pii_detection.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/demos/dqx_demo_pii_detection.py b/demos/dqx_demo_pii_detection.py index 4d59a6c96..065467f97 100644 --- a/demos/dqx_demo_pii_detection.py +++ b/demos/dqx_demo_pii_detection.py @@ -108,20 +108,20 @@ def contains_pii(batch: pd.Series) -> pd.Series: # COMMAND ---------- -def does_not_contain_pii(col_name: str) -> Column: +def does_not_contain_pii(column: str) -> Column: # Define a PII detection expression calling the pandas UDF: - pii_info = contains_pii(col(col_name)) + pii_info = contains_pii(col(column)) # Return the DQX condition that uses the PII detection expression: return make_condition( pii_info.isNotNull(), concat_ws( ' ', - lit(col_name), + lit(column), lit('contains pii with the following info:'), pii_info ), - f'{col_name}_contains_pii' + f'{column}_contains_pii' ) # Define the DQX rule: From 3e9c771ef343f3393fd21458bb8596382d16ded2 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 28 Jul 2025 23:26:16 -0400 Subject: [PATCH 59/91] Use waiter for pipeline updates --- tests/e2e/test_run_demos.py | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 077e674da..0f2b70fde 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -1,12 +1,11 @@ import logging import os -import time from datetime import timedelta from pathlib import Path from databricks.sdk import WorkspaceClient from databricks.sdk.service.workspace import ImportFormat -from databricks.sdk.service.pipelines import NotebookLibrary, PipelineLibrary, UpdateInfoState +from databricks.sdk.service.pipelines import GetPipelineResponse, NotebookLibrary, PipelineLibrary, UpdateInfoState from databricks.sdk.service.jobs import NotebookTask, Run, Task, TerminationTypeType logging.getLogger("tests").setLevel("DEBUG") @@ -125,15 +124,14 @@ def test_run_dqx_dlt_demo(make_notebook, make_pipeline): notebook_path = notebook.as_fuse().as_posix() pipeline = make_pipeline(libraries=[PipelineLibrary(notebook=NotebookLibrary(notebook_path))]) - update = ws.pipelines.start_update(pipeline.pipeline_id) + ws.pipelines.start_update(pipeline.pipeline_id) - while True: - update_details = ws.pipelines.get_update(pipeline_id=pipeline.pipeline_id, update_id=update.update_id) - if update_details.update.state in [UpdateInfoState.CANCELED, UpdateInfoState.COMPLETED, UpdateInfoState.FAILED]: - break - time.sleep(RETRY_INTERVAL_SECONDS) - - assert update_details.update.state == UpdateInfoState.COMPLETED, f"Run of pipeline '{pipeline.pipeline_id}' failed" + ws.pipelines.wait_get_pipeline_idle( + pipeline_id=pipeline.pipeline_id, + timeout=timedelta(minutes=30), + callback=validate_demo_update_status, + ) + logging.info(f"Pipeline {pipeline.pipeline_id} update completed successfully for dqx_dlt_demo") def test_run_dqx_demo_tool(installation_ctx, make_schema, make_notebook, make_job): @@ -190,3 +188,12 @@ def validate_demo_run_status(run: Run, client: WorkspaceClient | None) -> None: assert ( termination_details.type == TerminationTypeType.SUCCESS ), f"Run of '{task.task_key}' failed with message: {client.jobs.get_run_output(task.run_id).error}" + + +def validate_demo_update_status(pipeline: GetPipelineResponse) -> None: + """ + Validates that a demo pipeline update completed successfully. + :param pipeline: `GetPipelineResponse` object returned by the Databricks SDK + """ + update = pipeline.latest_updates[0] # updates are ordered by latest creation date + assert update.state == UpdateInfoState.COMPLETED, f"Run of pipeline {pipeline.pipeline_id} update failed" From 551c627a21aa34952cd93465a762fa3ae733def4 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Mon, 28 Jul 2025 23:34:47 -0400 Subject: [PATCH 60/91] Fix DLT demo execution --- tests/e2e/test_run_demos.py | 27 ++++++++++----------------- 1 file changed, 10 insertions(+), 17 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 0f2b70fde..f0c76494e 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -5,8 +5,8 @@ from pathlib import Path from databricks.sdk import WorkspaceClient from databricks.sdk.service.workspace import ImportFormat -from databricks.sdk.service.pipelines import GetPipelineResponse, NotebookLibrary, PipelineLibrary, UpdateInfoState -from databricks.sdk.service.jobs import NotebookTask, Run, Task, TerminationTypeType +from databricks.sdk.service.pipelines import NotebookLibrary, PipelineLibrary +from databricks.sdk.service.jobs import NotebookTask, PipelineTask, Run, Task, TerminationTypeType logging.getLogger("tests").setLevel("DEBUG") logging.getLogger("databricks.labs.dqx").setLevel("DEBUG") @@ -116,7 +116,7 @@ def test_run_dqx_demo_pii_detection(make_notebook, make_cluster, make_job): logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_pii_detection") -def test_run_dqx_dlt_demo(make_notebook, make_pipeline): +def test_run_dqx_dlt_demo(make_notebook, make_pipeline, make_job): path = Path(__file__).parent.parent.parent / "demos" / "dqx_dlt_demo.py" ws = WorkspaceClient() with open(path, "rb") as f: @@ -124,14 +124,16 @@ def test_run_dqx_dlt_demo(make_notebook, make_pipeline): notebook_path = notebook.as_fuse().as_posix() pipeline = make_pipeline(libraries=[PipelineLibrary(notebook=NotebookLibrary(notebook_path))]) - ws.pipelines.start_update(pipeline.pipeline_id) + pipeline_task = PipelineTask(pipeline_id=pipeline.pipeline_id) + job = make_job(tasks=[Task(task_key="dqx_dlt_demo", pipeline_task=pipeline_task)]) - ws.pipelines.wait_get_pipeline_idle( - pipeline_id=pipeline.pipeline_id, + waiter = ws.jobs.run_now_and_wait(job.job_id) + run = ws.jobs.wait_get_run_job_terminated_or_skipped( + run_id=waiter.run_id, timeout=timedelta(minutes=30), - callback=validate_demo_update_status, + callback=lambda r: validate_demo_run_status(r, client=ws), ) - logging.info(f"Pipeline {pipeline.pipeline_id} update completed successfully for dqx_dlt_demo") + logging.info(f"Job run {run.run_id} completed successfully for dqx_dlt_demo") def test_run_dqx_demo_tool(installation_ctx, make_schema, make_notebook, make_job): @@ -188,12 +190,3 @@ def validate_demo_run_status(run: Run, client: WorkspaceClient | None) -> None: assert ( termination_details.type == TerminationTypeType.SUCCESS ), f"Run of '{task.task_key}' failed with message: {client.jobs.get_run_output(task.run_id).error}" - - -def validate_demo_update_status(pipeline: GetPipelineResponse) -> None: - """ - Validates that a demo pipeline update completed successfully. - :param pipeline: `GetPipelineResponse` object returned by the Databricks SDK - """ - update = pipeline.latest_updates[0] # updates are ordered by latest creation date - assert update.state == UpdateInfoState.COMPLETED, f"Run of pipeline {pipeline.pipeline_id} update failed" From 55024ca6f2d3f5adabefd40b495d256bf6947d14 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 29 Jul 2025 09:01:08 -0400 Subject: [PATCH 61/91] Run DLT demo from current PR branch --- tests/e2e/test_run_demos.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index f0c76494e..fbb469941 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -5,7 +5,7 @@ from pathlib import Path from databricks.sdk import WorkspaceClient from databricks.sdk.service.workspace import ImportFormat -from databricks.sdk.service.pipelines import NotebookLibrary, PipelineLibrary +from databricks.sdk.service.pipelines import NotebookLibrary, PipelinesEnvironment, PipelineLibrary from databricks.sdk.service.jobs import NotebookTask, PipelineTask, Run, Task, TerminationTypeType logging.getLogger("tests").setLevel("DEBUG") @@ -123,7 +123,10 @@ def test_run_dqx_dlt_demo(make_notebook, make_pipeline, make_job): notebook = make_notebook(content=f, format=ImportFormat.SOURCE) notebook_path = notebook.as_fuse().as_posix() - pipeline = make_pipeline(libraries=[PipelineLibrary(notebook=NotebookLibrary(notebook_path))]) + pipeline = make_pipeline( + libraries=[PipelineLibrary(notebook=NotebookLibrary(notebook_path))], + environment=PipelinesEnvironment(dependencies=[TEST_LIBRARY_REF]), + ) pipeline_task = PipelineTask(pipeline_id=pipeline.pipeline_id) job = make_job(tasks=[Task(task_key="dqx_dlt_demo", pipeline_task=pipeline_task)]) From dadd36a6a14cd8e59ba932b469932949281d9f4d Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 29 Jul 2025 14:16:11 -0400 Subject: [PATCH 62/91] Dynamic retrieval of Git URL and branch --- demos/dqx_demo_tool.py | 2 +- demos/dqx_manufacturing_demo.py | 2 +- tests/conftest.py | 3 ++- tests/e2e/test_run_demos.py | 22 +++++++++++++++------- 4 files changed, 19 insertions(+), 10 deletions(-) diff --git a/demos/dqx_demo_tool.py b/demos/dqx_demo_tool.py index 5abbb30b7..5e115651a 100644 --- a/demos/dqx_demo_tool.py +++ b/demos/dqx_demo_tool.py @@ -62,7 +62,7 @@ dqx_wheel_files = glob.glob(dqx_wheel_files_path) try: dqx_latest_wheel = max(dqx_wheel_files, key=os.path.getctime) -except: +except ValueError: raise ValueError(f"No files in path: {dqx_wheel_files_path}") %pip install {dqx_latest_wheel} diff --git a/demos/dqx_manufacturing_demo.py b/demos/dqx_manufacturing_demo.py index 55cb3d8b2..213d1b639 100644 --- a/demos/dqx_manufacturing_demo.py +++ b/demos/dqx_manufacturing_demo.py @@ -806,7 +806,7 @@ print(status) assert not status.has_errors -# COMMAND ---------- +# COMMAND ---------- # save checks in a workspace location sensor_dq_rules_yaml = f"{quality_rules_path}/sensor_dq_rules.yml" diff --git a/tests/conftest.py b/tests/conftest.py index 658b6f5ba..306215d79 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -32,7 +32,8 @@ def product_info(): @pytest.fixture def set_utc_timezone(): """ - Set the timezone to UTC for the duration of the test to make sure spark timestamps are handled the same way regardless of the environment. + Set the timezone to UTC for the duration of the test to make sure spark + timestamps are handled the same way regardless of the environment. """ os.environ["TZ"] = "UTC" yield diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index fbb469941..5423b11dc 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -1,5 +1,5 @@ import logging -import os +import subprocess from datetime import timedelta from pathlib import Path @@ -12,8 +12,19 @@ logging.getLogger("databricks.labs.dqx").setLevel("DEBUG") logger = logging.getLogger(__name__) + +def get_git_url(repo_path: str = ".") -> str: + """Gets the Git remote URL""" + return subprocess.check_output(["git", "-C", repo_path, "config", "--get", "remote.origin.url"], text=True).strip() + + +def get_git_branch(repo_path: str = ".") -> str: + """Gets the Git branch name""" + return subprocess.check_output(["git", "-C", repo_path, "rev-parse", "--abbrev-ref", "HEAD"], text=True).strip() + + RETRY_INTERVAL_SECONDS = 30 -TEST_LIBRARY_REF = f"git+https://github.com/databrickslabs/dqx.git@refs/pull/{os.getenv('REF_NAME')}" +TEST_LIBRARY_REF = f"git+{get_git_url()}@{get_git_branch()}" logger.info(f"Running demo tests from {TEST_LIBRARY_REF}") @@ -178,15 +189,12 @@ def test_run_dqx_demo_tool(installation_ctx, make_schema, make_notebook, make_jo logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_tool") -def validate_demo_run_status(run: Run, client: WorkspaceClient | None) -> None: +def validate_demo_run_status(run: Run, client: WorkspaceClient) -> None: """ Validates that a demo run completed successfully. :param run: `Run` object returned from a `WorkspaceClient.jobs.submit(...)` command - :param client: Optional `WorkspaceClient` object for getting task output + :param client: `WorkspaceClient` object for getting task output """ - if client is None: - client = WorkspaceClient() - task = run.tasks[0] termination_details = run.status.termination_details From 323eba5f6ba1e8ec989b469a8ace0d0b73b25489 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 29 Jul 2025 17:22:50 -0400 Subject: [PATCH 63/91] Fix Git URL for nightly and acceptance tests --- demos/dqx_demo_tool.py | 2 +- demos/dqx_manufacturing_demo.py | 2 +- tests/conftest.py | 3 +-- tests/e2e/test_run_demos.py | 17 ++++------------- 4 files changed, 7 insertions(+), 17 deletions(-) diff --git a/demos/dqx_demo_tool.py b/demos/dqx_demo_tool.py index 5e115651a..5abbb30b7 100644 --- a/demos/dqx_demo_tool.py +++ b/demos/dqx_demo_tool.py @@ -62,7 +62,7 @@ dqx_wheel_files = glob.glob(dqx_wheel_files_path) try: dqx_latest_wheel = max(dqx_wheel_files, key=os.path.getctime) -except ValueError: +except: raise ValueError(f"No files in path: {dqx_wheel_files_path}") %pip install {dqx_latest_wheel} diff --git a/demos/dqx_manufacturing_demo.py b/demos/dqx_manufacturing_demo.py index 213d1b639..55cb3d8b2 100644 --- a/demos/dqx_manufacturing_demo.py +++ b/demos/dqx_manufacturing_demo.py @@ -806,7 +806,7 @@ print(status) assert not status.has_errors -# COMMAND ---------- +# COMMAND ---------- # save checks in a workspace location sensor_dq_rules_yaml = f"{quality_rules_path}/sensor_dq_rules.yml" diff --git a/tests/conftest.py b/tests/conftest.py index 306215d79..658b6f5ba 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -32,8 +32,7 @@ def product_info(): @pytest.fixture def set_utc_timezone(): """ - Set the timezone to UTC for the duration of the test to make sure spark - timestamps are handled the same way regardless of the environment. + Set the timezone to UTC for the duration of the test to make sure spark timestamps are handled the same way regardless of the environment. """ os.environ["TZ"] = "UTC" yield diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 5423b11dc..a53ca220d 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -1,5 +1,5 @@ import logging -import subprocess +import os from datetime import timedelta from pathlib import Path @@ -12,19 +12,10 @@ logging.getLogger("databricks.labs.dqx").setLevel("DEBUG") logger = logging.getLogger(__name__) - -def get_git_url(repo_path: str = ".") -> str: - """Gets the Git remote URL""" - return subprocess.check_output(["git", "-C", repo_path, "config", "--get", "remote.origin.url"], text=True).strip() - - -def get_git_branch(repo_path: str = ".") -> str: - """Gets the Git branch name""" - return subprocess.check_output(["git", "-C", repo_path, "rev-parse", "--abbrev-ref", "HEAD"], text=True).strip() - - RETRY_INTERVAL_SECONDS = 30 -TEST_LIBRARY_REF = f"git+{get_git_url()}@{get_git_branch()}" +TEST_LIBRARY_REF = f"git+https://github.com/databrickslabs/dqx" +if os.getenv("REF_NAME"): + TEST_LIBRARY_REF = f"{TEST_LIBRARY_REF}.git@refs/pull/{os.getenv('REF_NAME')}" logger.info(f"Running demo tests from {TEST_LIBRARY_REF}") From e131fa525f74f6ddd814536c6f993c8169ed960a Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Tue, 29 Jul 2025 17:29:11 -0400 Subject: [PATCH 64/91] Fmt --- tests/e2e/test_run_demos.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index a53ca220d..af5b0bf73 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -13,7 +13,7 @@ logger = logging.getLogger(__name__) RETRY_INTERVAL_SECONDS = 30 -TEST_LIBRARY_REF = f"git+https://github.com/databrickslabs/dqx" +TEST_LIBRARY_REF = "git+https://github.com/databrickslabs/dqx" if os.getenv("REF_NAME"): TEST_LIBRARY_REF = f"{TEST_LIBRARY_REF}.git@refs/pull/{os.getenv('REF_NAME')}" logger.info(f"Running demo tests from {TEST_LIBRARY_REF}") From 5699dd357899468896d5982350c838e8992b035a Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 30 Jul 2025 08:24:53 -0400 Subject: [PATCH 65/91] Comment for acceptance actions env var --- .github/workflows/acceptance.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index 246d45189..84bbc80b1 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -134,7 +134,7 @@ jobs: timeout: 2h codegen_path: tests/e2e/.codegen.json env: - REF_NAME: ${{ github.ref_name }} + REF_NAME: ${{ github.ref_name }} # NOTE: end-to-end tests use this to pip install from the current PR branch GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }} @@ -168,7 +168,7 @@ jobs: timeout: 2h codegen_path: tests/e2e/.codegen.json env: - REF_NAME: ${{ github.ref_name }} + REF_NAME: ${{ github.ref_name }} # NOTE: end-to-end tests use this to pip install from the current PR branch GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }} From 16e64cbb3607389a1cdc157debfc1970e1f29f36 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 12:28:14 -0400 Subject: [PATCH 66/91] Add test for asset bundle demo --- .github/workflows/acceptance.yml | 36 ++++++++++++++++++++++++++++++++ .github/workflows/nightly.yml | 36 ++++++++++++++++++++++++++++++++ tests/e2e/test_run_demos.py | 11 ++++++++++ 3 files changed, 83 insertions(+) diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index 84bbc80b1..6802b9edc 100644 --- a/.github/workflows/acceptance.yml +++ b/.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: diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 0d1b8c58a..9d4c67563 100644 --- a/.github/workflows/nightly.yml +++ b/.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: diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index af5b0bf73..04c1529b9 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -1,5 +1,6 @@ import logging import os +import subprocess from datetime import timedelta from pathlib import Path @@ -180,6 +181,16 @@ def test_run_dqx_demo_tool(installation_ctx, make_schema, make_notebook, make_jo logging.info(f"Job run {run.run_id} completed successfully for dqx_demo_tool") +def test_run_dqx_demo_asset_bundle(): + which_output = subprocess.run(["which", "databricks"], capture_output=True, text=True, check=True) + cli_path = which_output.stdout.strip() + + subprocess.run([cli_path, "bundle", "validate"], check=True) + subprocess.run([cli_path, "bundle", "deploy"], check=True) + subprocess.run([cli_path, "bundle", "run", "dqx_demo_job"], check=True) + subprocess.run([cli_path, "bundle", "destroy"], check=True) + + def validate_demo_run_status(run: Run, client: WorkspaceClient) -> None: """ Validates that a demo run completed successfully. From 5fc8bd01fa2fe4d6683a8f58a4527129f248bdd2 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 12:39:22 -0400 Subject: [PATCH 67/91] Remove os import --- tests/e2e/test_run_demos.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 022ae5d98..097959927 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -1,5 +1,4 @@ import logging -import os import subprocess from datetime import timedelta From 106f99e0774ef2f071bf5dae5a6dbf1d4b2b1fdd Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 12:40:53 -0400 Subject: [PATCH 68/91] Update nightly actions --- .github/workflows/nightly.yml | 36 +++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 77ff9dd4d..076727225 100644 --- a/.github/workflows/nightly.yml +++ b/.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: From 5d6831e254084006dd1d49c743056ecbc5de810d Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 12:42:38 -0400 Subject: [PATCH 69/91] Fix actions --- .github/workflows/acceptance.yml | 3 ++- .github/workflows/nightly.yml | 2 ++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index 92e0f07db..d4def8ad1 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -202,8 +202,9 @@ jobs: with: vault_uri: ${{ secrets.VAULT_URI }} timeout: 2h - codegen_path: tests/integration/.codegen.json + codegen_path: tests/e2e/.codegen.json env: + REF_NAME: ${{ github.ref_name }} # NOTE: end-to-end tests use this to pip install from the current PR branch GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }} diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 076727225..2f0460edd 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -142,6 +142,7 @@ jobs: create_issues: true codegen_path: tests/e2e/.codegen.json env: + REF_NAME: ${{ github.ref_name }} # NOTE: end-to-end tests use this to pip install from the current PR branch GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }} @@ -194,6 +195,7 @@ jobs: create_issues: true codegen_path: tests/e2e/.codegen.json env: + REF_NAME: ${{ github.ref_name }} # NOTE: end-to-end tests use this to pip install from the current PR branch GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }} From 07e968372d97104b5c7399aa375ed16ce137cdf6 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 12:43:40 -0400 Subject: [PATCH 70/91] Fix nightly actions --- .github/workflows/nightly.yml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 2f0460edd..076727225 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -142,7 +142,6 @@ jobs: create_issues: true codegen_path: tests/e2e/.codegen.json env: - REF_NAME: ${{ github.ref_name }} # NOTE: end-to-end tests use this to pip install from the current PR branch GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }} @@ -195,7 +194,6 @@ jobs: create_issues: true codegen_path: tests/e2e/.codegen.json env: - REF_NAME: ${{ github.ref_name }} # NOTE: end-to-end tests use this to pip install from the current PR branch GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }} ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }} From 7de797916176ea9d88b56addd24bfae2a9536850 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 13:06:38 -0400 Subject: [PATCH 71/91] Login --- tests/e2e/test_run_demos.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 097959927..631c156e8 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -234,6 +234,7 @@ def test_run_dqx_demo_asset_bundle(): which_output = subprocess.run(["which", "databricks"], capture_output=True, text=True, check=True) cli_path = which_output.stdout.strip() + subprocess.run([cli_path, "auth", "login"], check=True) subprocess.run([cli_path, "bundle", "validate"], check=True) subprocess.run([cli_path, "bundle", "deploy"], check=True) subprocess.run([cli_path, "bundle", "run", "dqx_demo_job"], check=True) From 4900c52d714a00ac5115426accdbc4f6636a89ac Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 13:16:20 -0400 Subject: [PATCH 72/91] Raise stderr output on failure --- tests/e2e/test_run_demos.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 631c156e8..e1b292183 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -234,11 +234,14 @@ def test_run_dqx_demo_asset_bundle(): which_output = subprocess.run(["which", "databricks"], capture_output=True, text=True, check=True) cli_path = which_output.stdout.strip() - subprocess.run([cli_path, "auth", "login"], check=True) - subprocess.run([cli_path, "bundle", "validate"], check=True) - subprocess.run([cli_path, "bundle", "deploy"], check=True) - subprocess.run([cli_path, "bundle", "run", "dqx_demo_job"], check=True) - subprocess.run([cli_path, "bundle", "destroy"], check=True) + try: + subprocess.run([cli_path, "auth", "login"], check=True, capture_output=True) + subprocess.run([cli_path, "bundle", "validate"], check=True, capture_output=True) + subprocess.run([cli_path, "bundle", "deploy"], check=True, capture_output=True) + subprocess.run([cli_path, "bundle", "run", "dqx_demo_job"], check=True, capture_output=True) + subprocess.run([cli_path, "bundle", "destroy"], check=True, capture_output=True) + except subprocess.CalledProcessError as ex: + raise Exception(ex.stderr) def validate_run_status(run: Run, client: WorkspaceClient) -> None: From dc6201054a36ad47ed2b3dd15a431e8e27308561 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 13:31:24 -0400 Subject: [PATCH 73/91] Add host --- tests/e2e/test_run_demos.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index e1b292183..27e218382 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -1,4 +1,5 @@ import logging +import os import subprocess from datetime import timedelta @@ -235,13 +236,15 @@ def test_run_dqx_demo_asset_bundle(): cli_path = which_output.stdout.strip() try: - subprocess.run([cli_path, "auth", "login"], check=True, capture_output=True) + subprocess.run( + [cli_path, "auth", "login", "--host", os.environ["DATABRICKS_HOST"]], check=True, capture_output=True + ) subprocess.run([cli_path, "bundle", "validate"], check=True, capture_output=True) subprocess.run([cli_path, "bundle", "deploy"], check=True, capture_output=True) subprocess.run([cli_path, "bundle", "run", "dqx_demo_job"], check=True, capture_output=True) subprocess.run([cli_path, "bundle", "destroy"], check=True, capture_output=True) except subprocess.CalledProcessError as ex: - raise Exception(ex.stderr) + raise AssertionError(ex.stderr) from ex def validate_run_status(run: Run, client: WorkspaceClient) -> None: From 4b69ac8db5b775580d10b00f7bd49198a9b317dd Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 14:02:44 -0400 Subject: [PATCH 74/91] Remove auth --- tests/e2e/test_run_demos.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 27e218382..22f1bf23f 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -236,9 +236,6 @@ def test_run_dqx_demo_asset_bundle(): cli_path = which_output.stdout.strip() try: - subprocess.run( - [cli_path, "auth", "login", "--host", os.environ["DATABRICKS_HOST"]], check=True, capture_output=True - ) subprocess.run([cli_path, "bundle", "validate"], check=True, capture_output=True) subprocess.run([cli_path, "bundle", "deploy"], check=True, capture_output=True) subprocess.run([cli_path, "bundle", "run", "dqx_demo_job"], check=True, capture_output=True) From 96e4f3f1aa7df89959d74ce10b9b41c05d476963 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 14:11:49 -0400 Subject: [PATCH 75/91] Run from the correct directory --- tests/e2e/test_run_demos.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 22f1bf23f..9df93b1b0 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -1,5 +1,4 @@ import logging -import os import subprocess from datetime import timedelta @@ -234,8 +233,10 @@ def test_run_dqx_streaming_demo_diy(make_notebook, make_job, tmp_path, library_r def test_run_dqx_demo_asset_bundle(): which_output = subprocess.run(["which", "databricks"], capture_output=True, text=True, check=True) cli_path = which_output.stdout.strip() + path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_asset_bundle" try: + subprocess.run([f"cd {path.as_posix()}"], check=True, capture_output=True) subprocess.run([cli_path, "bundle", "validate"], check=True, capture_output=True) subprocess.run([cli_path, "bundle", "deploy"], check=True, capture_output=True) subprocess.run([cli_path, "bundle", "run", "dqx_demo_job"], check=True, capture_output=True) From 0ccb4ee2f1c1896e40f3718972ca017fd5c32783 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 14:28:23 -0400 Subject: [PATCH 76/91] Update DAB path --- tests/e2e/test_run_demos.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 9df93b1b0..bed98d3d3 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -233,7 +233,7 @@ def test_run_dqx_streaming_demo_diy(make_notebook, make_job, tmp_path, library_r def test_run_dqx_demo_asset_bundle(): which_output = subprocess.run(["which", "databricks"], capture_output=True, text=True, check=True) cli_path = which_output.stdout.strip() - path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_asset_bundle" + path = Path(__file__).parent.parent.parent.parent / "demos" / "dqx_demo_asset_bundle" try: subprocess.run([f"cd {path.as_posix()}"], check=True, capture_output=True) From cee08d8f8865fc4591b536b3f01f38819169b7d1 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 15:04:12 -0400 Subject: [PATCH 77/91] Refactor --- tests/e2e/test_run_demos.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index bed98d3d3..ab21e9c77 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -1,4 +1,5 @@ import logging +import shutil import subprocess from datetime import timedelta @@ -231,18 +232,17 @@ def test_run_dqx_streaming_demo_diy(make_notebook, make_job, tmp_path, library_r def test_run_dqx_demo_asset_bundle(): - which_output = subprocess.run(["which", "databricks"], capture_output=True, text=True, check=True) - cli_path = which_output.stdout.strip() - path = Path(__file__).parent.parent.parent.parent / "demos" / "dqx_demo_asset_bundle" + cli_path = shutil.which("databricks") + path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_asset_bundle" try: - subprocess.run([f"cd {path.as_posix()}"], check=True, capture_output=True) - subprocess.run([cli_path, "bundle", "validate"], check=True, capture_output=True) - subprocess.run([cli_path, "bundle", "deploy"], check=True, capture_output=True) - subprocess.run([cli_path, "bundle", "run", "dqx_demo_job"], check=True, capture_output=True) - subprocess.run([cli_path, "bundle", "destroy"], check=True, capture_output=True) + subprocess.run([cli_path, "bundle", "validate"], check=True, capture_output=True, cwd=path) + subprocess.run([cli_path, "bundle", "deploy"], 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 + finally: + subprocess.run([cli_path, "bundle", "destroy"], check=True, capture_output=True, cwd=path) def validate_run_status(run: Run, client: WorkspaceClient) -> None: From 4d18c7bfff19e40d81a6d45456de5a70a0410960 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 15:08:21 -0400 Subject: [PATCH 78/91] Refactor --- tests/e2e/test_run_demos.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index ab21e9c77..108def2fe 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -239,10 +239,9 @@ def test_run_dqx_demo_asset_bundle(): subprocess.run([cli_path, "bundle", "validate"], check=True, capture_output=True, cwd=path) subprocess.run([cli_path, "bundle", "deploy"], check=True, capture_output=True, cwd=path) subprocess.run([cli_path, "bundle", "run", "dqx_demo_job"], check=True, capture_output=True, cwd=path) + subprocess.run([cli_path, "bundle", "destroy"], check=True, capture_output=True, cwd=path) except subprocess.CalledProcessError as ex: raise AssertionError(ex.stderr) from ex - finally: - subprocess.run([cli_path, "bundle", "destroy"], check=True, capture_output=True, cwd=path) def validate_run_status(run: Run, client: WorkspaceClient) -> None: From 4fa50d1fd2b1ece62b7c7a32ab7a674b7a8a2657 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 15:57:55 -0400 Subject: [PATCH 79/91] Run DABs test with branch library version --- demos/dqx_demo_asset_bundle/databricks.yml | 24 ++++++++++++++-------- tests/e2e/test_run_demos.py | 9 ++++++-- 2 files changed, 23 insertions(+), 10 deletions(-) diff --git a/demos/dqx_demo_asset_bundle/databricks.yml b/demos/dqx_demo_asset_bundle/databricks.yml index 4258129fe..dd7baa51e 100644 --- a/demos/dqx_demo_asset_bundle/databricks.yml +++ b/demos/dqx_demo_asset_bundle/databricks.yml @@ -14,6 +14,9 @@ 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: @@ -21,6 +24,9 @@ resources: name: "[${workspace.current_user.userName}] DQX Demo Job" tasks: - task_key: dqx_demo_notebook + job_cluster_key: dqx_demo_cluster + libraries: + - pypi: "${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 +35,14 @@ 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: + cluster_name: dqx_demo_cluster + is_single_node: true + spark_version: 15.4.x-scala2.12 + node_type_id: Standard_DS3_v2 + data_security_mode: SINGLE_USER + queue: - enabled: true - environments: - - environment_key: Default - spec: - client: "2" - dependencies: - - databricks-labs-dqx \ No newline at end of file + enabled: true \ No newline at end of file diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 108def2fe..909f92db0 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -231,13 +231,18 @@ 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(): +def test_run_dqx_demo_asset_bundle(library_ref): cli_path = shutil.which("databricks") path = Path(__file__).parent.parent.parent / "demos" / "dqx_demo_asset_bundle" try: subprocess.run([cli_path, "bundle", "validate"], check=True, capture_output=True, cwd=path) - subprocess.run([cli_path, "bundle", "deploy"], check=True, capture_output=True, cwd=path) + subprocess.run( + [cli_path, "bundle", "deploy", f"--var='library_ref={library_ref}'"], + check=True, + capture_output=True, + cwd=path, + ) subprocess.run([cli_path, "bundle", "run", "dqx_demo_job"], check=True, capture_output=True, cwd=path) subprocess.run([cli_path, "bundle", "destroy"], check=True, capture_output=True, cwd=path) except subprocess.CalledProcessError as ex: From d9094ee6a8087ebd9b9b520702e6ab252e912c48 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 16:07:37 -0400 Subject: [PATCH 80/91] Update bundle config --- demos/dqx_demo_asset_bundle/databricks.yml | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/demos/dqx_demo_asset_bundle/databricks.yml b/demos/dqx_demo_asset_bundle/databricks.yml index dd7baa51e..b797b9979 100644 --- a/demos/dqx_demo_asset_bundle/databricks.yml +++ b/demos/dqx_demo_asset_bundle/databricks.yml @@ -26,7 +26,8 @@ resources: - task_key: dqx_demo_notebook job_cluster_key: dqx_demo_cluster libraries: - - pypi: "${var.library_ref}" + - 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 @@ -38,10 +39,8 @@ resources: job_clusters: - job_cluster_key: dqx_demo_cluster new_cluster: - cluster_name: dqx_demo_cluster - is_single_node: true + num_workers: 1 spark_version: 15.4.x-scala2.12 - node_type_id: Standard_DS3_v2 data_security_mode: SINGLE_USER queue: From 5b530bf21d443f1dced1340cc1176fbe798136a0 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 16:18:26 -0400 Subject: [PATCH 81/91] Fix DABs deployment --- demos/dqx_demo_asset_bundle/databricks.yml | 1 + tests/e2e/test_run_demos.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/demos/dqx_demo_asset_bundle/databricks.yml b/demos/dqx_demo_asset_bundle/databricks.yml index b797b9979..ad85165f4 100644 --- a/demos/dqx_demo_asset_bundle/databricks.yml +++ b/demos/dqx_demo_asset_bundle/databricks.yml @@ -40,6 +40,7 @@ resources: - 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 diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 909f92db0..dff63b191 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -238,7 +238,7 @@ def test_run_dqx_demo_asset_bundle(library_ref): 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}'"], + [cli_path, "bundle", "deploy", f'--var="library_ref={library_ref}"'], check=True, capture_output=True, cwd=path, From 7dbe8d1ce0109ef12e6971323b78417e420c88b0 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 16:35:49 -0400 Subject: [PATCH 82/91] Parameterize the demo catalog and schema --- tests/e2e/test_run_demos.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index dff63b191..d6ca2b282 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -231,14 +231,23 @@ 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(library_ref): +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}"'], + [ + cli_path, + "bundle", + "deploy", + f'--var="library_ref={library_ref}"', + f'--var="demo_catalog={catalog}"', + f'--var="demo_schema={schema}"', + ], check=True, capture_output=True, cwd=path, From f9be6e848f068342135a7d2ffed93c73989c80b0 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 16:44:27 -0400 Subject: [PATCH 83/91] Remove create catalog and schema statements --- demos/dqx_demo_asset_bundle/dqx_demo_notebook.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/demos/dqx_demo_asset_bundle/dqx_demo_notebook.py b/demos/dqx_demo_asset_bundle/dqx_demo_notebook.py index ba8d53768..abf96778b 100644 --- a/demos/dqx_demo_asset_bundle/dqx_demo_notebook.py +++ b/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" From 96158e9ef72b697527571784362dc4dcf487dd76 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 16:54:44 -0400 Subject: [PATCH 84/91] Update sample data --- demos/dqx_demo_asset_bundle/dqx_demo_notebook.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/demos/dqx_demo_asset_bundle/dqx_demo_notebook.py b/demos/dqx_demo_asset_bundle/dqx_demo_notebook.py index abf96778b..a999cc2b2 100644 --- a/demos/dqx_demo_asset_bundle/dqx_demo_notebook.py +++ b/demos/dqx_demo_asset_bundle/dqx_demo_notebook.py @@ -223,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, @@ -277,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, @@ -332,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, @@ -346,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, From 1670855567532e8445cabb8d23d535b75819641e Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 17:06:06 -0400 Subject: [PATCH 85/91] Auto-approve commands --- tests/e2e/test_run_demos.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index d6ca2b282..c6bc59687 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -247,13 +247,14 @@ def test_run_dqx_demo_asset_bundle(make_schema, library_ref): 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) - subprocess.run([cli_path, "bundle", "destroy"], check=True, capture_output=True, cwd=path) + subprocess.run([cli_path, "bundle", "destroy", "--auto-approve"], check=True, capture_output=True, cwd=path) except subprocess.CalledProcessError as ex: raise AssertionError(ex.stderr) from ex From 0f104135cd043c5676f13dfc7f4aa60061ef3e02 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Wed, 20 Aug 2025 18:45:29 -0400 Subject: [PATCH 86/91] Don't destroy bundle in tests --- tests/e2e/test_run_demos.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index c6bc59687..6a71491fc 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -254,7 +254,6 @@ def test_run_dqx_demo_asset_bundle(make_schema, library_ref): cwd=path, ) subprocess.run([cli_path, "bundle", "run", "dqx_demo_job"], check=True, capture_output=True, cwd=path) - subprocess.run([cli_path, "bundle", "destroy", "--auto-approve"], check=True, capture_output=True, cwd=path) except subprocess.CalledProcessError as ex: raise AssertionError(ex.stderr) from ex From 7e22e40f53d8c5826e704b8806af27a2c067d9d2 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Thu, 21 Aug 2025 12:10:20 -0400 Subject: [PATCH 87/91] Add target environments for each cloud provider --- demos/dqx_demo_asset_bundle/README.md | 8 ++++---- demos/dqx_demo_asset_bundle/databricks.yml | 13 ++++++++++++- tests/e2e/test_run_demos.py | 7 +++++-- 3 files changed, 21 insertions(+), 7 deletions(-) diff --git a/demos/dqx_demo_asset_bundle/README.md b/demos/dqx_demo_asset_bundle/README.md index 88ea36470..f900faae2 100644 --- a/demos/dqx_demo_asset_bundle/README.md +++ b/demos/dqx_demo_asset_bundle/README.md @@ -23,16 +23,16 @@ The directory includes: Navigate to `/dqx_demo_asset_bundle` and execute the following commands in your console: ``` -databricks bundle deploy +databricks bundle deploy --target ``` -***Note:** We can set the Databricks profile for running CLI commands by specifying the `--profile` option.* +***Notes:** Set the Databricks profile for running CLI commands by specifying the `--profile` option.* ## Running the job To run the job, execute another CLI command: ``` -databricks bundle run +databricks bundle run --target ``` ## Removing the job, notebook, and files @@ -40,7 +40,7 @@ databricks bundle run To tear-down the job, notebook, and files, run a final CLI command: ``` -databricks bundle destroy +databricks bundle destroy --target ``` ***Note:** Catalogs, schemas, or tables created by running the notebook should be removed manually.* \ No newline at end of file diff --git a/demos/dqx_demo_asset_bundle/databricks.yml b/demos/dqx_demo_asset_bundle/databricks.yml index ad85165f4..221e38f70 100644 --- a/demos/dqx_demo_asset_bundle/databricks.yml +++ b/demos/dqx_demo_asset_bundle/databricks.yml @@ -1,6 +1,14 @@ bundle: name: dqx_demo_bundle +targets: + aws: + variables: + instance_type: i3.2xlarge + azure: + variables: + instance_type: Standard_DS3_v2 + variables: # Sets the default value of several notebook parameters demo_catalog: description: Name of the catalog where demo data will be written @@ -17,6 +25,9 @@ variables: # Sets the default value of several notebook parameters library_ref: description: PyPi package name for installing DQX default: databricks-labs-dqx + instance_type: + description: Instance type used to run the demo job. + default: Standard_DS3_v2 resources: jobs: @@ -40,7 +51,7 @@ resources: - job_cluster_key: dqx_demo_cluster new_cluster: num_workers: 1 - node_type_id: Standard_DS3_v2 + node_type_id: "${var.instance_type}" spark_version: 15.4.x-scala2.12 data_security_mode: SINGLE_USER diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 6a71491fc..5ce495bef 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -238,7 +238,7 @@ def test_run_dqx_demo_asset_bundle(make_schema, library_ref): 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", "validate", "--target azure"], check=True, capture_output=True, cwd=path) subprocess.run( [ cli_path, @@ -247,13 +247,16 @@ def test_run_dqx_demo_asset_bundle(make_schema, library_ref): f'--var="library_ref={library_ref}"', f'--var="demo_catalog={catalog}"', f'--var="demo_schema={schema}"', + "--target azure", "--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) + subprocess.run( + [cli_path, "bundle", "run", "dqx_demo_job", "--target azure"], check=True, capture_output=True, cwd=path + ) except subprocess.CalledProcessError as ex: raise AssertionError(ex.stderr) from ex From d09fc40b2f0e1716047a9f9c6e5b08f37a8cc9fb Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Thu, 21 Aug 2025 12:36:24 -0400 Subject: [PATCH 88/91] Refactor test --- tests/e2e/test_run_demos.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 5ce495bef..8e462b1a1 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -238,7 +238,7 @@ def test_run_dqx_demo_asset_bundle(make_schema, library_ref): schema = make_schema(catalog_name=catalog).name try: - subprocess.run([cli_path, "bundle", "validate", "--target azure"], check=True, capture_output=True, cwd=path) + subprocess.run([cli_path, "bundle", "validate", '--target "azure"'], check=True, capture_output=True, cwd=path) subprocess.run( [ cli_path, @@ -247,7 +247,7 @@ def test_run_dqx_demo_asset_bundle(make_schema, library_ref): f'--var="library_ref={library_ref}"', f'--var="demo_catalog={catalog}"', f'--var="demo_schema={schema}"', - "--target azure", + '--target "azure"', "--auto-approve", ], check=True, @@ -255,7 +255,7 @@ def test_run_dqx_demo_asset_bundle(make_schema, library_ref): cwd=path, ) subprocess.run( - [cli_path, "bundle", "run", "dqx_demo_job", "--target azure"], check=True, capture_output=True, cwd=path + [cli_path, "bundle", "run", "dqx_demo_job", '--target "azure"'], check=True, capture_output=True, cwd=path ) except subprocess.CalledProcessError as ex: raise AssertionError(ex.stderr) from ex From 5a6467ed7a28f99239f819da90040494fbef28b0 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Thu, 21 Aug 2025 12:41:15 -0400 Subject: [PATCH 89/91] Fix test --- tests/e2e/test_run_demos.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 8e462b1a1..be0d5f6f5 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -238,7 +238,7 @@ def test_run_dqx_demo_asset_bundle(make_schema, library_ref): schema = make_schema(catalog_name=catalog).name try: - subprocess.run([cli_path, "bundle", "validate", '--target "azure"'], check=True, capture_output=True, cwd=path) + subprocess.run([cli_path, "bundle", "validate"], check=True, capture_output=True, cwd=path) subprocess.run( [ cli_path, @@ -247,7 +247,7 @@ def test_run_dqx_demo_asset_bundle(make_schema, library_ref): f'--var="library_ref={library_ref}"', f'--var="demo_catalog={catalog}"', f'--var="demo_schema={schema}"', - '--target "azure"', + "-t azure", "--auto-approve", ], check=True, @@ -255,7 +255,7 @@ def test_run_dqx_demo_asset_bundle(make_schema, library_ref): cwd=path, ) subprocess.run( - [cli_path, "bundle", "run", "dqx_demo_job", '--target "azure"'], check=True, capture_output=True, cwd=path + [cli_path, "bundle", "run", "dqx_demo_job", "-t azure"], check=True, capture_output=True, cwd=path ) except subprocess.CalledProcessError as ex: raise AssertionError(ex.stderr) from ex From bf369babfcae7c9a13103565764c98f566e7bfb8 Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Thu, 21 Aug 2025 13:12:19 -0400 Subject: [PATCH 90/91] Use a default target --- demos/dqx_demo_asset_bundle/databricks.yml | 1 + tests/e2e/test_run_demos.py | 3 +-- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/demos/dqx_demo_asset_bundle/databricks.yml b/demos/dqx_demo_asset_bundle/databricks.yml index 221e38f70..a4233d5b8 100644 --- a/demos/dqx_demo_asset_bundle/databricks.yml +++ b/demos/dqx_demo_asset_bundle/databricks.yml @@ -6,6 +6,7 @@ targets: variables: instance_type: i3.2xlarge azure: + default: true variables: instance_type: Standard_DS3_v2 diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index be0d5f6f5..53bdba5db 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -247,7 +247,6 @@ def test_run_dqx_demo_asset_bundle(make_schema, library_ref): f'--var="library_ref={library_ref}"', f'--var="demo_catalog={catalog}"', f'--var="demo_schema={schema}"', - "-t azure", "--auto-approve", ], check=True, @@ -255,7 +254,7 @@ def test_run_dqx_demo_asset_bundle(make_schema, library_ref): cwd=path, ) subprocess.run( - [cli_path, "bundle", "run", "dqx_demo_job", "-t azure"], check=True, capture_output=True, cwd=path + [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 From 88092e886a62e3073ca8497f88454bc8e1241e0a Mon Sep 17 00:00:00 2001 From: Greg Hansen Date: Thu, 21 Aug 2025 13:23:35 -0400 Subject: [PATCH 91/91] Fmt --- tests/e2e/test_run_demos.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py index 53bdba5db..6a71491fc 100644 --- a/tests/e2e/test_run_demos.py +++ b/tests/e2e/test_run_demos.py @@ -253,9 +253,7 @@ def test_run_dqx_demo_asset_bundle(make_schema, library_ref): capture_output=True, cwd=path, ) - subprocess.run( - [cli_path, "bundle", "run", "dqx_demo_job"], 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