phanikumv commented on code in PR #69272:
URL: https://github.com/apache/airflow/pull/69272#discussion_r3518549639
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -893,12 +898,78 @@ def _get_required_client_id(self) -> str:
)
return client_id
+ def _get_federation_subject_token(self) -> tuple[str, str | None]:
+ """
+ Resolve the OIDC JWT to exchange for a Databricks token (RFC 8693
``subject_token``).
+
+ Two subject-token sources are supported:
+
+ * ``federated_token_provider`` -- a dotted path to a ``Callable[[],
str]`` that returns
+ the JWT. The token is obtained in-process and never written to disk,
so a control
+ plane can vend a short-lived, per-workload identity token for the
exchange. ``client_id``
+ is optional here: supply it in the extra for a service principal
federation policy, or
+ omit it for an account-wide federation policy.
+ * Kubernetes service account (default) -- read from the pod.
``client_id`` is required
+ because Kubernetes service account tokens cannot carry custom
claims, so only
+ service-principal-level federation is possible; it is validated
before the token is read.
+
+ :return: a ``(subject_token, client_id)`` tuple; ``client_id`` is
``None`` when the exchange
+ should omit it (account-wide federation policy).
+ """
+ provider =
self.databricks_conn.extra_dejson.get("federated_token_provider")
+ if provider:
+ return self._resolve_supplied_subject_token(provider),
self.databricks_conn.extra_dejson.get(
+ "client_id"
+ )
+ client_id = self._get_required_client_id()
+ return self._get_k8s_jwt_token(), client_id
+
+ async def _a_get_federation_subject_token(self) -> tuple[str, str | None]:
+ """Async version of :meth:`_get_federation_subject_token`."""
+ provider =
self.databricks_conn.extra_dejson.get("federated_token_provider")
+ if provider:
+ # The provider is a synchronous callable that typically makes a
blocking network call to
+ # mint the token. Offload it to a worker thread so it can't stall
the triggerer event loop.
+ loop = asyncio.get_running_loop()
+ subject_token = await loop.run_in_executor(None,
self._resolve_supplied_subject_token, provider)
+ return subject_token,
self.databricks_conn.extra_dejson.get("client_id")
+ client_id = self._get_required_client_id()
+ return await self._a_get_k8s_jwt_token(), client_id
+
+ def _resolve_supplied_subject_token(self, provider: str) -> str:
+ """
+ Import and invoke the ``federated_token_provider`` callable, returning
its OIDC JWT.
+
+ The dotted path is resolved and executed in-process; its return value
is the RFC 8693
+ ``subject_token`` and is never written to disk. Raises when the path
does not resolve to a
+ callable returning a non-empty string, so a misconfigured provider
fails with a clear error
+ rather than posting an empty ``subject_token`` to the exchange.
+ """
+ token_provider: Callable[[], str] = import_string(provider)
+ token = token_provider()
+ if not isinstance(token, str) or not token:
Review Comment:
Can we add .strip() , because the condition catches None, "", 0 and
non-strings, but a whitespace-only return (" ", "\t") is a non-empty str, so it
slips through...
--
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]
