shahar1 commented on code in PR #75248:
URL: https://github.com/apache/airflow/pull/75248#discussion_r4503271743
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -612,6 +624,107 @@ async def
_a_get_aad_token_for_default_az_credential(self, resource: str) -> str
return token.token
+ @staticmethod
+ def _is_jwt_token_valid(token: str) -> bool:
+ """
+ Check if a JWT token hasn't expired yet.
+
+ Only the ``exp`` claim of the payload is inspected; the signature is
not verified
+ here because this check only decides whether the cached token can be
reused --
+ the token itself is verified by Databricks.
+
+ :param token: JWT token string
+ :return: true if the token is well-formed and expires later than the
refresh
+ lead time, false otherwise
+ """
+ try:
+ payload = token.split(".")[1]
+ payload += "=" * (-len(payload) % 4)
+ decoded = json.loads(base64.urlsafe_b64decode(payload))
+ return int(decoded["exp"]) > (int(time.time()) +
TOKEN_REFRESH_LEAD_TIME)
+ except (IndexError, KeyError, TypeError, ValueError):
+ return False
+
+ def _get_google_id_token(self) -> str:
+ """
+ Get a Google-signed OpenID Connect ID token.
+
+ The token is issued for the audience configured in the
``google_id_token_audience``
+ connection extra. If ``google_impersonated_service_account`` is set,
the default
+ credentials (e.g. from ``GOOGLE_APPLICATION_CREDENTIALS``) impersonate
that service
+ account through the IAM Credentials API; otherwise the token is
fetched for the
+ default credentials themselves.
+
+ :return: Google ID token, or raise an exception
+ """
+ audience =
self.databricks_conn.extra_dejson.get("google_id_token_audience")
+ if not audience:
+ raise AirflowException(
+ "`google_id_token_audience` must be present in Connection
extra "
+ "when `use_google_id_token` is enabled"
+ )
+ impersonated_service_account = self.databricks_conn.extra_dejson.get(
+ "google_impersonated_service_account"
+ )
+ resource = f"{audience}-{impersonated_service_account}"
+
+ google_id_token = self.jwt_tokens.get(resource)
+ if google_id_token and self._is_jwt_token_valid(google_id_token):
+ return google_id_token
+
+ self.log.info("Existing Google ID token is expired, or going to expire
soon. Refreshing...")
+ try:
+ import google.auth
+ from google.auth.transport.requests import AuthorizedSession,
Request
+ from google.oauth2 import id_token
+
+ for attempt in self._get_retry_object():
+ with attempt:
+ if impersonated_service_account:
+ credentials, _ =
google.auth.default(scopes=DEFAULT_GCP_SCOPES)
+ authed_session = AuthorizedSession(credentials)
+ resp = authed_session.post(
Review Comment:
**minor**: this is a hand-rolled `generateIdToken` call, using its own URL
constant and the `AuthorizedSession` POST. Together with the custom JWT decoder
in `_is_jwt_token_valid`, it re-implements something google-auth already offers
as public API:
```python
from google.auth import impersonated_credentials
source_credentials, _ = google.auth.default(scopes=DEFAULT_GCP_SCOPES)
target = impersonated_credentials.Credentials(
source_credentials=source_credentials,
target_principal=impersonated_service_account,
target_scopes=DEFAULT_GCP_SCOPES,
)
id_creds = impersonated_credentials.IDTokenCredentials(target,
target_audience=audience, include_email=True)
id_creds.refresh(Request())
# id_creds.token, id_creds.expiry
```
The Databricks SDK's own `google-id` credentials provider does exactly this
(`databricks/sdk/credentials_provider.py`). Both the impersonated and the
`fetch_id_token_credentials` paths return credentials objects with an `expiry`
attribute. If you cache that object (or the token together with its expiry),
you can drop `_is_jwt_token_valid` and the `GOOGLE_IMPERSONATED_ID_TOKEN_URL`
constant. This is not a project rule, just less code to maintain; your call.
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -612,6 +624,107 @@ async def
_a_get_aad_token_for_default_az_credential(self, resource: str) -> str
return token.token
+ @staticmethod
+ def _is_jwt_token_valid(token: str) -> bool:
+ """
+ Check if a JWT token hasn't expired yet.
+
+ Only the ``exp`` claim of the payload is inspected; the signature is
not verified
+ here because this check only decides whether the cached token can be
reused --
+ the token itself is verified by Databricks.
+
+ :param token: JWT token string
+ :return: true if the token is well-formed and expires later than the
refresh
+ lead time, false otherwise
+ """
+ try:
+ payload = token.split(".")[1]
+ payload += "=" * (-len(payload) % 4)
+ decoded = json.loads(base64.urlsafe_b64decode(payload))
+ return int(decoded["exp"]) > (int(time.time()) +
TOKEN_REFRESH_LEAD_TIME)
+ except (IndexError, KeyError, TypeError, ValueError):
+ return False
+
+ def _get_google_id_token(self) -> str:
+ """
+ Get a Google-signed OpenID Connect ID token.
+
+ The token is issued for the audience configured in the
``google_id_token_audience``
+ connection extra. If ``google_impersonated_service_account`` is set,
the default
+ credentials (e.g. from ``GOOGLE_APPLICATION_CREDENTIALS``) impersonate
that service
+ account through the IAM Credentials API; otherwise the token is
fetched for the
+ default credentials themselves.
+
+ :return: Google ID token, or raise an exception
+ """
+ audience =
self.databricks_conn.extra_dejson.get("google_id_token_audience")
+ if not audience:
+ raise AirflowException(
Review Comment:
**major**: this is one of four new `raise AirflowException` in this PR (also
lines 698, 707 and 710). `generated/known_airflow_exceptions.txt` allows 44 in
`databricks_base.py`, and the file now has 48, so the
`check-no-new-airflow-exceptions` static check will fail.
> **Flag any new `raise AirflowException` usage.** The community is reducing
direct `AirflowException` usage, not increasing it; new ones are not allowed
(enforced by the `check-no-new-airflow-exceptions` prek hook) — prefer Python's
standard exceptions (`ValueError`, `TypeError`, `OSError`), or a dedicated
class in the appropriate `exceptions.py`.
([code-review
instructions](https://github.com/apache/airflow/blob/main/.github/instructions/code-review.instructions.md#code-quality-rules))
Suggestion: raise `ValueError` for the missing `google_id_token_audience`.
For the fetch failures (empty token, exhausted retries, HTTP error), add a
dedicated class to `providers/databricks/exceptions.py`, next to
`DatabricksApiError`. The existing Azure branches still use `AirflowException`,
but they are already on the allowlist. Please don't copy that pattern into new
code.
##########
providers/databricks/tests/unit/databricks/hooks/test_databricks_google_id_token.py:
##########
@@ -0,0 +1,215 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import base64
+import json
+import time
+from unittest import mock
+
+import pytest
+
+from airflow.models import Connection
+from airflow.providers.common.compat.sdk import AirflowException,
AirflowOptionalProviderFeatureException
+from airflow.providers.databricks.hooks.databricks_base import (
+ DEFAULT_GCP_SCOPES,
+ GOOGLE_ID_TOKEN_SETTING_KEY,
+ BaseDatabricksHook,
+)
+
+AUDIENCE = "test-audience"
+IMPERSONATED_SERVICE_ACCOUNT =
"[email protected]"
+IMPERSONATE_URL = (
+ "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/";
+ f"{IMPERSONATED_SERVICE_ACCOUNT}:generateIdToken"
+)
+
+
+def _make_jwt(exp: int) -> str:
+ """Build an unsigned JWT string carrying only an ``exp`` claim (test
fixture)."""
+
+ def _b64(payload: dict) -> str:
+ return
base64.urlsafe_b64encode(json.dumps(payload).encode()).decode().rstrip("=")
+
+ return f"{_b64({'alg': 'none', 'typ': 'JWT'})}.{_b64({'exp':
exp})}.signature"
+
+
+def _valid_jwt() -> str:
+ return _make_jwt(int(time.time()) + 3600)
+
+
+def _expired_jwt() -> str:
+ return _make_jwt(int(time.time()) - 3600)
+
+
+def _hook_with_extra(extra: dict) -> BaseDatabricksHook:
+ hook = BaseDatabricksHook()
+ hook.databricks_conn = Connection(host="example.databricks.com",
extra=extra)
+ return hook
+
+
+class TestGoogleIdTokenDispatch:
+ @mock.patch(
+
"airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn",
+ new_callable=mock.PropertyMock,
+ )
+ @mock.patch.object(BaseDatabricksHook, "_get_google_id_token")
+ def test_get_token_google_id_token_success(self, mock_get_google_id_token,
mock_conn):
+ mock_conn.return_value =
Connection(extra={GOOGLE_ID_TOKEN_SETTING_KEY: True})
+ mock_get_google_id_token.return_value = "google_id_token"
+ hook = BaseDatabricksHook()
+ with mock.patch.object(hook.log, "debug") as mock_log_debug:
+ token = hook._get_token()
+ assert token == "google_id_token"
+ mock_get_google_id_token.assert_called_once_with()
+ mock_log_debug.assert_called_once_with("Using Google ID Token.")
+
+ @pytest.mark.asyncio
+ @mock.patch(
+
"airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook._a_get_google_id_token",
+ new_callable=mock.AsyncMock,
+ )
+ @mock.patch(
+
"airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn",
+ new_callable=mock.PropertyMock,
+ )
+ async def test_a_get_token_google_id_token_success(self, mock_conn,
mock_a_get_google_id_token):
+ mock_conn.return_value =
Connection(extra={GOOGLE_ID_TOKEN_SETTING_KEY: True})
+ mock_a_get_google_id_token.return_value = "google_id_token"
+ hook = BaseDatabricksHook()
+ with mock.patch.object(hook.log, "debug") as mock_log_debug:
+ token = await hook._a_get_token()
+ assert token == "google_id_token"
+ mock_a_get_google_id_token.assert_awaited_once_with()
+ mock_log_debug.assert_called_once_with("Using Google ID Token.")
+
+
+class TestGetGoogleIdToken:
Review Comment:
**minor**: three error branches of `_get_google_id_token` have no test: the
empty or missing `token` in the response (`"Google ID token couldn't be
obtained"`), exhausted retries (`RetryError`, for example a 503 on every
attempt), and a non-retryable `HTTPError` (for example 403 from IAM
Credentials). Reverting any of them leaves the suite green.
> **Flag any changed or added behaviour without a corresponding test, and
flag tests that a reviewer cannot fail by reverting the PR's change.** The
target is exactly 100% coverage of what the PR changes — no more, no less.
([code-review
instructions](https://github.com/apache/airflow/blob/main/.github/instructions/code-review.instructions.md#testing-requirements))
##########
providers/databricks/docs/connections/databricks.rst:
##########
@@ -137,6 +142,28 @@ Extra (optional)
a user inside workspace)
* ``azure_managed_identity_client_id``: optional client ID of the
user-assigned managed identity. This parameter is only required if you're using
a user-assigned managed identity. If not specified, the hook will attempt to
authenticate using a system-assigned managed identity.
+ Following parameters are necessary if using authentication with a
Google-signed OpenID Connect ID token
+ (requires the ``google`` extra, ``pip install
apache-airflow-providers-databricks[google]``):
+
+ * ``use_google_id_token``: required boolean flag to specify that a Google
ID token should be used
+ * ``google_id_token_audience``: required audience (``aud`` claim) of the
ID token, as configured for
Review Comment:
**nit**: for workspace-level API calls, the audience Databricks expects is
the workspace URL. The Databricks SDK's `google-id` provider passes
`target_audience=cfg.host`. You could default `google_id_token_audience` to
`https://<host>` and make it optional, which also makes the placeholder in the
JSON example below more concrete. Feel free to keep it required if you prefer
an explicit value.
##########
providers/databricks/tests/unit/databricks/hooks/test_databricks_google_id_token.py:
##########
@@ -0,0 +1,215 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import base64
+import json
+import time
+from unittest import mock
+
+import pytest
+
+from airflow.models import Connection
+from airflow.providers.common.compat.sdk import AirflowException,
AirflowOptionalProviderFeatureException
+from airflow.providers.databricks.hooks.databricks_base import (
+ DEFAULT_GCP_SCOPES,
+ GOOGLE_ID_TOKEN_SETTING_KEY,
+ BaseDatabricksHook,
+)
+
+AUDIENCE = "test-audience"
+IMPERSONATED_SERVICE_ACCOUNT =
"[email protected]"
+IMPERSONATE_URL = (
+ "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/";
+ f"{IMPERSONATED_SERVICE_ACCOUNT}:generateIdToken"
+)
+
+
+def _make_jwt(exp: int) -> str:
+ """Build an unsigned JWT string carrying only an ``exp`` claim (test
fixture)."""
+
+ def _b64(payload: dict) -> str:
+ return
base64.urlsafe_b64encode(json.dumps(payload).encode()).decode().rstrip("=")
+
+ return f"{_b64({'alg': 'none', 'typ': 'JWT'})}.{_b64({'exp':
exp})}.signature"
+
+
+def _valid_jwt() -> str:
+ return _make_jwt(int(time.time()) + 3600)
+
+
+def _expired_jwt() -> str:
+ return _make_jwt(int(time.time()) - 3600)
+
+
+def _hook_with_extra(extra: dict) -> BaseDatabricksHook:
+ hook = BaseDatabricksHook()
+ hook.databricks_conn = Connection(host="example.databricks.com",
extra=extra)
+ return hook
+
+
+class TestGoogleIdTokenDispatch:
+ @mock.patch(
+
"airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn",
+ new_callable=mock.PropertyMock,
+ )
+ @mock.patch.object(BaseDatabricksHook, "_get_google_id_token")
+ def test_get_token_google_id_token_success(self, mock_get_google_id_token,
mock_conn):
+ mock_conn.return_value =
Connection(extra={GOOGLE_ID_TOKEN_SETTING_KEY: True})
+ mock_get_google_id_token.return_value = "google_id_token"
+ hook = BaseDatabricksHook()
+ with mock.patch.object(hook.log, "debug") as mock_log_debug:
+ token = hook._get_token()
+ assert token == "google_id_token"
+ mock_get_google_id_token.assert_called_once_with()
+ mock_log_debug.assert_called_once_with("Using Google ID Token.")
+
+ @pytest.mark.asyncio
+ @mock.patch(
+
"airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook._a_get_google_id_token",
+ new_callable=mock.AsyncMock,
+ )
+ @mock.patch(
+
"airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn",
+ new_callable=mock.PropertyMock,
+ )
+ async def test_a_get_token_google_id_token_success(self, mock_conn,
mock_a_get_google_id_token):
+ mock_conn.return_value =
Connection(extra={GOOGLE_ID_TOKEN_SETTING_KEY: True})
+ mock_a_get_google_id_token.return_value = "google_id_token"
+ hook = BaseDatabricksHook()
+ with mock.patch.object(hook.log, "debug") as mock_log_debug:
+ token = await hook._a_get_token()
+ assert token == "google_id_token"
+ mock_a_get_google_id_token.assert_awaited_once_with()
+ mock_log_debug.assert_called_once_with("Using Google ID Token.")
+
+
+class TestGetGoogleIdToken:
+ def test_missing_audience_raises(self):
+ hook = _hook_with_extra({GOOGLE_ID_TOKEN_SETTING_KEY: True})
+ with pytest.raises(AirflowException, match="google_id_token_audience"):
+ hook._get_google_id_token()
+
+ @mock.patch("google.oauth2.id_token.fetch_id_token")
+ @mock.patch("google.auth.transport.requests.AuthorizedSession")
+ @mock.patch("google.auth.default")
+ def test_impersonated_service_account(self, mock_default,
mock_session_class, mock_fetch_id_token):
+ credentials = mock.Mock()
Review Comment:
**nit**: this mock has no spec. The same applies to the
`mock.patch.object(BaseDatabricksHook, "_get_google_id_token")` patches, which
have no `autospec=True`.
> **Flag any `mock.Mock()` or `mock.MagicMock()` without `spec` or
`autospec`.** Unspec'd mocks silently accept any attribute access, hiding real
bugs.
([code-review
instructions](https://github.com/apache/airflow/blob/main/.github/instructions/code-review.instructions.md#testing-requirements))
```suggestion
credentials = mock.Mock(spec=google.auth.credentials.Credentials)
```
(This needs `import google.auth.credentials` at the top of the test module.)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
