Allow installation in a custom folder

.github/workflows/nightly.yml

@@ -97,7 +97,6 @@ jobs:
           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:
@@ -147,7 +146,6 @@ jobs:
           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:

docs/dqx/docs/installation.mdx

@@ -98,18 +98,29 @@ It is recommended to use serverless clusters for the workflows, as it allows for
 If serverless clusters are not used, a default cluster configuration will be used for the workflows.
 Alternatively, you can override the cluster configuration for each workflow in the `config.yml` file after the installation, to provide existing clusters to use.
 
-#### User vs Global Installation
+#### Installation Options
 
-DQX is installed by default in the user home directory (under `/Users/<user>/.dqx`). You can also install DQX globally
-by setting the 'DQX_FORCE_INSTALL' environment variable. The following options are available:
+DQX offers flexible installation options. By default, DQX is installed in the user home directory (under `/Users/<user>/.dqx`).
 
-* `DQX_FORCE_INSTALL=global databricks labs install dqx`: will force the installation to be for root only (`/Applications/dqx`)
-* `DQX_FORCE_INSTALL=user databricks labs install dqx`: will force the installation to be for user only (`/Users/<user>/.dqx`)
+**Custom Workspace Folder:**
+If you provide a custom path during installation (e.g., `/Shared/dqx-team` or `/Users/shared-user/dqx-project`), DQX will be installed there.
+
+**Environment Variable Override:**
+You can also force global or user installation using the 'DQX_FORCE_INSTALL' environment variable:
+  * `DQX_FORCE_INSTALL=global databricks labs install dqx`: forces installation to `/Applications/dqx`
+  * `DQX_FORCE_INSTALL=user databricks labs install dqx`: forces installation to `/Users/<user>/.dqx` (default behaviour)
+
+<Admonition type="info" title="Complete Profiling Guide">
+If a custom folder is provided during installation, the installation folder will take precedence over any environment variables (e.g. `DQX_FORCE_INSTALL`).
+</Admonition> 
 
 #### Configuration file
 
 DQX configuration file can contain multiple run configurations for different pipelines or projects, each defining specific input, output and quarantine locations, etc.
-By default, the config is created in the installation directory under `/Users/<user>/.dqx/config.yml` or `/Applications/dqx/config.yml` if installed globally.
+The configuration file is created in the installation directory depending on installation options (see above):
+- User home (default): `/Users/<user>/.dqx/config.yml`
+- Global: `/Applications/dqx/config.yml` (when DQX_FORCE_INSTALL=global)
+- Custom folder: `<your-installation-folder>/config.yml`
 The "default" run configuration is created during the installation. When DQX is upgraded, the configuration is preserved.
 The configuration can be updated / extended manually by the user after the installation. Each run config defines configuration for one specific input and output location.
 

docs/dqx/docs/reference/cli.mdx

@@ -24,8 +24,9 @@ databricks labs uninstall dqx
 ```
 
 <Admonition type="info" title="Where is DQX installed?">
-By default, DQX is installed under the user's home (for example `/Users/<user>/.dqx`).
-Use the `DQX_FORCE_INSTALL` env var to force a global or user install. See the installation guide for details.
+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.
 </Admonition>
 
 ## Configuration helpers

src/databricks/labs/dqx/checks_storage.py

@@ -262,9 +262,11 @@ def _get_storage_handler_and_config(
         self, config: InstallationChecksStorageConfig
     ) -> tuple[ChecksStorageHandler, InstallationChecksStorageConfig]:
         run_config = self._run_config_loader.load_run_config(
-            config.run_config_name, config.assume_user, config.product_name
+            run_config_name=config.run_config_name, assume_user=config.assume_user, product_name=config.product_name
+        )
+        installation = self._run_config_loader.get_installation(
+            config.assume_user, config.product_name, config.install_folder
         )
-        installation = self._run_config_loader.get_installation(config.assume_user, config.product_name)
 
         config.location = run_config.checks_location
 

src/databricks/labs/dqx/config.py

@@ -219,9 +219,15 @@ class InstallationChecksStorageConfig(
         run_config_name: The name of the run configuration to use for checks (default is 'default').
         product_name: The product name for retrieving checks from the installation (default is 'dqx').
         assume_user: Whether to assume the user is the owner of the checks (default is True).
+        install_folder: The installation folder where DQX is installed.
+        Default to /Users/<your_user>/.dqx if not provided.
     """
 
     location: str = "installation"  # retrieved from the installation config
     run_config_name: str = "default"  # to retrieve run config
     product_name: str = "dqx"
     assume_user: bool = True
+    # DQX will be installed in a default directory if no custom folder is provided:
+    # * Default: "/Users/<your_user>/.dqx"
+    # * if `DQX_FORCE_INSTALL=global`: "/Applications/dqx"
+    install_folder: str | None = None

src/databricks/labs/dqx/config_loader.py

@@ -1,5 +1,6 @@
 from databricks.labs.blueprint.installation import Installation
 from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import NotFound
 
 from databricks.labs.dqx.config import RunConfig, WorkspaceConfig
 
@@ -13,36 +14,67 @@ def __init__(self, workspace_client: WorkspaceClient):
         self.ws = workspace_client
 
     def load_run_config(
-        self, run_config_name: str | None, assume_user: bool = True, product_name: str = "dqx"
+        self,
+        run_config_name: str | None,
+        install_folder: str | None = None,
+        assume_user: bool = True,
+        product_name: str = "dqx",
     ) -> RunConfig:
         """
         Load run configuration from the installation.
 
         Args:
             run_config_name: name of the run configuration to use
+            install_folder: optional installation folder
             assume_user: if True, assume user installation
             product_name: name of the product
         """
-        installation = self.get_installation(assume_user, product_name)
+        installation = self.get_installation(assume_user, product_name, install_folder)
         return self._load_run_config(installation, run_config_name)
 
-    def get_installation(self, assume_user: bool, product_name: str) -> Installation:
+    def get_installation(self, assume_user: bool, product_name: str, install_folder: str | None = None) -> Installation:
         """
         Get the installation for the given product name.
 
         Args:
             assume_user: if True, assume user installation
             product_name: name of the product
+            install_folder: optional installation folder
         """
+
+        if install_folder:
+            installation = self.get_custom_installation(self.ws, product_name, install_folder)
+            return installation
+
         if assume_user:
             installation = Installation.assume_user_home(self.ws, product_name)
         else:
             installation = Installation.assume_global(self.ws, product_name)
 
-        # verify the installation
         installation.current(self.ws, product_name, assume_user=assume_user)
         return installation
 
+    @staticmethod
+    def get_custom_installation(ws: WorkspaceClient, product_name: str, install_folder: str) -> Installation:
+        """
+        Creates an Installation instance for a custom folder, similar to assume_user_home and assume_global.
+        This ensures the custom folder is created in the workspace when the installation is accessed.
+
+        Args:
+            ws: Databricks SDK `WorkspaceClient`
+            product_name: The product name
+            install_folder: The custom installation folder path
+
+        Returns:
+            An Installation instance for the custom folder
+        """
+        try:
+            ws.workspace.get_status(install_folder)
+        except NotFound:
+            ws.workspace.mkdirs(install_folder)
+
+        return Installation(ws, product_name, install_folder=install_folder)
+
     @staticmethod
     def _load_run_config(installation: Installation, run_config_name: str | None) -> RunConfig:
         """

src/databricks/labs/dqx/engine.py

@@ -629,11 +629,15 @@ def save_results_in_table(
             None
         """
         if output_df is not None and output_config is None:
-            run_config = self._run_config_loader.load_run_config(run_config_name, assume_user, product_name)
+            run_config = self._run_config_loader.load_run_config(
+                run_config_name=run_config_name, assume_user=assume_user, product_name=product_name
+            )
             output_config = run_config.output_config
 
         if quarantine_df is not None and quarantine_config is None:
-            run_config = self._run_config_loader.load_run_config(run_config_name, assume_user, product_name)
+            run_config = self._run_config_loader.load_run_config(
+                run_config_name=run_config_name, assume_user=assume_user, product_name=product_name
+            )
             quarantine_config = run_config.quarantine_config
 
         if output_df is not None and output_config is not None:

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

@@ -1,3 +1,4 @@
+import os
 import json
 import logging
 from databricks.labs.blueprint.tui import Prompts
@@ -17,12 +18,21 @@ def __init__(self, prompts: Prompts, warehouse_configurator: WarehouseInstaller)
         self._prompts = prompts
         self._warehouse_configurator = warehouse_configurator
 
-    def prompt_new_installation(self) -> WorkspaceConfig:
+    def prompt_new_installation(self, install_folder: str | None = None) -> WorkspaceConfig:
         logger.info(
             "Please answer a couple of questions to provide default DQX run configuration. "
             "The configuration can also be updated manually after the installation."
         )
 
+        # Show installation folder information
+        if install_folder:
+            logger.info(f"DQX will be installed in folder '{install_folder}'")
+        else:
+            install_path = (
+                "/Applications/dqx" if os.getenv("DQX_FORCE_INSTALL") == "global" else "/Users/<your_user>/.dqx"
+            )
+            logger.info(f"DQX will be installed in the default location: '{install_path}'")
+
         log_level = self._prompts.question("Log level", default="INFO").upper()
         is_streaming = self._prompts.confirm("Should the input data be read using streaming?")
         input_config = self._prompt_input_config(is_streaming)

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

@@ -13,7 +13,7 @@
 from databricks.labs.blueprint.parallel import ManyError, Threads
 from databricks.labs.blueprint.tui import Prompts
 from databricks.labs.blueprint.upgrades import Upgrades
-from databricks.labs.blueprint.wheels import ProductInfo
+from databricks.labs.blueprint.wheels import ProductInfo, WheelsV2
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.core import with_user_agent_extra
 from databricks.sdk.errors import (
@@ -22,6 +22,7 @@
     PermissionDenied,
 )
 
+from databricks.labs.dqx.config_loader import RunConfigLoader
 from databricks.labs.dqx.installer.config_provider import ConfigProvider
 from databricks.labs.dqx.installer.dashboard_installer import DashboardInstaller
 from databricks.labs.dqx.installer.version_checker import VersionChecker
@@ -45,14 +46,16 @@ class WorkspaceInstaller(WorkspaceContext):
     Args:
         environ: Optional dictionary of environment variables.
         ws: The WorkspaceClient instance.
+        install_folder: Optional custom workspace folder path for installation.
     """
 
-    def __init__(self, ws: WorkspaceClient, environ: dict[str, str] | None = None):
+    def __init__(self, ws: WorkspaceClient, environ: dict[str, str] | None = None, install_folder: str | None = None):
         super().__init__(ws)
         if not environ:
             environ = dict(os.environ.items())
 
         self._force_install = environ.get("DQX_FORCE_INSTALL")
+        self._install_folder = install_folder
 
         if "DATABRICKS_RUNTIME_VERSION" in environ:
             msg = "WorkspaceInstaller is not supposed to be executed in Databricks Runtime"
@@ -79,6 +82,11 @@ def installation(self):
         Raises:
             NotFound: If the installation is not found.
         """
+        if self._install_folder:
+            return RunConfigLoader.get_custom_installation(
+                self.workspace_client, self.product_info.product_name(), self._install_folder
+            )
+
         try:
             return self.product_info.current_installation(self.workspace_client)
         except NotFound:
@@ -204,7 +212,7 @@ def _is_testing(self):
     def _prompt_for_new_installation(self) -> WorkspaceConfig:
         configurator = WarehouseInstaller(self.workspace_client, self.prompts)
         prompter = ConfigProvider(self.prompts, configurator)
-        return prompter.prompt_new_installation()
+        return prompter.prompt_new_installation(self._install_folder)
 
     def _confirm_force_install(self) -> bool:
         if not self._force_install:
@@ -263,7 +271,7 @@ def __init__(
         self._ws = ws
         self._prompts = prompts
         self._product_info = product_info
-        self._wheels = product_info.wheels(ws)
+        self._wheels = WheelsV2(self._installation, product_info)
 
     @classmethod
     def current(cls, ws: WorkspaceClient):
@@ -282,7 +290,7 @@ def current(cls, ws: WorkspaceClient):
         config = installation.load(WorkspaceConfig)
         run_config_name = config.get_run_config().name
         prompts = Prompts()
-        wheels = product_info.wheels(ws)
+        wheels = WheelsV2(installation, product_info)
         tasks = WorkflowsRunner.all(config).tasks()
         workflow_installer = WorkflowDeployment(
             config, run_config_name, installation, install_state, ws, wheels, product_info, tasks
@@ -366,5 +374,16 @@ def _create_dashboard(self) -> None:
     if is_in_debug():
         logging.getLogger("databricks").setLevel(logging.DEBUG)
 
-    workspace_installer = WorkspaceInstaller(WorkspaceClient(product="dqx", product_version=__version__))
+    installer_prompts = Prompts()
+    custom_folder = installer_prompts.question(
+        "Enter a workspace path for DQX installation (leave empty for default behavior)",
+        default="",
+        valid_regex=r"^(/.*)?$",
+    ).strip()
+
+    custom_install_folder = custom_folder if custom_folder else None
+
+    workspace_installer = WorkspaceInstaller(
+        WorkspaceClient(product="dqx", product_version=__version__), install_folder=custom_install_folder
+    )
     workspace_installer.run()

tests/conftest.py

@@ -6,7 +6,7 @@
 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.blueprint.wheels import ProductInfo, WheelsV2
 from databricks.labs.dqx.__about__ import __version__
 from databricks.labs.dqx.config import WorkspaceConfig, RunConfig
 from databricks.labs.dqx.contexts.workflow_context import WorkflowContext
@@ -78,14 +78,16 @@ def __init__(
         ws: WorkspaceClient,
         checks_location,
         serverless_clusters: bool = True,
+        install_folder: str | None = None,
     ):
         super().__init__(env_or_skip_fixture, ws)
         self.checks_location = checks_location
         self.serverless_clusters = serverless_clusters
+        self.install_folder = install_folder
 
     @cached_property
     def installation(self):
-        return Installation(self.workspace_client, self.product_info.product_name())
+        return Installation(self.workspace_client, self.product_info.product_name(), install_folder=self.install_folder)
 
     @cached_property
     def environ(self) -> dict[str, str]:
@@ -96,6 +98,7 @@ def workspace_installer(self):
         return WorkspaceInstaller(
             self.workspace_client,
             self.environ,
+            self.install_folder,
         ).replace(prompts=self.prompts, installation=self.installation, product_info=self.product_info)
 
     @cached_property
@@ -130,7 +133,7 @@ def workflows_deployment(self) -> WorkflowDeployment:
             self.installation,
             self.install_state,
             self.workspace_client,
-            self.product_info.wheels(self.workspace_client),
+            WheelsV2(self.installation, self.product_info),
             self.product_info,
             self.tasks,
         )
@@ -188,6 +191,18 @@ def serverless_installation_ctx(
     ctx.installation_service.uninstall()
 
 
+@pytest.fixture
+def installation_ctx_custom_install_folder(
+    ws: WorkspaceClient, make_directory, env_or_skip: Callable[[str], str], checks_location="checks.yml"
+) -> Generator[MockInstallationContext, None, None]:
+    custom_folder = str(make_directory().absolute())
+    ctx = MockInstallationContext(
+        env_or_skip, ws, checks_location, serverless_clusters=False, install_folder=custom_folder
+    )
+    yield ctx.replace(workspace_client=ws)
+    ctx.installation_service.uninstall()
+
+
 @pytest.fixture
 def checks_yaml_content():
     return """- criticality: error

tests/integration/test_config.py

@@ -1,6 +1,9 @@
 from unittest.mock import patch
 
 from databricks.labs.blueprint.installation import Installation
+from databricks.labs.blueprint.wheels import ProductInfo
+
+from databricks.labs.dqx.config import WorkspaceConfig
 from databricks.labs.dqx.config_loader import RunConfigLoader
 
 
@@ -29,3 +32,14 @@ def test_load_run_config_from_global_installation(ws, installation_ctx):
         )
 
         assert run_config == expected_run_config
+
+
+def test_get_custom_installation(ws, make_directory):
+    product_info = ProductInfo.for_testing(WorkspaceConfig)
+    custom_folder = str(make_directory().absolute())
+
+    custom_installation = RunConfigLoader.get_custom_installation(ws, product_info.product_name(), custom_folder)
+    custom_installation.install_folder()
+
+    assert custom_installation.install_folder() == custom_folder
+    assert ws.workspace.get_status(custom_folder)