Support Custom Folder Installation for CLI Commands

docs/dqx/docs/reference/cli.mdx

@@ -26,7 +26,9 @@ databricks labs uninstall dqx
 <Admonition type="info" title="Where is DQX installed?">
 By default, DQX is installed under the user's home folder (for example `/Users/<user>/.dqx`).
 Use the `DQX_FORCE_INSTALL` env var to force a global or user install. Provide a custom installation
-folder during installation to override the default location. See the installation guide for details.
+folder during installation to override the default location. See the installation guide for details. 
+Set the `--install-folder` argument When running commands for a DQX installation with a custom 
+installation folder.
 </Admonition>
 
 ## Configuration helpers
@@ -35,8 +37,14 @@ folder during installation to override the default location. See the installatio
 # Open the installation config in the workspace
 databricks labs dqx open-remote-config
 
+# Open the installation config in the workspace with a custom installation folder
+databricks labs dqx open-remote-config --install-folder "/Workspace/my_dqx_folder"
+
 # Open dashboards folder for the installation
 databricks labs dqx open-dashboards
+
+# Open dashboards folder for an installation with a custom installation folder
+databricks labs dqx open-dashboards --install-folder "/Workspace/my_dqx_folder"
 ```
 
 ## Workflows and logs
@@ -45,8 +53,14 @@ databricks labs dqx open-dashboards
 # List installed workflows with their latest run state
 databricks labs dqx workflows
 
+# List installed workflows with their latest run state with a custom installation folder
+databricks labs dqx workflows --install-folder "/Workspace/my_dqx_folder"
+
 # Show logs for the latest run of a workflow (profiler, quality-checker, e2e)
 databricks labs dqx logs --workflow quality-checker
+
+# Show logs for the latest run of a workflow (profiler, quality-checker, e2e) with a custom installation folder
+databricks labs dqx logs --workflow quality-checker --install-folder "/Workspace/my_dqx_folder"
 ```
 
 ## Running workflows
@@ -55,6 +69,9 @@ databricks labs dqx logs --workflow quality-checker
 # Run profiler for all run configs in the configuration file: profile data and generate quality check candidates
 databricks labs dqx profile --timeout-minutes 20
 
+# Run profiler for all run configs in the configuration file with a custom installation folder: profile data and generate quality check candidates
+databricks labs dqx profile --timeout-minutes 20 --install-folder "/Workspace/my_dqx_folder"
+
 # Run profiler for a single run config in the configuration file: profile data and generate quality check candidates
 databricks labs dqx profile --run-config default --timeout-minutes 20
 
@@ -72,6 +89,9 @@ databricks labs dqx profile --run-config "default" --patterns "main.product001.*
 # Run quality checker for all run configs in the configuration file: apply quality checks and write results and optionally quarantine
 databricks labs dqx apply-checks --timeout-minutes 20
 
+# Run quality checker for all run configs in the configuration file with a custom installation folder: apply quality checks and write results and optionally quarantine
+databricks labs dqx apply-checks --timeout-minutes 20 --install-folder "/Workspace/my_dqx_folder"
+
 # Run quality checker for a single run config in the configuration file: apply quality checks and write results and optionally quarantine
 databricks labs dqx apply-checks --run-config default --timeout-minutes 20
 
@@ -81,6 +101,9 @@ databricks labs dqx apply-checks --run-config "default" --patterns "main.product
 # Run e2e (end-to-end) workflows for all run configs in the configuration file: profile data > generate quality checks > apply checks > write results and optionally quarantine
 databricks labs dqx e2e --timeout-minutes 20
 
+# Run e2e (end-to-end) workflows for all run configs in the configuration file with a custom installation folder: profile data > generate quality checks > apply checks > write results and optionally quarantine
+databricks labs dqx e2e --timeout-minutes 20 --install-folder "/Workspace/my_dqx_folder"
+
 # Run e2e (end-to-end) workflows for a single run config in the configuration file: profile data > generate quality checks > apply checks > write results and optionally quarantine
 databricks labs dqx e2e --run-config default --timeout-minutes 20
 
@@ -94,6 +117,9 @@ databricks labs dqx e2e --run-config "default" --patterns "main.product001.*;mai
 # Validate checks stored in the installation (defined in `checks_location`) for a single run config in the configuration file:
 databricks labs dqx validate-checks
 
+# Validate checks stored in a custom installation folder (defined in `checks_location`) for a single run config in the configuration file:
+databricks labs dqx validate-checks --install-folder "/Workspace/my_dqx_folder"
+
 # Validate checks stored in the installation (defined in `checks_location`) for all run configs in the configuration file:
 databricks labs dqx validate-checks --run-config default
 ```

labs.yml

@@ -11,14 +11,23 @@ min_python: 3.10
 commands:
   - name: open-remote-config
     description: Opens remote configuration in the browser.
+    flags:
+      - name: install-folder
+        description: (Optional) Custom installation folder. If not provided, the default installation folder is used.
   - name: open-dashboards
     description: Opens remote dashboards directory.
+    flags:
+      - name: install-folder
+        description: (Optional) Custom installation folder. If not provided, the default installation folder is used.
   - name: installations
     description: Show installations by different users on the same workspace.
     table_template: |-
       Path\tVersion\Input
       {{range .}}{{.path}}\t{{.version}}\t{{.input}}
       {{end}}
+    flags:
+      - name: install-folder
+        description: (Optional) Custom installation folder. If not provided, the default installation folder is used.
   - name: validate-checks
     description: Validate checks
     flags:
@@ -27,6 +36,8 @@ commands:
       - name: validate-custom-check-functions
         description: Whether to validate custom check functions.
         default: true
+      - name: install-folder
+        description: (Optional) Custom installation folder. If not provided, the default installation folder is used.
     table_template: |-
       Run Config\tErrors
       {{range .}}{{.run_config}}\t{{.error}}
@@ -42,6 +53,8 @@ commands:
         description: (Optional) Semicolon separated list of location patterns to exclude from profiling. Useful if wanting to exclude existing output tables.
       - name: timeout-minutes
         description: The timeout in minutes for the cli to wait for the workflow to complete.
+      - name: install-folder
+        description: (Optional) Custom installation folder. If not provided, the default installation folder is used.
   - name: apply-checks
     description: Apply data quality checks to the input data, and save the results to the output table and optionally quarantine table (based on the run config).
     flags:
@@ -59,6 +72,8 @@ commands:
         default: "_dq_quarantine"
       - name: timeout-minutes
         description: The timeout in minutes for the cli to wait for the workflow to complete.
+      - name: install-folder
+        description: (Optional) Custom installation folder. If not provided, the default installation folder is used.
   - name: e2e
     description: Run end to end workflow to profile input data and apply quality checks, and save the results to the output table and optionally quarantine table (based on the run config).
     flags:
@@ -76,14 +91,21 @@ commands:
         default: "_dq_quarantine"
       - name: timeout-minutes
         description: The timeout in minutes for the cli to wait for the workflow to complete.
+      - name: install-folder
+        description: (Optional) Custom installation folder. If not provided, the default installation folder is used.
   - name: workflows
     description: Show deployed workflows and their latest run state.
     table_template: |-
       Workflow\tWorkflow ID\tState\tStarted
       {{range .}}{{.workflow}}\t{{.workflow_id}}\t{{.state}}\t{{.started}}
       {{end}}
+    flags:
+      - name: install-folder
+        description: (Optional) Custom installation folder. If not provided, the default installation folder is used.
   - name: logs
     description: Show logs from the latest job run
     flags:
       - name: workflow
-        description: Name of the workflow to show logs for, e.g. profiler.
+        description: Name of the workflow to show logs for, e.g. profiler.
+      - name: install-folder
+        description: (Optional) Custom installation folder. If not provided, the default installation folder is used.

pyproject.toml

@@ -39,7 +39,7 @@ 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",
+    "sqlalchemy>=2.0,<3.0",
 ]
 
 [project.optional-dependencies]

src/databricks/labs/dqx/cli.py

@@ -23,29 +23,31 @@
 
 
 @dqx.command
-def open_remote_config(w: WorkspaceClient, *, ctx: WorkspaceContext | None = None):
+def open_remote_config(w: WorkspaceClient, *, install_folder: str = "", ctx: WorkspaceContext | None = None):
     """
     Opens remote configuration in the browser.
 
     Args:
         w: The WorkspaceClient instance to use for accessing the workspace.
+        install_folder: Optional custom installation folder path.
         ctx: The WorkspaceContext instance to use for accessing the workspace.
     """
-    ctx = ctx or WorkspaceContext(w)
+    ctx = ctx or WorkspaceContext(w, install_folder=install_folder or None)
     workspace_link = ctx.installation.workspace_link(WorkspaceConfig.__file__)
     webbrowser.open(workspace_link)
 
 
 @dqx.command
-def open_dashboards(w: WorkspaceClient, *, ctx: WorkspaceContext | None = None):
+def open_dashboards(w: WorkspaceClient, *, install_folder: str = "", ctx: WorkspaceContext | None = None):
     """
     Opens remote dashboard directory in the browser.
 
     Args:
         w: The WorkspaceClient instance to use for accessing the workspace.
+        install_folder: Optional custom installation folder path.
         ctx: The WorkspaceContext instance to use for accessing the workspace.
     """
-    ctx = ctx or WorkspaceContext(w)
+    ctx = ctx or WorkspaceContext(w, install_folder=install_folder or None)
     workspace_link = ctx.installation.workspace_link("")
     webbrowser.open(f"{workspace_link}dashboards/")
 
@@ -85,6 +87,7 @@ def validate_checks(
     *,
     run_config: str = "",
     validate_custom_check_functions: bool = True,
+    install_folder: str = "",
     ctx: WorkspaceContext | None = None,
 ) -> list[dict]:
     """
@@ -94,9 +97,10 @@ def validate_checks(
         w: The WorkspaceClient instance to use for accessing the workspace.
         run_config: The name of the run configuration to use. If not provided, run it for all run configs.
         validate_custom_check_functions: Whether to validate custom check functions (default is True).
+        install_folder: Optional custom installation folder path.
         ctx: The WorkspaceContext instance to use for accessing the workspace.
     """
-    ctx = ctx or WorkspaceContext(w)
+    ctx = ctx or WorkspaceContext(w, install_folder=install_folder or None)
     config = ctx.installation.load(WorkspaceConfig)
 
     errors_list = []
@@ -154,6 +158,7 @@ def profile(
     patterns: str = "",
     exclude_patterns: str = "",
     timeout_minutes: int = 30,
+    install_folder: str = "",
     ctx: WorkspaceContext | None = None,
 ) -> None:
     """
@@ -168,10 +173,11 @@ def profile(
         exclude_patterns: Semicolon-separated list of location patterns to exclude.
             Useful to skip existing output and quarantine tables based on suffixes.
         timeout_minutes: The timeout for the workflow run in minutes (default is 30).
+        install_folder: Optional custom installation folder path.
         ctx: The WorkspaceContext instance to use for accessing the workspace.
     """
     timeout = timedelta(minutes=timeout_minutes)
-    ctx = ctx or WorkspaceContext(w)
+    ctx = ctx or WorkspaceContext(w, install_folder=install_folder or None)
     ctx.deployed_workflows.run_workflow(
         workflow="profiler",
         run_config_name=run_config,
@@ -191,6 +197,7 @@ def apply_checks(
     output_table_suffix: str = "_dq_output",
     quarantine_table_suffix: str = "_dq_quarantine",
     timeout_minutes: int = 30,
+    install_folder: str = "",
     ctx: WorkspaceContext | None = None,
 ) -> None:
     """
@@ -207,10 +214,11 @@ def apply_checks(
         output_table_suffix: Suffix to append to the output table names (default is "_dq_output").
         quarantine_table_suffix: Suffix to append to the quarantine table names (default is "_dq_quarantine").
         timeout_minutes: The timeout for the workflow run in minutes (default is 30).
+        install_folder: Optional custom installation folder path.
         ctx: The WorkspaceContext instance to use for accessing the workspace.
     """
     timeout = timedelta(minutes=timeout_minutes)
-    ctx = ctx or WorkspaceContext(w)
+    ctx = ctx or WorkspaceContext(w, install_folder=install_folder or None)
     ctx.deployed_workflows.run_workflow(
         workflow="quality-checker",
         run_config_name=run_config,
@@ -232,6 +240,7 @@ def e2e(
     output_table_suffix: str = "_dq_output",
     quarantine_table_suffix: str = "_dq_quarantine",
     timeout_minutes: int = 60,
+    install_folder: str = "",
     ctx: WorkspaceContext | None = None,
 ) -> None:
     """
@@ -251,10 +260,11 @@ def e2e(
         output_table_suffix: Suffix to append to the output table names (default is "_dq_output").
         quarantine_table_suffix: Suffix to append to the quarantine table names (default is "_dq_quarantine").
         timeout_minutes: The timeout for the workflow run in minutes (default is 60).
+        install_folder: Optional custom installation folder path.
         ctx: The WorkspaceContext instance to use for accessing the workspace.
     """
     timeout = timedelta(minutes=timeout_minutes)
-    ctx = ctx or WorkspaceContext(w)
+    ctx = ctx or WorkspaceContext(w, install_folder=install_folder or None)
     ctx.deployed_workflows.run_workflow(
         workflow="e2e",
         run_config_name=run_config,
@@ -267,32 +277,36 @@ def e2e(
 
 
 @dqx.command
-def workflows(w: WorkspaceClient, *, ctx: WorkspaceContext | None = None):
+def workflows(w: WorkspaceClient, *, ctx: WorkspaceContext | None = None, install_folder: str = ""):
     """
     Show deployed workflows and their state
 
     Args:
         w: The WorkspaceClient instance to use for accessing the workspace.
         ctx: The WorkspaceContext instance to use for accessing the workspace.
+        install_folder: Optional custom installation folder path.
     """
-    ctx = ctx or WorkspaceContext(w)
+    ctx = ctx or WorkspaceContext(w, install_folder=install_folder or None)
     logger.info("Fetching deployed jobs...")
     latest_job_status = ctx.deployed_workflows.latest_job_status()
     print(json.dumps(latest_job_status))
     return latest_job_status
 
 
 @dqx.command
-def logs(w: WorkspaceClient, *, workflow: str | None = None, ctx: WorkspaceContext | None = None):
+def logs(
+    w: WorkspaceClient, *, workflow: str | None = None, install_folder: str = "", ctx: WorkspaceContext | None = None
+):
     """
     Show logs of the latest job run.
 
     Args:
         w: The WorkspaceClient instance to use for accessing the workspace.
         workflow: The name of the workflow to show logs for.
+        install_folder: Optional custom installation folder path.
         ctx: The WorkspaceContext instance to use for accessing the workspace
     """
-    ctx = ctx or WorkspaceContext(w)
+    ctx = ctx or WorkspaceContext(w, install_folder=install_folder or None)
     ctx.deployed_workflows.relay_logs(workflow)
 
 

src/databricks/labs/dqx/contexts/global_context.py

@@ -15,10 +15,11 @@ class GlobalContext(abc.ABC):
     GlobalContext class that provides a global context, including workspace client,
     """
 
-    def __init__(self, named_parameters: dict[str, str] | None = None):
+    def __init__(self, named_parameters: dict[str, str] | None = None, install_folder: str | None = None):
         if not named_parameters:
             named_parameters = {}
         self._named_parameters = named_parameters
+        self._install_folder = install_folder
 
     def replace(self, **kwargs):
         """
@@ -48,6 +49,10 @@ def product_info(self) -> ProductInfo:
 
     @cached_property
     def installation(self) -> Installation:
+        if self._install_folder:
+            return Installation(
+                self.workspace_client, self.product_info.product_name(), install_folder=self._install_folder
+            )
         return Installation.current(self.workspace_client, self.product_info.product_name())
 
     @cached_property

src/databricks/labs/dqx/contexts/workspace_context.py

@@ -10,8 +10,10 @@ class WorkspaceContext(CliContext):
     WorkspaceContext class that extends CliContext to provide workspace-specific functionality.
     """
 
-    def __init__(self, ws: WorkspaceClient, named_parameters: dict[str, str] | None = None):
-        super().__init__(named_parameters)
+    def __init__(
+        self, ws: WorkspaceClient, named_parameters: dict[str, str] | None = None, install_folder: str | None = None
+    ):
+        super().__init__(named_parameters, install_folder)
         self._ws = ws
 
     @cached_property

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

@@ -50,7 +50,7 @@ class WorkspaceInstaller(WorkspaceContext, InstallationMixin):
     """
 
     def __init__(self, ws: WorkspaceClient, environ: dict[str, str] | None = None, install_folder: str | None = None):
-        super().__init__(ws)
+        super().__init__(ws, install_folder=install_folder)
         if not environ:
             environ = dict(os.environ.items())
 

tests/integration/conftest.py

@@ -183,6 +183,9 @@ def setup_workflows_with_custom_folder(
     Set up the workflows with installation in the custom install folder.
     """
 
+    if os.getenv("DATABRICKS_SERVERLESS_COMPUTE_ID"):
+        pytest.skip()
+
     def create(_spark, **kwargs):
         installation_ctx_custom_install_folder.installation_service.run()