Skip to content

Commit 5262855

Browse files
Raise DB-API errors for transport failures and empty access tokens
- ThriftDatabricksClient.make_request re-raised urllib3 HTTP errors of every request except GetOperationStatus unchanged, so an unreachable workspace (refused connection, proxy/tunnel failure) surfaced from connect() and execute() as a raw urllib3 MaxRetryError instead of a DB-API error. Report it through the normal non-retryable path as a RequestError (OperationalError). - connect(access_token="") dropped the empty token and fell back to the interactive browser OAuth flow, blocking on a local callback server. An explicitly empty token now raises 'No valid authentication settings!' at once. Signed-off-by: Amin Ghadersohi <amin.ghadersohi@gmail.com>
1 parent 01564c7 commit 5262855

5 files changed

Lines changed: 48 additions & 3 deletions

File tree

‎src/databricks/sql/auth/auth.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,10 @@ def get_auth_provider(cfg: ClientContext, http_client):
4343
cfg.auth_type,
4444
)
4545
elif cfg.access_token is not None:
46+
if not cfg.access_token:
47+
# An explicitly empty token is a missing credential; never fall
48+
# back to an interactive browser login for it.
49+
raise RuntimeError("No valid authentication settings! access_token is empty")
4650
base_provider = AccessTokenAuthProvider(cfg.access_token)
4751
elif cfg.use_cert_as_auth and cfg.tls_client_cert_file:
4852
# no op authenticator. authentication is performed using ssl certificate outside of headers

‎src/databricks/sql/backend/thrift_backend.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -459,7 +459,11 @@ def attempt_request(attempt):
459459
f"GetOperationStatus failed with HTTP error and will be retried: {str(err)}"
460460
)
461461
else:
462-
raise err
462+
# Not retried here (urllib3 already applied the retry policy),
463+
# but surfaced as a DB-API RequestError rather than a raw
464+
# urllib3 exception.
465+
error = err
466+
error_message = str(err)
463467
except OSError as err:
464468
error = err
465469
error_message = str(err)

‎src/databricks/sql/client.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -372,7 +372,7 @@ def read(self) -> Optional[OAuthToken]:
372372
http_path,
373373
)
374374

375-
if access_token:
375+
if access_token is not None:
376376
access_token_kv = {"access_token": access_token}
377377
kwargs = {**kwargs, **access_token_kv}
378378

‎tests/unit/test_auth.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,13 @@ def test_get_python_sql_connector_auth_provider_access_token(self):
152152
auth_provider.add_headers(headers)
153153
self.assertEqual(headers["Authorization"], "Bearer dpi123")
154154

155+
def test_get_python_sql_connector_auth_provider_empty_access_token(self):
156+
"""An explicitly empty token must not start an interactive OAuth login."""
157+
with self.assertRaisesRegex(RuntimeError, "No valid authentication settings!"):
158+
get_python_sql_connector_auth_provider(
159+
"moderakh-test.cloud.databricks.com", MagicMock(), access_token=""
160+
)
161+
155162
def test_get_python_sql_connector_auth_provider_external(self):
156163
class MyProvider(CredentialsProvider):
157164
def auth_type(self) -> str:

‎tests/unit/test_thrift_backend.py‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1927,7 +1927,7 @@ def test_make_request_will_retry_GetOperationStatus(
19271927

19281928
import thrift, errno
19291929
from databricks.sql.thrift_api.TCLIService.TCLIService import Client
1930-
from databricks.sql.exc import RequestError
1930+
from databricks.sql.exc import RequestError, RequestError
19311931
from databricks.sql.utils import NoRetryReason
19321932

19331933
this_gos_name = "GetOperationStatus"
@@ -2047,6 +2047,36 @@ def test_make_request_will_retry_GetOperationStatus_for_http_error(
20472047
f"{EXPECTED_RETRIES}/{EXPECTED_RETRIES}", cm.exception.context["attempt"]
20482048
)
20492049

2050+
@patch("thrift.transport.THttpClient.THttpClient")
2051+
def test_make_request_wraps_urllib3_http_error_as_request_error(
2052+
self, t_transport_class
2053+
):
2054+
import urllib3
2055+
2056+
t_transport_instance = t_transport_class.return_value
2057+
t_transport_instance.code = None
2058+
t_transport_instance.headers = {}
2059+
mock_method = Mock()
2060+
mock_method.__name__ = "OpenSession"
2061+
mock_method.side_effect = urllib3.exceptions.MaxRetryError(
2062+
None, "/", "Tunnel connection failed: 503"
2063+
)
2064+
2065+
thrift_backend = ThriftDatabricksClient(
2066+
"foobar",
2067+
443,
2068+
"path",
2069+
[],
2070+
auth_provider=AuthProvider(),
2071+
ssl_options=SSLOptions(),
2072+
http_client=MagicMock(),
2073+
)
2074+
2075+
with self.assertRaises(RequestError) as cm:
2076+
thrift_backend.make_request(mock_method, Mock())
2077+
2078+
self.assertIn("Tunnel connection failed", str(cm.exception.message_with_context()))
2079+
20502080
@patch("thrift.transport.THttpClient.THttpClient")
20512081
def test_make_request_wont_retry_if_error_code_not_429_or_503(
20522082
self, t_transport_class

0 commit comments

Comments
 (0)