If the token is null, the connection hangs (#458)

Signed-off-by: peco-engineer-bot[bot] <3815206+peco-engineer-bot[bot]@users.noreply.github.com>

Author: peco-engineer-bot[bot]

src/databricks/sql/auth/oauth.py

@@ -62,6 +62,12 @@ def refresh(self) -> Token:
 
 
 class OAuthManager:
+    # Maximum time (in seconds) to wait for the browser OAuth redirect callback
+    # before giving up. Without this, the local callback server would block
+    # forever in a headless environment (e.g. a notebook/job with no browser),
+    # making the connection appear to hang indefinitely. See issue #458.
+    REDIRECT_CALLBACK_TIMEOUT_SECONDS = 60 * 5
+
     def __init__(
         self,
         port_range: List[int],
@@ -133,6 +139,10 @@ def __get_authorization_code(self, client, auth_url, scope, state, challenge):
         for port in self.port_range:
             try:
                 with HTTPServer(("", port), handler) as httpd:
+                    # Bound how long we wait for the browser redirect callback so
+                    # that a headless environment (no browser to complete the
+                    # flow) fails with a clear error instead of hanging forever.
+                    httpd.timeout = self.REDIRECT_CALLBACK_TIMEOUT_SECONDS
                     redirect_url = OAuthManager.__get_redirect_url(port)
                     auth_req_uri, _, _ = client.prepare_authorization_request(
                         authorization_url=auth_url,

tests/unit/test_auth.py

@@ -207,6 +207,63 @@ def test_get_python_sql_connector_default_auth(self, mock__initial_get_token):
 
         self.assertEqual(auth_provider.external_provider._client_id, PYSQL_OAUTH_CLIENT_ID)
 
+    @patch("databricks.sql.auth.oauth.webbrowser.open_new")
+    @patch.object(OAuthManager, "_OAuthManager__fetch_well_known_config")
+    def test_get_tokens_does_not_hang_when_no_callback_received(
+        self, mock_fetch_config, mock_open_new
+    ):
+        """When the U2M browser OAuth callback never arrives (e.g. a headless
+        notebook/job with a null token), the local redirect server must not
+        block forever. It should time out and surface a clear error rather than
+        hang. See issue #458."""
+        mock_fetch_config.return_value = {
+            "authorization_endpoint": "https://foo.cloud.databricks.com/oidc/oauth2/v2.0/authorize",
+            "token_endpoint": "https://foo.cloud.databricks.com/oidc/oauth2/v2.0/token",
+        }
+        # Do not actually launch a browser during the test.
+        mock_open_new.return_value = True
+
+        oauth_manager = OAuthManager(
+            port_range=[8030],
+            client_id="mock-id",
+            idp_endpoint=InHouseOAuthEndpointCollection(),
+            http_client=MagicMock(),
+        )
+        # Keep the test fast: shorten the callback wait. The production default
+        # is minutes; the bug is that WITHOUT any bound the wait is infinite.
+        oauth_manager.REDIRECT_CALLBACK_TIMEOUT_SECONDS = 2
+
+        # No callback is ever delivered to the redirect server. Run the flow in
+        # a daemon thread and join with a wall-clock bound so a regression to
+        # the infinite-block behaviour fails this test loudly instead of hanging
+        # the whole suite.
+        import threading
+
+        result = {}
+
+        def run():
+            try:
+                oauth_manager.get_tokens(
+                    hostname="foo.cloud.databricks.com", scope="offline_access sql"
+                )
+                result["outcome"] = "returned"
+            except BaseException as e:  # noqa: BLE001
+                result["outcome"] = "raised"
+                result["error"] = e
+
+        worker = threading.Thread(target=run, daemon=True)
+        worker.start()
+        worker.join(timeout=30)
+
+        self.assertFalse(
+            worker.is_alive(),
+            "OAuth callback server blocked indefinitely waiting for a callback "
+            "that never arrives (issue #458)",
+        )
+        self.assertEqual(result.get("outcome"), "raised")
+        self.assertIsInstance(result.get("error"), RuntimeError)
+        self.assertIn("No path parameters were returned", str(result.get("error")))
+
 
 class TestClientCredentialsTokenSource:
     @pytest.fixture