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

src/databricks/sql/auth/oauth.py

@@ -62,18 +62,41 @@ def refresh(self) -> Token:
 
 
 class OAuthManager:
+    # Default maximum time (in seconds) to wait for the browser OAuth redirect
+    # callback before giving up. Without a bound, 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. This is generous for interactive logins (slow MFA, IdP re-auth, SSO
+    # redirects) and is a fixed ceiling for connector end users: there is no
+    # public ``connect()``/``Connection`` parameter to change it. The
+    # ``redirect_callback_timeout_seconds`` argument below is an internal-only
+    # override for callers constructing ``OAuthManager`` directly (private API);
+    # it is not plumbed through ``DatabricksOAuthProvider`` or the public
+    # connection kwargs.
+    REDIRECT_CALLBACK_TIMEOUT_SECONDS = 60 * 5
+
     def __init__(
         self,
         port_range: List[int],
         client_id: str,
         idp_endpoint: OAuthEndpointCollection,
         http_client,
+        redirect_callback_timeout_seconds: Optional[int] = None,
     ):
         self.port_range = port_range
         self.client_id = client_id
         self.redirect_port = None
         self.idp_endpoint = idp_endpoint
         self.http_client = http_client
+        # Fall back to the class default when not overridden. Kept as a
+        # lowercase instance attribute (distinct from the ALL_CAPS class
+        # constant that provides the default) so the rest of the flow reads a
+        # single, clearly per-instance source of truth.
+        self._redirect_callback_timeout_seconds = (
+            redirect_callback_timeout_seconds
+            if redirect_callback_timeout_seconds is not None
+            else self.REDIRECT_CALLBACK_TIMEOUT_SECONDS
+        )
 
     @staticmethod
     def __token_urlsafe(nbytes=32):
@@ -128,11 +151,40 @@ def __get_challenge():
 
     def __get_authorization_code(self, client, auth_url, scope, state, challenge):
         handler = OAuthHttpSingleRequestHandler("Databricks Sql Connector")
+        # Bound the read of the accepted connection too (not just the accept
+        # wait below). StreamRequestHandler.setup() applies this to the accepted
+        # socket via settimeout(), so a client that connects but never completes
+        # the request line can no longer block indefinitely.
+        handler.timeout = self._redirect_callback_timeout_seconds
 
         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.
+                    # NOTE: httpd.timeout bounds only the wait to ACCEPT an
+                    # incoming connection (it is the select() timeout used by
+                    # handle_request). The subsequent read of the HTTP request
+                    # line is bounded separately by the handler.timeout set
+                    # above (applied to the accepted socket by
+                    # StreamRequestHandler.setup()), so a client that connects
+                    # but never completes the request can no longer block. The
+                    # headless "no connection ever arrives" case (issue #458) is
+                    # covered by this accept-wait timeout.
+                    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 +216,18 @@ 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}. The login flow was not completed in time — "
+                    "either the interactive login was not finished within the "
+                    "timeout, or this is a headless environment (e.g. a notebook "
+                    "or job with no browser) where no browser can complete the "
+                    "flow. 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,76 @@ 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.
+        # Keep the test fast: shorten the callback wait via the constructor
+        # keyword (the actual new surface this PR introduces) rather than
+        # mutating the private attribute after the fact. This also exercises the
+        # Optional[int] fall-back logic. The production default is minutes; the
+        # bug is that WITHOUT any bound the wait is infinite.
+        oauth_manager = OAuthManager(
+            port_range=[0],
+            client_id="mock-id",
+            idp_endpoint=InHouseOAuthEndpointCollection(),
+            http_client=MagicMock(),
+            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