ai: apply changes for #876 (1 review thread)

Addresses:
  - #3635880990 at src/databricks/sql/auth/oauth.py:145

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

Author: peco-engineer-bot[bot]

src/databricks/sql/auth/oauth.py

@@ -136,13 +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,
@@ -174,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

@@ -262,7 +262,13 @@ def run():
         )
         self.assertEqual(result.get("outcome"), "raised")
         self.assertIsInstance(result.get("error"), RuntimeError)
-        self.assertIn("No path parameters were returned", str(result.get("error")))
+        # 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: