Use single lakebase instance for all integration tests and add retry logic (#618)

mwojtyczka · tlgnr · web-flow · commit 6fbf5ae533f4 · 2025-10-18T23:37:05.000+02:00
## Changes
&lt;!-- Summary of your changes that are easy to understand. Add
screenshots when necessary --&gt;

* Create a single lakebase instance for all tests.
* Implement retry logic in case workspace quota limit for the number of
lakebase instances is exceeded.
* Bumped pytester version.
* Bumped hatch version to avoid click compatibility issues.

### Tests
&lt;!-- How is this tested? Please see the checklist below and also
describe any other relevant tests --&gt;

- [x] manually tested
- [ ] added unit tests
- [x] added integration tests
- [ ] added end-to-end tests
- [ ] added performance tests

---------

Co-authored-by: Manuel Tilgner &lt;manuel.tilgner@databricks.com&gt;
diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml
@@ -41,8 +41,7 @@ jobs:
           python-version: '3.12'
 
       - name: Install hatch
-        # click 8.3+ introduced bug for hatch
-        run: pip install "hatch==1.13.0" "click<8.3"
+        run: pip install hatch==1.15.0
 
       - name: Run unit tests and generate test coverage report
         run: make test
@@ -94,8 +93,7 @@ jobs:
           python-version: '3.12'
 
       - name: Install hatch
-        # click 8.3+ introduced bug for hatch
-        run: pip install "hatch==1.13.0" "click<8.3"
+        run: pip install hatch==1.15.0
 
       - name: Run integration tests on serverless cluster
         uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.4
@@ -127,8 +125,7 @@ jobs:
           python-version: '3.12'
 
       - name: Install hatch
-        # click 8.3+ introduced bug for hatch
-        run: pip install "hatch==1.13.0" "click<8.3"
+        run: pip install hatch==1.15.0
 
       - name: Install Databricks CLI
         run: |
@@ -180,8 +177,7 @@ jobs:
           python-version: '3.12'
 
       - name: Install hatch
-        # click 8.3+ introduced bug for hatch
-        run: pip install "hatch==1.13.0" "click<8.3"
+        run: pip install hatch==1.15.0
 
       - name: Install Databricks CLI
         run: |
diff --git a/.github/workflows/docs-release.yml b/.github/workflows/docs-release.yml
@@ -28,8 +28,7 @@ jobs:
 
       - name: Install Hatch
         run: |
-          # click 8.3+ introduced bug for hatch
-          pip install "hatch==1.13.0" "click<8.3"
+          pip install hatch==1.15.0
 
       - uses: actions/setup-node@v4
         with:
diff --git a/.github/workflows/downstreams.yml b/.github/workflows/downstreams.yml
@@ -43,8 +43,7 @@ jobs:
 
       - name: Install toolchain
         run: |
-          # click 8.3+ introduced bug for hatch
-          pip install "hatch==1.13.0" "click<8.3"
+          pip install hatch==1.15.0
 
       - name: Check downstream compatibility
         uses: databrickslabs/sandbox/downstreams@downstreams/v0.0.1
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
@@ -32,8 +32,7 @@ jobs:
           python-version: '3.12'
 
       - name: Install hatch
-        # click 8.3+ introduced bug for hatch
-        run: pip install "hatch==1.13.0" "click<8.3"
+        run: pip install hatch==1.15.0
 
       - name: Run unit tests and generate test coverage report
         run: make test
@@ -82,8 +81,7 @@ jobs:
           python-version: '3.12'
 
       - name: Install hatch
-        # click 8.3+ introduced bug for hatch
-        run: pip install "hatch==1.13.0" "click<8.3"
+        run: pip install hatch==1.15.0
 
       - name: Run integration tests on serverless cluster
         uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.4
@@ -115,8 +113,7 @@ jobs:
           python-version: '3.12'
 
       - name: Install hatch
-        # click 8.3+ introduced bug for hatch
-        run: pip install "hatch==1.13.0" "click<8.3"
+        run: pip install hatch==1.15.0
 
       - name: Install Databricks CLI
         run: |
@@ -167,8 +164,7 @@ jobs:
           python-version: '3.12'
 
       - name: Install hatch
-        # click 8.3+ introduced bug for hatch
-        run: pip install "hatch==1.13.0" "click<8.3"
+        run: pip install hatch==1.15.0
 
       - name: Install Databricks CLI
         run: |
@@ -221,8 +217,7 @@ jobs:
           python-version: '3.12'
 
       - name: Install hatch
-        # click 8.3+ introduced bug for hatch
-        run: pip install "hatch==1.13.0" "click<8.3"
+        run: pip install hatch==1.15.0
 
       - name: Login to Azure for azure-cli authentication
         uses: azure/login@v2
diff --git a/.github/workflows/performance.yml b/.github/workflows/performance.yml
@@ -43,8 +43,7 @@ jobs:
           cache-dependency-path: '**/pyproject.toml'
 
       - name: Install hatch
-        # click 8.3+ introduced bug for hatch
-        run: pip install "hatch==1.13.0" "click<8.3"
+        run: pip install hatch==1.15.0
 
       - name: Login to Azure for azure-cli authentication
         uses: azure/login@v2
diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml
@@ -37,8 +37,7 @@ jobs:
 
       - name: Run unit tests
         run: |
-          # click 8.3+ introduced bug for hatch
-          pip install "hatch==1.13.0" "click<8.3"
+          pip install hatch==1.15.0
           make test
 
   fmt:
@@ -70,8 +69,7 @@ jobs:
 
       - name: Install Hatch
         run: |
-          # click 8.3+ introduced bug for hatch
-          pip install "hatch==1.13.0" "click<8.3"
+          pip install hatch==1.15.0
 
       - uses: actions/setup-node@v4
         with:
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -27,8 +27,7 @@ jobs:
       
       - name: Build wheels
         run: |
-          # click 8.3+ introduced bug for hatch
-          pip install "hatch==1.13.0" "click<8.3"
+          pip install hatch==1.15.0
           hatch build
       
       - name: Github release
diff --git a/docs/dqx/docs/guide/quality_checks_apply.mdx b/docs/dqx/docs/guide/quality_checks_apply.mdx
@@ -15,6 +15,7 @@ You can apply quality checks to your data using the following approaches:
 
 DQX offers a set of predefined built-in quality rules (checks) as described [here](/docs/guide/quality_checks_definition).
 All DQX rules are evaluated independently and in parallel. If one rule fails, the others are still processed.
+All rule types, including row-level and dataset-level rules, can be applied together in a single execution.
 
 ## Programmatic approach
 
diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx
@@ -12,6 +12,7 @@ This page provides a reference for the quality check functions available in DQX.
 The terms `checks` and `rules` are used interchangeably in the documentation, as they are synonymous in DQX.
 
 All DQX rules are evaluated independently and in parallel. If one rule fails, the others are still processed.
+All rule types, including row-level and dataset-level rules, can be defined and applied simultaneously in a single execution.
 
 You can explore the implementation details of the check functions [here](https://github.com/databrickslabs/dqx/blob/v0.9.3/src/databricks/labs/dqx/check_funcs.py).
 
@@ -1372,7 +1373,7 @@ checks = [
 
 Dataset-level checks are applied to group of rows in a PySpark DataFrame.
 Similar to row-level checks, the results of these quality checks are reported for each individual row in the result columns.
-That means you can define and apply row-level and dataset-level checks at once.
+That means you can define and apply row-level and dataset-level checks together.
 
 Dataset-level checks are useful for validating aggregated values of the entire DataFrame or groups of rows within the DataFrame,
 such as ensuring that the number of rows does not exceed a certain limit or that the sum of a column is within a specified range.
diff --git a/pyproject.toml b/pyproject.toml
@@ -83,7 +83,7 @@ dependencies = [
     "chispa~=0.10.1",
     "coverage[toml]~=7.4.4",
     "databricks-labs-pylint~=0.5",
-    "databricks-labs-pytester~=0.7.2",
+    "databricks-labs-pytester~=0.7.4",
     "mypy~=1.9.0",
     "pylint~=3.3.1",
     "pylint-per-file-ignores~=1.3",
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -1,4 +1,5 @@
 import os
+from datetime import timedelta
 from typing import Any
 import re
 import logging
@@ -8,6 +9,8 @@
 
 import pytest
 from pyspark.sql.types import StructType, StructField, StringType, IntegerType
+from databricks.sdk.errors import BadRequest, NotFound, RequestLimitExceeded, TooManyRequests
+from databricks.sdk.retries import retried
 from databricks.labs.blueprint.installation import Installation, MockInstallation
 from databricks.labs.blueprint.tui import MockPrompts
 from databricks.labs.blueprint.wheels import ProductInfo, WheelsV2
@@ -29,7 +32,7 @@
 
 @pytest.fixture(scope="session")
 def debug_env_name():
-    return "ws"  # Specify the name of the debug environment from ~/.databricks/debug-env.json
+    return "ws2"  # Specify the name of the debug environment from ~/.databricks/debug-env.json
 
 
 @pytest.fixture
@@ -677,47 +680,49 @@ def delete(volume_file_path: str) -> None:
     yield from factory("file", create, delete)
 
 
-@pytest.fixture(scope="session")
-def lakebase_user():
+@pytest.fixture
+def lakebase_user(ws):
     """
-    Get the Lakebase user (service principal client ID) from environment variable.
+    Get a Lakebase user.
 
     This fixture reads ARM_CLIENT_ID which is set in the CI/CD pipeline.
-    For local development, you can set this environment variable to your service principal's client ID.
+    For local development where the ARM_CLIENT_ID is not set, the user is fetched from the workspace client.
 
     Returns:
-        The client ID to use as the Lakebase user.
-
-    Raises:
-        pytest.skip: If ARM_CLIENT_ID is not set in the environment.
+        The client ID or use name to use as the Lakebase user.
     """
     client_id = os.environ.get("ARM_CLIENT_ID")
     if not client_id:
-        pytest.skip("ARM_CLIENT_ID environment variable not set. Required for Lakebase tests.")
+        return ws.current_user.me().user_name
     return client_id
 
 
 @dataclass
 class LakebaseInstance:
     name: str
-    catalog: str
-    location: str
+    catalog_name: str
+    database_name: str
 
 
 @pytest.fixture
 def make_lakebase_instance(ws, make_random):
-
     def create() -> LakebaseInstance:
-        instance_name = f"dqxtest-{make_random(10).lower()}"
+        run_id = os.getenv("GITHUB_RUN_ID", "local")
+        instance_name = f"dqx-test-{run_id}-{make_random(10).lower()}"
         database_name = "dqx"  # does not need to be random
-        catalog_name = f"dqxtest-{make_random(10).lower()}"
-        capacity = "CU_2"
-        schema_name = "config"  # auto-created when saving checks, can be static since each test create new db
-        table_name = "checks"  # auto-created when saving checks, can be static since each test create new db
+        catalog_name = f"dqx-test-{run_id}-{make_random(10).lower()}"
+        capacity = "CU_1"
+
+        # Retry logic handles BadRequest exceptions when database instance creation fails due to workspace quota limits.
+        # Retries are performed until the timeout is reached.
+        @retried(on=[BadRequest, TooManyRequests, RequestLimitExceeded], timeout=timedelta(minutes=20))
+        def _create_database_instance():
+            ws.database.create_database_instance_and_wait(
+                database_instance=DatabaseInstance(name=instance_name, capacity=capacity)
+            )
+
+        _create_database_instance()
 
-        _ = ws.database.create_database_instance_and_wait(
-            database_instance=DatabaseInstance(name=instance_name, capacity=capacity)
-        )
         logger.info(f"Successfully created database instance: {instance_name}")
 
         ws.database.create_database_catalog(
@@ -730,23 +735,28 @@ def create() -> LakebaseInstance:
         )
         logger.info(f"Successfully created database catalog: {catalog_name}")
 
-        return LakebaseInstance(
-            name=instance_name, catalog=catalog_name, location=f"{database_name}.{schema_name}.{table_name}"
-        )
+        return LakebaseInstance(name=instance_name, catalog_name=catalog_name, database_name=database_name)
 
     def delete(instance: LakebaseInstance) -> None:
-        try:
-            ws.database.delete_database_catalog(name=instance.catalog)
-            logger.info(f"Successfully deleted database catalog: {instance.catalog}")
-        except Exception as e:
-            logger.warning(f"Failed to delete database catalog {instance.catalog}: {e}")
-
-        try:
-            ws.database.delete_database_instance(name=instance.name)
-            logger.info(f"Successfully deleted database instance: {instance.name}")
-        except Exception as e:
-            logger.error(f"Failed to delete database instance {instance.name}: {e}")
-            raise
+        # Delete catalog first, then instance.
+        @retried(on=[TooManyRequests, RequestLimitExceeded], timeout=timedelta(minutes=2))
+        def _delete_catalog():
+            try:
+                ws.database.delete_database_catalog(name=instance.catalog_name)
+                logger.info(f"Successfully deleted database catalog: {instance.catalog_name}")
+            except NotFound:
+                logger.info(f"Database catalog {instance.catalog_name} not found (already deleted)")
+
+        @retried(on=[TooManyRequests, RequestLimitExceeded], timeout=timedelta(minutes=2))
+        def _delete_instance():
+            try:
+                ws.database.delete_database_instance(name=instance.name)
+                logger.info(f"Successfully deleted database instance: {instance.name}")
+            except NotFound:
+                logger.info(f"Database instance {instance.name} not found (already deleted)")
+
+        _delete_catalog()
+        _delete_instance()
 
     yield from factory("lakebase", create, delete)
 
diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py
@@ -11,8 +11,6 @@
 logging.getLogger("tests").setLevel("DEBUG")
 logging.getLogger("databricks.labs.dqx").setLevel("DEBUG")
 
-RETRY_INTERVAL_SECONDS = 30
-
 
 @pytest.fixture
 def library_ref() -> str:
diff --git a/tests/e2e/test_run_demos.py b/tests/e2e/test_run_demos.py
@@ -12,7 +12,6 @@
 
 
 logger = logging.getLogger(__name__)
-RETRY_INTERVAL_SECONDS = 30
 
 
 def test_run_dqx_demo_library(make_notebook, make_schema, make_job, library_ref):
diff --git a/tests/integration/test_save_and_load_checks_from_lakebase_table.py b/tests/integration/test_save_and_load_checks_from_lakebase_table.py
