Support Custom Folder Installation for CLI Commands

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 used for opening the remote configuration. 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 used for opening the remote dashboards. 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 used for opening the remote configuration. 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 used for opening the remote configuration. 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 used for opening the remote configuration. 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 used for opening the remote configuration. 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,6 +91,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 used for opening the remote configuration. If not provided, the default installation folder is used.
   - name: workflows
     description: Show deployed workflows and their latest run state.
     table_template: |-
@@ -86,4 +103,6 @@ commands:
     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 used for opening the remote configuration. If not provided, the default installation folder is used.

pyproject.toml

@@ -35,11 +35,10 @@ classifiers = [
     "Topic :: Utilities",
 ]
 # 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",
+dependencies = ["databricks-labs-blueprint>=0.9.1,<0.10",
+                "databricks-sdk~=0.67",
+                "databricks-labs-lsql>=0.5,<=0.16",
+                "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,
@@ -283,16 +293,19 @@ def workflows(w: WorkspaceClient, *, ctx: WorkspaceContext | None = None):
 
 
 @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()