diff --git a/CHANGELOG.md b/CHANGELOG.md index 4a0380734..58dcb94f3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ ## dbt-databricks 1.12.5 (TBD) +### Features + +- Support `auth_type: env-oidc` and `auth_type: file-oidc` for authenticating via workload identity federation, with a new `oidc_token_filepath` profile config for the latter ([#1666](https://github.com/databricks/dbt-databricks/pull/1666)) + ### Fixes - Replace an existing table or view with a metric view using backup-and-create instead of `CREATE OR REPLACE VIEW ... WITH METRICS` ([#1640](https://github.com/databricks/dbt-databricks/pull/1640) resolves [#1639](https://github.com/databricks/dbt-databricks/issues/1639)) diff --git a/dbt/adapters/databricks/credentials.py b/dbt/adapters/databricks/credentials.py index b8add643f..dc17e94c6 100644 --- a/dbt/adapters/databricks/credentials.py +++ b/dbt/adapters/databricks/credentials.py @@ -52,6 +52,9 @@ class DatabricksCredentials(Credentials): connection_parameters: Optional[dict[str, Any]] = None auth_type: Optional[str] = None + # Path of the file holding the OIDC ID token, for `auth_type: file-oidc`. + oidc_token_filepath: Optional[str] = None + # Named compute resources specified in the profile. Used for # creating a connection when a model specifies a compute resource. compute: Optional[dict[str, Any]] = None @@ -140,9 +143,18 @@ def validate_creds(self) -> None: for key in ["host", "http_path"]: if not getattr(self, key): raise DbtConfigError(f"The config '{key}' is required to connect to Databricks") - if not self.token and self.auth_type != "oauth": + if not self.token and self.auth_type not in ("oauth", "env-oidc", "file-oidc"): + raise DbtConfigError( + "The config `auth_type` must be one of `oauth`, `env-oidc`, or `file-oidc` " + "when not using an access token" + ) + + # Without an explicit client_id the SDK is handed the `dbt-databricks` + # public client, which carries no federation policy, and fails with a 401. + if not self.token and self.auth_type in ("env-oidc", "file-oidc") and not self.client_id: raise DbtConfigError( - "The config `auth_type: oauth` is required when not using access token" + "The config 'client_id' is required to connect to Databricks " + f"with 'auth_type: {self.auth_type}'" ) if not self.client_id and self.client_secret: @@ -281,6 +293,7 @@ class DatabricksCredentialManager(DataClassDictMixin): oauth_scopes: list[str] = field(default_factory=lambda: SCOPES) token: Optional[str] = None auth_type: Optional[str] = None + oidc_token_filepath: Optional[str] = None workspace_id: Optional[str] = None @classmethod @@ -295,6 +308,7 @@ def create_from(cls, credentials: DatabricksCredentials) -> "DatabricksCredentia oauth_redirect_url=credentials.oauth_redirect_url or REDIRECT_URL, oauth_scopes=credentials.oauth_scopes or SCOPES, auth_type=credentials.auth_type, + oidc_token_filepath=credentials.oidc_token_filepath, workspace_id=extract_workspace_id(credentials.http_path), ) @@ -323,6 +337,16 @@ def authenticate_with_oauth_m2m(self) -> Config: ) ) + def authenticate_with_oidc(self) -> Config: + kwargs = self._config_kwargs( + host=self.host, + client_id=self.client_id, + auth_type=self.auth_type, + ) + if self.oidc_token_filepath: + kwargs["oidc_token_filepath"] = self.oidc_token_filepath + return Config(**kwargs) + def authenticate_with_external_browser(self) -> Config: return Config( **self._config_kwargs( @@ -373,6 +397,8 @@ def _ensure_config(self) -> Config: if self.token: self._config = self.authenticate_with_pat() + elif self.auth_type in ("env-oidc", "file-oidc"): + self._config = self.authenticate_with_oidc() elif self.azure_client_id and self.azure_client_secret: self._config = self.authenticate_with_azure_client_secret() elif not self.client_secret: diff --git a/dbt/adapters/databricks/handle.py b/dbt/adapters/databricks/handle.py index 3da799c9c..e94005f0b 100644 --- a/dbt/adapters/databricks/handle.py +++ b/dbt/adapters/databricks/handle.py @@ -439,9 +439,10 @@ def _add_kernel_auth_arguments( PAT, Databricks OAuth M2M, and OAuth U2M — we forward the raw credentials it understands so the connector owns the token lifecycle and refresh. - The kernel has no Azure-AD flow, so an Azure service principal cannot connect - through it; any auth other than the three above is rejected here with a clear - error directing the user to the default backend. + The kernel has no Azure-AD or OIDC federation flow, so neither an Azure + service principal nor a federated identity can connect through it; any auth + other than the three above is rejected here with a clear error directing the + user to the default backend. """ # PAT: forward the token directly. if creds_manager.token: @@ -464,8 +465,13 @@ def _add_kernel_auth_arguments( # Databricks OAuth M2M: forward OAuth creds so the kernel owns refresh. The # resolved auth_type distinguishes genuine Databricks OAuth from an Azure SP - # supplied through the client_id/client_secret fields. - if not creds_manager.azure_client_secret and creds_manager.config.auth_type == "oauth-m2m": + # in the client_id/client_secret fields. The OIDC term must stay first: + # resolving `.config` for a federated auth_type triggers a token exchange. + if ( + creds.auth_type not in ("env-oidc", "file-oidc") + and not creds_manager.azure_client_secret + and creds_manager.config.auth_type == "oauth-m2m" + ): args["oauth_client_id"] = creds_manager.client_id args["oauth_client_secret"] = creds_manager.client_secret args["oauth_scopes"] = creds_manager.oauth_scopes @@ -474,6 +480,6 @@ def _add_kernel_auth_arguments( raise DbtConfigError( "use_kernel=True supports only personal access tokens and Databricks " "OAuth (M2M/U2M); the configured authentication (e.g. an Azure service " - "principal) is not supported by the kernel backend. Remove use_kernel to " - "use the default backend." + "principal, or OIDC workload identity federation) is not supported by the " + "kernel backend. Remove use_kernel to use the default backend." ) diff --git a/tests/unit/test_auth.py b/tests/unit/test_auth.py index 4e03ddfc6..75d609783 100644 --- a/tests/unit/test_auth.py +++ b/tests/unit/test_auth.py @@ -3,6 +3,7 @@ from unittest import mock import pytest +from dbt_common.exceptions import DbtConfigError from dbt.adapters.databricks.credentials import ( DatabricksCredentialManager, @@ -52,6 +53,39 @@ def test_azure_client_secret_credentials_init_does_not_call_config(self): ) mock_config.assert_not_called() + def test_env_oidc_credentials_init_does_not_call_config(self): + with mock.patch("dbt.adapters.databricks.credentials.Config") as mock_config: + DatabricksCredentials(client_id="cid", auth_type="env-oidc", **_COMMON_KWARGS) + mock_config.assert_not_called() + + +class TestValidateCreds: + """`validate_creds` runs at connect time and is the only place that rejects + an unusable combination of profile fields before the SDK is involved.""" + + def test_oidc_auth_types_need_no_token(self): + for auth_type in ("env-oidc", "file-oidc"): + creds = DatabricksCredentials(client_id="cid", auth_type=auth_type, **_COMMON_KWARGS) + creds.validate_creds() + + def test_unknown_auth_type_without_token_raises(self): + creds = DatabricksCredentials( + client_id="cid", auth_type="not-a-real-auth-type", **_COMMON_KWARGS + ) + with pytest.raises(DbtConfigError, match="must be one of"): + creds.validate_creds() + + @pytest.mark.parametrize("auth_type", ["env-oidc", "file-oidc"]) + def test_oidc_without_client_id_raises(self, auth_type): + creds = DatabricksCredentials(auth_type=auth_type, **_COMMON_KWARGS) + with pytest.raises(DbtConfigError, match="'client_id' is required"): + creds.validate_creds() + + @pytest.mark.parametrize("auth_type", ["env-oidc", "file-oidc"]) + def test_token_removes_the_client_id_requirement(self, auth_type): + creds = DatabricksCredentials(token="foo", auth_type=auth_type, **_COMMON_KWARGS) + creds.validate_creds() + class TestEnsureConfigTriggersTheRightAuth: """Connect-time counterpart to TestParseTimeIsOffline: when something @@ -124,6 +158,50 @@ def test_non_dose_secret_tries_legacy_azure_first(self): auth_type="azure-client-secret", ) + def test_env_oidc_uses_oidc_auth(self): + creds = DatabricksCredentials(client_id="cid", auth_type="env-oidc", **_COMMON_KWARGS) + with mock.patch("dbt.adapters.databricks.credentials.Config") as mock_config: + creds.authenticate().config + mock_config.assert_called_once_with( + host=_COMMON_KWARGS["host"], + client_id="cid", + auth_type="env-oidc", + ) + + def test_file_oidc_uses_oidc_auth(self): + creds = DatabricksCredentials(client_id="cid", auth_type="file-oidc", **_COMMON_KWARGS) + with mock.patch("dbt.adapters.databricks.credentials.Config") as mock_config: + creds.authenticate().config + mock_config.assert_called_once_with( + host=_COMMON_KWARGS["host"], + client_id="cid", + auth_type="file-oidc", + ) + + def test_file_oidc_forwards_token_filepath(self): + creds = DatabricksCredentials( + client_id="cid", + auth_type="file-oidc", + oidc_token_filepath="/var/run/secrets/token", + **_COMMON_KWARGS, + ) + with mock.patch("dbt.adapters.databricks.credentials.Config") as mock_config: + creds.authenticate().config + mock_config.assert_called_once_with( + host=_COMMON_KWARGS["host"], + client_id="cid", + auth_type="file-oidc", + oidc_token_filepath="/var/run/secrets/token", + ) + + def test_token_takes_precedence_over_oidc_auth_type(self): + creds = DatabricksCredentials( + token="foo", client_id="cid", auth_type="env-oidc", **_COMMON_KWARGS + ) + with mock.patch("dbt.adapters.databricks.credentials.Config") as mock_config: + creds.authenticate().config + mock_config.assert_called_once_with(host=_COMMON_KWARGS["host"], token="foo") + def test_falls_back_to_second_method_when_first_raises(self): creds = DatabricksCredentials( client_id="cid", diff --git a/tests/unit/test_handle.py b/tests/unit/test_handle.py index 8f58c25a6..348b6491a 100644 --- a/tests/unit/test_handle.py +++ b/tests/unit/test_handle.py @@ -320,6 +320,18 @@ def test_prepare_connection_arguments__kernel_legacy_azure_sp_raises(self): with pytest.raises(DbtConfigError, match="use_kernel"): SqlUtils.prepare_connection_arguments(creds, manager, _KERNEL_HTTP_PATH, {}) + @pytest.mark.parametrize("auth_type", ["env-oidc", "file-oidc"]) + def test_prepare_connection_arguments__kernel_oidc_raises(self, auth_type): + """The kernel has no OIDC flow. Rejected on the raw auth_type before the + oauth-m2m check, which would otherwise resolve `.config` and trigger the + SDK's OIDC token exchange here.""" + with pytest.raises(DbtConfigError, match="use_kernel"): + _prepare_connection_args_without_config_auth( + client_id="cid", + auth_type=auth_type, + connection_parameters={"use_kernel": True}, + ) + def test_prepare_connection_arguments__kernel_u2m_explicit_client_id(self): """use_kernel with OAuth U2M and an explicit client_id forwards that client_id and translates dbt's auth_type='oauth' to the kernel's