Work around DATABRICKS_HOST normalisation issue during install and uninstall (#2052)

## Changes

This PR introduces a workaround for a `DATABRICKS_HOST` normalisation
issue in the underlying Python SDK (databricks/databricks-sdk-py#996).
We fixed this for the CLI entry points in #1706 but missed the `install`
and `uninstall` hooks, which this PR addresses.

### Relevant implementation details

The actual workaround is implemented in Blueprint, but a bit hard to get
to… the easiest way to get to it without modifying Blueprint is to
introduce the `App` subclass.

### Linked issues

Resolves #2037

### Functionality

- modified existing commands:

  - `databricks labs install lakebridge`
  - `databricks labs uninstall lakebridge`

### Tests

Manually tested, see logs below:

 - Normal situation, with `host=` set to just a host:

- [`databricks labs install lakebridge
--debug`](https://github.com/user-attachments/files/22517659/install-normal.log)
- [`databricks labs uninstall lakebridge
--debug`](https://github.com/user-attachments/files/22517616/uninstall-normal.log)

 - Situation with `host=` set to a login URL:

- [`databricks labs install lakebridge
--debug`](https://github.com/user-attachments/files/22517675/install-with-workaround.log)
- [`databricks labs uninstall lakebridge
--debug`](https://github.com/user-attachments/files/22517697/uninstall-with-workaround.log)

Author: asnare

src/databricks/labs/lakebridge/base_install.py

@@ -2,10 +2,9 @@
 
 from databricks.labs.blueprint.logger import install_logger
 from databricks.labs.blueprint.entrypoint import get_logger
-from databricks.sdk import WorkspaceClient
 from databricks.sdk.core import with_user_agent_extra
 
-from databricks.labs.lakebridge import __version__
+from databricks.labs.lakebridge.cli import lakebridge
 from databricks.labs.lakebridge.install import installer as _installer
 from databricks.labs.lakebridge.transpiler.repository import TranspilerRepository
 
@@ -18,7 +17,7 @@ def main() -> None:
     logger.setLevel("INFO")
 
     installer = _installer(
-        WorkspaceClient(product="lakebridge", product_version=__version__),
+        ws=lakebridge.create_workspace_client(),
         transpiler_repository=TranspilerRepository.user_home(),
         is_interactive=sys.stdin.isatty(),
     )

src/databricks/labs/lakebridge/cli.py

@@ -42,7 +42,19 @@
 
 from databricks.labs.lakebridge.transpiler.transpile_status import ErrorSeverity
 
-lakebridge = App(__file__)
+
+# Subclass to allow controlled access to protected methods.
+class Lakebridge(App):
+    def create_workspace_client(self) -> WorkspaceClient:
+        """Create a workspace client, with the appropriate product and version information.
+
+        This is intended only for use by the install/uninstall hooks.
+        """
+        self._patch_databricks_host()
+        return self._workspace_client()
+
+
+lakebridge = Lakebridge(__file__)
 logger = get_logger(__file__)
 
 

src/databricks/labs/lakebridge/uninstall.py

@@ -1,9 +1,8 @@
 import logging
 
 from databricks.labs.blueprint.entrypoint import is_in_debug
-from databricks.sdk import WorkspaceClient
 
-from databricks.labs.lakebridge.__about__ import __version__
+from databricks.labs.lakebridge.cli import lakebridge
 from databricks.labs.lakebridge.contexts.application import ApplicationContext
 
 logger = logging.getLogger("databricks.labs.lakebridge.install")
@@ -18,11 +17,4 @@ def run(context: ApplicationContext):
     if is_in_debug():
         logging.getLogger("databricks").setLevel(logging.DEBUG)
 
-    run(
-        ApplicationContext(
-            WorkspaceClient(
-                product="lakebridge",
-                product_version=__version__,
-            )
-        )
-    )
+    run(ApplicationContext(ws=lakebridge.create_workspace_client()))