Document for private PYPI and Auto Upload dependencies for workspace installation

docs/dqx/docs/installation.mdx

@@ -66,6 +66,14 @@ pip install 'databricks-labs-dqx[pii]'
 
 This installs additional dependencies including Microsoft Presidio library.
 
+### Installing DQX with company-hosted PyPI mirror
+
+Some enterprises block the public PyPI index and host a company-controlled PyPI mirror. To install DQX while using a company-hosted PyPI mirror for finding its dependencies, add all DQX dependencies to the company-hosted PyPI mirror (see "dependencies" in [`pyproject.toml`](https://github.com/databrickslabs/dqx/blob/v0.10.0/pyproject.toml)) and set the environment variable `PIP_INDEX_URL` to the company-hosted PyPI mirror URL while installing DQX:
+
+```commandline
+PIP_INDEX_URL="https://url-to-company-hosted-pypi.internal" pip install databricks-labs-dqx
+```
+
 ## DQX installation as a Tool in a Databricks Workspace
 
 If you install DQX via PyPI and use it purely as a library, you don’t need to pre-install DQX in the workspace.
@@ -358,6 +366,26 @@ resources:
             #    package: databricks-labs-dqx==0.8.0
 ```
 
+### Installing DQX with company-hosted PyPI mirror
+
+Some enterprises block the public PyPI index and host a company-controlled PyPI mirror. To install DQX while using a company-hosted PyPI mirror for finding its dependencies, add all DQX dependencies to the company-hosted PyPI mirror (see "dependencies" in [`pyproject.toml`](https://github.com/databrickslabs/dqx/blob/v0.10.0/pyproject.toml)) and set the environment variable `PIP_INDEX_URL` to the company-hosted PyPI mirror URL while installing DQX:
+
+```commandline
+PIP_INDEX_URL="https://url-to-company-hosted-pypi.internal" databricks labs install dqx
+```
+
+During DQX installation as a tool in the workspace reply *yes* to the question "Does the given workspace block Internet access"?
+
+If the installation host has no access to GitHub, then dqx installation will not be able to download the files locally and will fail.
+In order to address that, follow the steps below
+ - install dqx on a host which has access to github using the instruction above
+ - zip the installation from ~/.databricks/labs/dqx
+ - copy the zip file to the target host and unzip
+Now the installation can be done in offline mode (ensure Databricks CLI is upgraded to version v0.244.0 or higher)
+```commandline
+PIP_INDEX_URL="https://url-to-company-hosted-pypi.internal" databricks labs install dqx --offline=true
+```
+
 ### Upgrade DQX in the Databricks workspace
 
 Verify that DQX is installed:

docs/dqx/docs/reference/benchmarks.mdx

@@ -9,9 +9,11 @@ sidebar_position: 13
 # Performance Benchmarks Report
 
 ## Specification
+
 * 100 million rows are used for each test.
-* Tests are run sequentially to reduce variability.
+* DQX rules are executed in parallel and in distributed manner using Spark. To minimize test variability, each type of check is executed sequentially in this benchmark. Specific tests, such as `test_benchmark_apply_checks_all_dataset_checks` and `test_benchmark_apply_checks_all_row_checks`, execute all types of checks in parallel simultaneously.
 * Benchmarks are created using Databricks Serverless cluster.
+* The provided benchmarks are indicative. You should always consider benchmarking results in the context of your own data and environment.
 
 ## Results
 | Test | Mean (s) | Median (s) | Min (s) | Max (s) | Stddev (s) | iqr (s) | q1 (s) | q3 (s) | Rounds | iqr outliers | stddev outliers | Ops/s |

pyproject.toml

@@ -34,10 +34,12 @@ classifiers = [
     "Topic :: Software Development :: Libraries",
     "Topic :: Utilities",
 ]
-dependencies = ["databricks-labs-blueprint>=0.9.1,<0.10",
-                "databricks-sdk~=0.71",
-                "databricks-labs-lsql>=0.5,<=0.16",
-                "sqlalchemy>=1.4,<3.0",
+# also update fallback dependencies in workflow installer when changing dependencies!
+dependencies = [
+    "databricks-labs-blueprint>=0.9.1,<0.10",
+    "databricks-sdk~=0.71",
+    "databricks-labs-lsql>=0.5,<=0.16",
+    "sqlalchemy>=1.4,<3.0",
 ]
 
 [project.optional-dependencies]

src/databricks/labs/dqx/config.py

@@ -141,6 +141,9 @@ class WorkspaceConfig:
 
     # whether to use serverless clusters for the jobs, only used during workspace installation
     serverless_clusters: bool = True
+    upload_dependencies: bool = (
+        False  # whether to upload dependencies to the workspace during installation to enable DQX in restricted (no-internet) environments
+    )
     extra_params: ExtraParams | None = None  # extra parameters to pass to the jobs, e.g. result_column_names
 
     # cluster configuration for the jobs (applicable for non-serverless clusters only)

src/databricks/labs/dqx/engine.py

@@ -1168,6 +1168,7 @@ def get_streaming_metrics_listener(
         )
         return StreamingMetricsListener(metrics_config, metrics_observation, self.spark, target_query_id)
 
+    @telemetry_logger("engine", "apply_checks_for_run_config")
     def _apply_checks_for_run_config(self, run_config: RunConfig) -> None:
         """
         Applies checks based on a given RunConfig.

src/databricks/labs/dqx/installer/config_provider.py

@@ -98,6 +98,9 @@ def prompt_new_installation(self, install_folder: str | None = None) -> Workspac
 
         warehouse_id = self._warehouse_configurator.create()
 
+        # Ask if the workspace blocks Internet access to determine if dependencies should be uploaded
+        upload_dependencies = self._prompts.confirm("Does the given workspace block Internet access?")
+
         return WorkspaceConfig(
             log_level=log_level,
             run_configs=[
@@ -114,6 +117,7 @@ def prompt_new_installation(self, install_folder: str | None = None) -> Workspac
                 )
             ],
             serverless_clusters=serverless_clusters,
+            upload_dependencies=upload_dependencies,
             profiler_spark_conf=profiler_spark_conf,
             profiler_override_clusters=profiler_override_clusters,
             quality_checker_spark_conf=quality_checker_spark_conf,

src/databricks/labs/dqx/installer/workflow_installer.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import importlib.metadata
 import logging
 import os.path
 import re
@@ -487,9 +488,82 @@ def _library_dep_order(library: str):
             case _:
                 return 2
 
+    @staticmethod
+    def _extract_dependency_prefix(requirement: str) -> str | None:
+        """
+        Extract package name prefix from a requirement string.
+
+        Args:
+            requirement: Requirement string (e.g., "databricks-sdk>=0.71")
+
+        Returns:
+            Package name prefix with underscores, or None if invalid.
+        """
+        match = re.match(r'^([a-zA-Z0-9-]+)', requirement)
+        if match:
+            pkg_name = match.group(1)
+            return pkg_name.replace('-', '_')
+        return None
+
+    @staticmethod
+    def _get_fallback_dependencies() -> list[str]:
+        """
+        Get fallback dependency prefixes when metadata is unavailable. The list should match the dependency list
+        from the pyproject.toml file in the project root.
+
+        Returns:
+            List of core dependency prefixes.
+        """
+        return [
+            "databricks_labs_blueprint",
+            "databricks_sdk",
+            "databricks_labs_lsql",
+            "sqlalchemy",
+        ]
+
+    @staticmethod
+    def _get_dependency_prefixes() -> list[str]:
+        """
+        Dynamically retrieve dependency prefixes from package metadata.
+
+        Includes both core and optional (extras) dependencies to ensure workflows
+        have access to all required packages.
+
+        Returns:
+            List of dependency package name prefixes for wheel files.
+        """
+        try:
+            requires = importlib.metadata.requires('databricks-labs-dqx')
+        except importlib.metadata.PackageNotFoundError:
+            logger.warning("databricks-labs-dqx package metadata not found, using fallback dependencies")
+            return WorkflowDeployment._get_fallback_dependencies()
+
+        if not requires:
+            logger.warning("No dependencies found in package metadata")
+            return []
+
+        prefixes = []
+        for req in requires:
+            prefix = WorkflowDeployment._extract_dependency_prefix(req)
+            if prefix:
+                prefixes.append(prefix)
+
+        # Remove duplicates while preserving order
+        unique_prefixes = list(dict.fromkeys(prefixes))
+        logger.info(f"Discovered {len(unique_prefixes)} dependencies from package metadata (including extras)")
+        return unique_prefixes
+
     def _upload_wheel(self):
         wheel_paths = []
         with self._wheels:
+            # Upload dependencies if workspace blocks Internet access
+            if self._config.upload_dependencies:
+                logger.info("Uploading dependencies to workspace...")
+                dependency_prefixes = self._get_dependency_prefixes()
+                # TODO the _build_wheel in the upload wheel dependencies method
+                #  currently does not handle installation of extras, therefore they are not uploaded
+                for whl in self._wheels.upload_wheel_dependencies(dependency_prefixes):
+                    wheel_paths.append(whl)
             wheel_paths.sort(key=WorkflowDeployment._library_dep_order)
             wheel_paths.append(self._wheels.upload_to_wsfs())
             wheel_paths = [f"/Workspace{wheel}" for wheel in wheel_paths]

tests/e2e/test_run_demos.py

@@ -142,9 +142,9 @@ def test_run_dqx_demo_tool(installation_ctx, make_schema, make_notebook, make_jo
     schema = make_schema(catalog_name=catalog).name
     installation_ctx.replace(
         extend_prompts={
-            r'Provide location for the input data .*': '/databricks-datasets/delta-sharing/samples/nyctaxi_2019',
-            r'Provide output table .*': f'{catalog}.{schema}.output_table',
-            r'Provide quarantined table .*': f'{catalog}.{schema}.quarantine_table',
+            r"Provide location for the input data .*": "/databricks-datasets/delta-sharing/samples/nyctaxi_2019",
+            r"Provide output table .*": f"{catalog}.{schema}.output_table",
+            r"Provide quarantined table .*": f"{catalog}.{schema}.quarantine_table",
         },
     )
     installation_ctx.workspace_installer.run(installation_ctx.config)

tests/integration/test_ai_rules_generator.py

@@ -9,25 +9,11 @@
 
 # Sample user input with specific requirements to avoid flakiness in tests
 USER_INPUT = """
-Username should not start with 's' and should not contain more than 20 letters if user id is provided. Error message must be: "Username should not start with 's' and should not contain more than 20 letters if user id is provided"
-Apply validation when user_id is not null. Use SQL expression for this check with the following expression: NOT (username LIKE 's%') AND LENGTH(username) <= 20.
-Users at age 18 or above must have a valid email address.
-Age should be between 0 and 120.
+Users at age 18 or above must have a valid email address checked using regex.
+Age should be between 0 and 120. Output the rules in the given order.
 """
 
 EXPECTED_CHECKS = [
-    {
-        "check": {
-            "arguments": {
-                "columns": ["username"],
-                "expression": "NOT (username LIKE 's%') AND LENGTH(username) <= 20",
-                "msg": "Username should not start with 's' and should not contain more than 20 letters if user id is provided",
-            },
-            "function": "sql_expression",
-        },
-        "criticality": "error",
-        "filter": "user_id IS NOT NULL",
-    },
     {
         "check": {
             "arguments": {"column": "email", "regex": "^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}$"},

tests/integration/test_installation.py

@@ -42,9 +42,9 @@ def factory(
 
         prompts = MockPrompts(
             {
-                r'Provide location for the input data .*': '/',
-                r'Provide output table .*': 'main.dqx_test.output_table',
-                r'Do you want to uninstall DQX .*': 'yes',
+                r"Provide location for the input data .*": "/",
+                r"Provide output table .*": "main.dqx_test.output_table",
+                r"Do you want to uninstall DQX .*": "yes",
                 r".*PRO or SERVERLESS SQL warehouse.*": "1",
                 r".*": "",
             }