ai: apply changes for #876 (2 review threads)

Addresses:
  - #3636172662 at src/databricks/sql/auth/oauth.py:173
  - #3636172666 at src/databricks/sql/auth/oauth.py:187

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

Author: peco-engineer-bot[bot]

src/databricks/sql/auth/authenticators.py

@@ -1,6 +1,6 @@
 import abc
 import logging
-from typing import Callable, Dict, List
+from typing import Callable, Dict, List, Optional
 from databricks.sql.common.http import HttpHeader
 from databricks.sql.auth.oauth import (
     OAuthManager,
@@ -63,6 +63,7 @@ def __init__(
         scopes: List[str],
         http_client,
         auth_type: str = "databricks-oauth",
+        redirect_callback_timeout_seconds: Optional[int] = None,
     ):
         try:
             idp_endpoint = get_oauth_endpoints(hostname, auth_type == "azure-oauth")
@@ -74,11 +75,15 @@ def __init__(
             # Convert to the corresponding scopes in the corresponding IdP
             cloud_scopes = idp_endpoint.get_scopes_mapping(scopes)
 
+            # Pass through the (optional) callback-timeout override so this
+            # private provider can surface it later if a public connection
+            # knob is added; None falls back to OAuthManager's default.
             self.oauth_manager = OAuthManager(
                 port_range=redirect_port_range,
                 client_id=client_id,
                 idp_endpoint=idp_endpoint,
                 http_client=http_client,
+                redirect_callback_timeout_seconds=redirect_callback_timeout_seconds,
             )
             self._hostname = hostname
             self._scopes_as_str = DatabricksOAuthProvider.SCOPE_DELIM.join(cloud_scopes)

src/databricks/sql/auth/oauth.py

@@ -1,4 +1,5 @@
 import base64
+import errno
 import hashlib
 import json
 import logging
@@ -204,7 +205,11 @@ def _on_timeout():
                 self.redirect_port = port
                 break
             except OSError as e:
-                if e.errno == 48:
+                # errno.EADDRINUSE resolves to the platform's value (48 on
+                # macOS, 98 on Linux), so a port-in-use bind failure is
+                # recognized cross-platform. Otherwise last_error stays None on
+                # Linux and `raise last_error` below becomes `raise None`.
+                if e.errno == errno.EADDRINUSE:
                     logger.info(f"Port {port} is in use")
                     last_error = e
             except Exception as e:

tests/unit/test_auth.py

@@ -350,10 +350,9 @@ def run():
             # closing the probe socket and the HTTPServer bind inside
             # __get_authorization_code, another process can claim the port. That
             # would take the can't-bind branch instead of the stalled-request-
-            # line path this test exercises (and the source's errno==48 check is
-            # macOS-only, so on Linux the bind failure surfaces as a TypeError
-            # rather than the RuntimeError asserted below). Re-probe a fresh port
-            # and retry on that rare race so the test stays deterministic in CI.
+            # line path this test exercises, surfacing a bind error rather than
+            # the RuntimeError asserted below. Re-probe a fresh port and retry
+            # on that rare race so the test stays deterministic in CI.
             worker, result = None, {}
             for _ in range(5):
                 probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)