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

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],
@@ -130,9 +136,24 @@ def __get_authorization_code(self, client, auth_url, scope, state, challenge):
         handler = OAuthHttpSingleRequestHandler("Databricks Sql Connector")
 
         last_error = None
+        callback_timed_out = False
         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
+
+                    # HTTPServer.handle_request() returns normally (via
+                    # handle_timeout()) when the wait elapses without a
+                    # connection, so record that case to distinguish it from a
+                    # received-but-empty callback below.
+                    def _on_timeout():
+                        nonlocal callback_timed_out
+                        callback_timed_out = True
+
+                    httpd.handle_timeout = _on_timeout
                     redirect_url = OAuthManager.__get_redirect_url(port)
                     auth_req_uri, _, _ = client.prepare_authorization_request(
                         authorization_url=auth_url,
@@ -164,7 +185,16 @@ def __get_authorization_code(self, client, auth_url, scope, state, challenge):
             raise last_error
 
         if not handler.request_path:
-            msg = f"No path parameters were returned to the callback at {redirect_url}"
+            if callback_timed_out:
+                msg = (
+                    f"Timed out after {self.REDIRECT_CALLBACK_TIMEOUT_SECONDS} "
+                    f"seconds waiting for the OAuth redirect callback at "
+                    f"{redirect_url}. No browser completed the login flow — this "
+                    "is expected in a headless environment (e.g. a notebook or "
+                    "job with no browser). See issue #458."
+                )
+            else:
+                msg = f"No path parameters were returned to the callback at {redirect_url}"
             logger.error(msg)
             raise RuntimeError(msg)
         # This is a kludge because the parsing library expects https callbacks

tests/unit/test_auth.py

@@ -207,6 +207,73 @@ 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
+
+        # Bind to an OS-assigned ephemeral port (0) rather than a fixed port so
+        # the test does not depend on a specific port being free. A fixed port
+        # that happens to be occupied would fail to bind and take the
+        # can't-find-free-port branch instead of the timeout path we exercise.
+        oauth_manager = OAuthManager(
+            port_range=[0],
+            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)
+        # The headless timeout path must surface a clear, timeout-specific error
+        # (not the generic received-but-empty callback message). See issue #458.
+        error_message = str(result.get("error"))
+        self.assertIn("Timed out", error_message)
+        self.assertIn(
+            str(oauth_manager.REDIRECT_CALLBACK_TIMEOUT_SECONDS), error_message
+        )
+
 
 class TestClientCredentialsTokenSource:
     @pytest.fixture