Improve error messages in case of connection errors (#2210)

## Changes
Improve error messages in case of connection errors

### Linked issues
Partially resolves #1323

### Functionality

- [x] modified existing command: `databricks labs ucx (un)install`


### Tests

- [x] manually tested
- [x] added integration tests

Author: JCZuurmond

README.md

@@ -1350,7 +1350,7 @@ from a VPC, or from a specific IP range.
 
 **Solution:** Please check that your laptop has network connectivity to
 the Databricks account and workspace. If not, you might need to be
-connected to a VPN to access your workspace.
+connected to a VPN or configure an HTTP proxy to access your workspace.
 
 **From local machine to GitHub:** UCX needs internet access to connect
 to [<u>github.com</u>](https://github.com) (to download the tool) from

src/databricks/labs/ucx/install.py

@@ -11,6 +11,8 @@
 from pathlib import Path
 from typing import Any
 
+from requests.exceptions import ConnectionError as RequestsConnectionError
+
 import databricks.sdk.errors
 from databricks.labs.blueprint.entrypoint import get_logger, is_in_debug
 from databricks.labs.blueprint.installation import Installation, SerdeError
@@ -174,36 +176,43 @@ def run(
         config: WorkspaceConfig | None = None,
     ) -> WorkspaceConfig:
         logger.info(f"Installing UCX v{self.product_info.version()}")
-        if config is None:
-            config = self.configure(default_config)
-        if self._is_testing():
-            return config
-        workflows_deployment = WorkflowsDeployment(
-            config,
-            self.installation,
-            self.install_state,
-            self.workspace_client,
-            self.wheels,
-            self.product_info,
-            verify_timeout,
-            self._tasks,
-        )
-        workspace_installation = WorkspaceInstallation(
-            config,
-            self.installation,
-            self.install_state,
-            self.sql_backend,
-            self.workspace_client,
-            workflows_deployment,
-            self.prompts,
-            self.product_info,
-        )
         try:
+            if config is None:
+                config = self.configure(default_config)
+            if self._is_testing():
+                return config
+            workflows_deployment = WorkflowsDeployment(
+                config,
+                self.installation,
+                self.install_state,
+                self.workspace_client,
+                self.wheels,
+                self.product_info,
+                verify_timeout,
+                self._tasks,
+            )
+            workspace_installation = WorkspaceInstallation(
+                config,
+                self.installation,
+                self.install_state,
+                self.sql_backend,
+                self.workspace_client,
+                workflows_deployment,
+                self.prompts,
+                self.product_info,
+            )
             workspace_installation.run()
         except ManyError as err:
             if len(err.errs) == 1:
                 raise err.errs[0] from None
             raise err
+        except TimeoutError as err:
+            if isinstance(err.__cause__, RequestsConnectionError):
+                logger.warning(
+                    f"Cannot connect with {self.workspace_client.config.host} see "
+                    f"https://github.com/databrickslabs/ucx#network-connectivity-issues for help: {err}"
+                )
+            raise err
         return config
 
     def _is_testing(self):
@@ -282,6 +291,11 @@ def _confirm_force_install(self) -> bool:
         return False
 
     def configure(self, default_config: WorkspaceConfig | None = None) -> WorkspaceConfig:
+        """Configure the workspaces
+
+        Notes:
+        1. Connection errors are not handled within this configure method.
+        """
         try:
             config = self.installation.load(WorkspaceConfig)
             self._compare_remote_local_versions()

tests/integration/install/test_installation.py

@@ -3,9 +3,9 @@
 import logging
 from datetime import timedelta
 
-import pytest  # pylint: disable=wrong-import-order
-from databricks.labs.ucx.__about__ import __version__
+import pytest
 
+import databricks
 from databricks.labs.blueprint.installation import Installation
 from databricks.labs.blueprint.parallel import ManyError
 from databricks.labs.blueprint.tui import MockPrompts
@@ -18,16 +18,16 @@
     NotFound,
     ResourceConflict,
 )
-
 from databricks.sdk.retries import retried
 from databricks.sdk.service import compute
 from databricks.sdk.service.iam import PermissionLevel
 
-import databricks
+from databricks.labs.ucx.__about__ import __version__
 from databricks.labs.ucx.config import WorkspaceConfig
 from databricks.labs.ucx.install import WorkspaceInstaller
 from databricks.labs.ucx.workspace_access.groups import MigratedGroup
 
+
 logger = logging.getLogger(__name__)
 
 

tests/unit/test_install.py

@@ -2,10 +2,12 @@
 import json
 import logging
 from datetime import datetime, timedelta
+from typing import Any
 from unittest.mock import MagicMock, create_autospec, patch
 
 import pytest
 import yaml
+from requests.exceptions import ConnectionError as RequestsConnectionError
 from databricks.labs.blueprint.installation import Installation, MockInstallation
 from databricks.labs.blueprint.installer import InstallState
 from databricks.labs.blueprint.parallel import ManyError
@@ -14,6 +16,8 @@
 from databricks.labs.lsql.backends import MockBackend
 from databricks.labs.lsql.dashboards import DashboardMetadata
 from databricks.sdk import AccountClient, WorkspaceClient
+from databricks.sdk.core import Config
+from databricks.sdk.credentials_provider import credentials_strategy
 from databricks.sdk.errors import (  # pylint: disable=redefined-builtin
     AlreadyExists,
     InvalidParameterValue,
@@ -1929,3 +1933,33 @@ def test_upload_dependencies(ws, mock_installation):
     workspace_installation.run()
     wheels.upload_wheel_dependencies.assert_called_once()
     wheels.upload_to_wsfs.assert_called_once()
+
+
+@pytest.fixture
+def no_connection_ws() -> WorkspaceClient:
+    """Configure a workspace like it does not have an internet connection."""
+
+    @credentials_strategy("raise_connection_error", [])
+    def raise_connection_error(_: Any):
+        """Mock no internet access by raising a ConnectionError"""
+
+        def inner():
+            raise RequestsConnectionError("no internet")
+
+        return inner
+
+    config = Config(
+        host="https://adb-123456789.12.azuredatabricks.net/",
+        credentials_strategy=raise_connection_error,
+        retry_timeout_seconds=1,
+    )
+    return WorkspaceClient(config=config)
+
+
+@pytest.mark.parametrize("default_config", [None, WorkspaceConfig("ucx")])
+def test_workspace_installer_warns_about_connection_error(caplog, no_connection_ws, default_config):
+    """This test runs both with and without internet connection"""
+    workspace_installer = WorkspaceInstaller(no_connection_ws)
+    with pytest.raises(TimeoutError), caplog.at_level(logging.WARNING, logger="databricks.labs.ucx.source_code.jobs"):
+        workspace_installer.run(default_config=default_config)
+    assert "Cannot connect with" in caplog.text