ynpreet opened a new pull request, #75248:
URL: https://github.com/apache/airflow/pull/75248
## Description
Adds Google-signed OpenID Connect ID token authentication to the Databricks
provider, as requested in #57867, following the existing AAD token pattern in
`BaseDatabricksHook`.
When the new `use_google_id_token` connection extra is enabled, the hook
obtains
a Google ID token at runtime and uses it as the Bearer token for Databricks
API
calls:
- The token audience is configured with the `google_id_token_audience` extra
(required when the method is enabled).
- If `google_impersonated_service_account` is set, the default credentials
(e.g. from `GOOGLE_APPLICATION_CREDENTIALS`) impersonate that service
account
through the IAM Credentials `generateIdToken` endpoint
(`google.auth.default()` + `AuthorizedSession`), as in Google's documented
impersonation flow. Otherwise the token is fetched for the default
credentials
themselves via `google.oauth2.id_token.fetch_id_token`.
- Tokens are JWTs, so they are cached separately (`jwt_tokens`) and
validated by
their `exp` claim (`_is_jwt_token_valid`, local freshness check only — the
signature is verified by Databricks), mirroring the AAD cache/refresh
behavior
including the 120s refresh lead time, retry handling, and lazy
optional-import
errors (`AirflowOptionalProviderFeatureException` pointing at the existing
`google` provider extra).
The async twin `_a_get_google_id_token` offloads the blocking google-auth
calls
to a worker thread (the same approach `_a_get_federation_subject_token` uses
for
supplied subject tokens) rather than depending on google-auth's non-public
async
modules. Dispatch branches are added to both `_get_token()` and
`_a_get_token()`,
after the Azure branches. The new extras are read by name (allowlist) and
added
to `extra_parameters` so they are not forwarded to the Databricks SQL
connector.
Connection documentation is updated with the new method and extras.
Two naming notes for reviewers: the reporter's example used the extra keys
`audience` / `impersonated_service_account`; this PR uses `google_`-prefixed
keys
because plain `audience` is already used by the Kubernetes federation path
and
newer extras in this hook are namespaced (`aws_jwt_audience`,
`azure_tenant_id`).
Happy to switch if you prefer the shorter names.
Closes #57867
## Tests
New `tests/unit/databricks/hooks/test_databricks_google_id_token.py` (16
tests),
mirroring the AAD token tests — all google-auth entry points mocked, no real
GCP
calls:
- `_get_token` / `_a_get_token` dispatch to the new methods when the flag is
set
- missing `google_id_token_audience` raises
- impersonated fetch: asserts the IAM Credentials URL, request body
(`audience`, `includeEmail`), default-credential scopes, and token caching
- direct fetch via `fetch_id_token` with the configured audience
- valid cached JWT reused; expired JWT refetched
- missing google-auth raises `AirflowOptionalProviderFeatureException`
- `_is_jwt_token_valid` for valid / expired / malformed / `exp`-less tokens
Results (lightweight venv: released Airflow 3.3.2, provider installed
editable
from this branch, pytest `--noconftest`):
```
test_databricks_google_id_token.py: 16 passed
test_databricks_base.py (existing): 112 passed, 16 failed
```
The 16 existing-suite failures are identical on unmodified `main` in the same
environment (missing optional `aiofiles` / amazon provider packages in the
lightweight venv) — no regressions. `ruff check` / `ruff format --check` are
clean on the changed files. I could not run the full breeze suite locally;
CI should be the authoritative run.
--
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]
