Fix OAuth M2M auth using credentials_provider instead of unsupported client_id/client_secret kwargs for sql.connect() (addresses Codex P1)

This commit is contained in:
Sameer6305
2026-07-16 20:11:58 +05:30
parent 2d5bd18fa4
commit bdcbaa3173
2 changed files with 54 additions and 6 deletions
+27 -2
View File
@@ -50,11 +50,14 @@ from ..utils.progress_tracker import get_progress_tracker
try:
from databricks import sql as databricks_sql
from databricks.sdk import WorkspaceClient
from databricks.sdk.core import Config, oauth_service_principal
DATABRICKS_AVAILABLE = True
except (ImportError, OSError):
databricks_sql = None
WorkspaceClient = None
Config = None
oauth_service_principal = None
DATABRICKS_AVAILABLE = False
@@ -170,6 +173,13 @@ class DatabricksConnector:
used as a context manager), it is reused instead of opening a second
one, which would otherwise be left unclosed.
Authentication:
- Personal access token: pass ``access_token`` to ``sql.connect``.
- OAuth M2M (service principal): ``sql.connect`` does **not** accept
``client_id``/``client_secret`` directly. The correct mechanism is
``credentials_provider``, a callable that returns a header-factory
produced by ``databricks.sdk.core.oauth_service_principal``.
Returns:
Connection: databricks-sql-connector connection object
@@ -194,8 +204,23 @@ class DatabricksConnector:
}
if self.client_id and self.client_secret:
conn_params["client_id"] = self.client_id
conn_params["client_secret"] = self.client_secret
# databricks-sql-connector ≥2.5 requires OAuth M2M to be wired
# through a credentials_provider callable; passing client_id /
# client_secret as plain kwargs is silently ignored and causes
# the connector to fall back to an interactive browser flow.
_client_id = self.client_id
_client_secret = self.client_secret
_host = self.host
def _m2m_credentials_provider():
cfg = Config(
host=_host,
client_id=_client_id,
client_secret=_client_secret,
)
return oauth_service_principal(cfg)
conn_params["credentials_provider"] = _m2m_credentials_provider
else:
conn_params["access_token"] = self.token
+27 -4
View File
@@ -145,13 +145,23 @@ class TestDatabricksConnector:
assert call_kwargs["access_token"] == "test_token"
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
@patch("semantica.ingest.databricks_ingestor.oauth_service_principal")
@patch("semantica.ingest.databricks_ingestor.Config")
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
def test_connector_connect_oauth_m2m(self, mock_sql, mock_databricks_connection):
"""Test connection with OAuth M2M authentication."""
def test_connector_connect_oauth_m2m(
self, mock_sql, mock_config, mock_oauth_sp, mock_databricks_connection
):
"""Test connection with OAuth M2M authentication uses credentials_provider.
databricks-sql-connector does NOT accept client_id/client_secret as
direct kwargs to sql.connect(); the correct mechanism is a
credentials_provider callable wrapping oauth_service_principal().
"""
from semantica.ingest.databricks_ingestor import DatabricksConnector
mock_conn, _ = mock_databricks_connection
mock_sql.connect = Mock(return_value=mock_conn)
mock_oauth_sp.return_value = {"Authorization": "Bearer fake-token"}
connector = DatabricksConnector(
host="https://adb-xxx.azuredatabricks.net",
@@ -163,10 +173,23 @@ class TestDatabricksConnector:
connector.connect()
call_kwargs = mock_sql.connect.call_args[1]
assert call_kwargs["client_id"] == "test_client_id"
assert call_kwargs["client_secret"] == "test_client_secret"
# Must use credentials_provider, not bare client_id/client_secret
assert "credentials_provider" in call_kwargs
assert callable(call_kwargs["credentials_provider"])
assert "client_id" not in call_kwargs
assert "client_secret" not in call_kwargs
assert "access_token" not in call_kwargs
# Invoke the provider to verify it wires Config + oauth_service_principal
call_kwargs["credentials_provider"]()
mock_config.assert_called_once_with(
host="https://adb-xxx.azuredatabricks.net",
client_id="test_client_id",
client_secret="test_client_secret",
)
mock_oauth_sp.assert_called_once_with(mock_config.return_value)
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
def test_connector_connect_missing_http_path(self, mock_sql):