You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by as...@apache.org on 2025/03/18 21:40:34 UTC
(airflow) branch main updated: Validate JWT tokens in the ExecutionAPI for strong task identity (#47885)
This is an automated email from the ASF dual-hosted git repository.
ash pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new c4e96d4ae3e Validate JWT tokens in the ExecutionAPI for strong task identity (#47885)
c4e96d4ae3e is described below
commit c4e96d4ae3ee99653f4c190b79f5eef85b83545e
Author: Ash Berlin-Taylor <as...@apache.org>
AuthorDate: Tue Mar 18 21:40:27 2025 +0000
Validate JWT tokens in the ExecutionAPI for strong task identity (#47885)
* Validate JWT tokens in the ExecutionAPI for strong task identity
This adds in Bearer auth token validation for the entire Execution API server.
Some notes about how this is achieved:
- We make heavy use of FastAPI's Dependency Injection to make this work
- We use the [`svcs`][1] module to make managing the service, and viewing the
state of it, easy
- In `.execution_api.routes` there are now two routers used -- one which
ensures that _everything_ underneath it has a valid ExecutionAPI JWT, and
another one which is "public" (and should contain very few routes)
- We ran into an issue (linked in code) with `__future__ annotations` imports,
and had to disable it on one specific file.
- Refreshing of expiring Task Identity tokens is not handled yet (this will
need middleware to support, a Dependency cannot add the response headers)
[1]: pypi.org/project/svcs
* Set `api_auth/jwt_secret` at deploy-time in our Kube tests
Without this being set the API server now crash-loops on startup, which adds
an extra 2 minutes to the test time for no reason.
---
airflow/api_fastapi/auth/tokens.py | 54 +++++++++++----
airflow/api_fastapi/execution_api/app.py | 59 ++++++++++++++--
.../api_fastapi/execution_api/datamodels/token.py | 7 +-
airflow/api_fastapi/execution_api/deps.py | 76 ++++++++++++++++++--
.../api_fastapi/execution_api/routes/__init__.py | 26 ++++---
.../execution_api/routes/connections.py | 52 ++++++--------
airflow/api_fastapi/execution_api/routes/health.py | 22 +++++-
.../execution_api/routes/task_instances.py | 11 ++-
.../api_fastapi/execution_api/routes/variables.py | 62 +++++++----------
airflow/api_fastapi/execution_api/routes/xcoms.py | 80 ++++++++--------------
airflow/config_templates/config.yml | 44 ++++++++++--
airflow/executors/base_executor.py | 26 +++++++
airflow/executors/workloads.py | 14 +++-
airflow/jobs/scheduler_job_runner.py | 2 +-
.../airflow_breeze/commands/kubernetes_commands.py | 2 +
hatch_build.py | 1 +
kubernetes_tests/test_base.py | 15 ----
tests/api_fastapi/execution_api/conftest.py | 15 +++-
.../execution_api/routes/test_connections.py | 31 +++++++--
.../execution_api/routes/test_task_instances.py | 46 +++++++++++++
.../execution_api/routes/test_variables.py | 48 ++++++++++---
.../api_fastapi/execution_api/routes/test_xcoms.py | 51 +++++++++++---
22 files changed, 534 insertions(+), 210 deletions(-)
diff --git a/airflow/api_fastapi/auth/tokens.py b/airflow/api_fastapi/auth/tokens.py
index a85c9ff8f27..10f71078171 100644
--- a/airflow/api_fastapi/auth/tokens.py
+++ b/airflow/api_fastapi/auth/tokens.py
@@ -119,10 +119,9 @@ class JWKS:
@classmethod
def from_private_key(cls, *keys: AllowedPrivateKeys | tuple[AllowedPrivateKeys, str]):
obj = cls(url=os.devnull)
-
keyset = [
# Each `key` is either the key directly or `(key, "my-kid")`
- key_to_jwk_dict(*key if isinstance(key, tuple) else key)
+ key_to_jwk_dict(*key) if isinstance(key, tuple) else key_to_jwk_dict(key)
for key in keys
]
obj._jwks = jwt.PyJWKSet(keyset)
@@ -209,6 +208,19 @@ class JWKS:
# It didn't load!
raise KeyError(f"Key ID {kid} not found in keyset")
+ def status(self):
+ # https://svcs.hynek.me/en/stable/core-concepts.html#health-checks
+ if not self._should_fetch_jwks():
+ # Up-to-date, we are healthy
+ return
+
+ if self.fetched_at == 0:
+ raise RuntimeError("JWKS never fetched")
+
+ last_successful_fetch = time.monotonic() - self.fetched_at
+ if last_successful_fetch > 3 * self.refresh_interval_secs:
+ raise RuntimeError(f"JWKS last fetched {last_successful_fetch}s ago")
+
def _conf_factory(section, key, **kwargs):
def factory() -> str:
@@ -273,7 +285,7 @@ class JWTValidator:
def _get_kid_from_header(self, unvalidated: str) -> str:
header = jwt.get_unverified_header(unvalidated)
if "kid" not in header:
- raise ValueError("Missing 'kid' in token header")
+ raise jwt.InvalidTokenError("Missing 'kid' in token header")
return header["kid"]
async def _get_validation_key(self, unvalidated: str) -> str | jwt.PyJWK:
@@ -316,6 +328,10 @@ class JWTValidator:
return claims
+ def status(self):
+ if self.jwks:
+ self.jwks.status()
+
def _pem_to_key(pem_data: str | bytes | AllowedPrivateKeys) -> AllowedPrivateKeys:
if isinstance(pem_data, str):
@@ -472,7 +488,7 @@ def base64url_encode(payload):
return encode.decode("utf-8").rstrip("=")
-def get_signing_key(section: str, key: str) -> str:
+def get_signing_key(section: str, key: str, make_secret_key_if_needed: bool = True) -> str:
"""
Get a signing shared key from the config.
@@ -484,18 +500,23 @@ def get_signing_key(section: str, key: str) -> str:
secret_key = conf.get(section, key, fallback=sentinel)
if not secret_key or secret_key is sentinel:
- log.warning(
- "`%s/%s` was empty, using a generated one for now. Please set this in your config", section, key
- )
- secret_key = base64url_encode(os.urandom(16))
- # Set it back so any other callers get the same value for the duration of this process
- conf.set(section, key, secret_key)
+ if make_secret_key_if_needed:
+ log.warning(
+ "`%s/%s` was empty, using a generated one for now. Please set this in your config",
+ section,
+ key,
+ )
+ secret_key = base64url_encode(os.urandom(16))
+ # Set it back so any other callers get the same value for the duration of this process
+ conf.set(section, key, secret_key)
+ else:
+ raise ValueError(f"The value {section}/{key} must be set!")
# Mypy can't grock the `if not secret_key`
return secret_key # type: ignore[return-value]
-def get_signing_args() -> dict[str, Any]:
+def get_signing_args(make_secret_key_if_needed: bool = True) -> dict[str, Any]:
"""
Return the args to splat into JWTGenerator for private or secret key.
@@ -508,10 +529,10 @@ def get_signing_args() -> dict[str, Any]:
return {"private_key": priv}
# Don't call this unless we have to as it might issue a warning
- return {"secret_key": get_signing_key("api_auth", "jwt_secret")}
+ return {"secret_key": get_signing_key("api_auth", "jwt_secret", make_secret_key_if_needed)}
-def get_sig_validation_args() -> dict[str, Any]:
+def get_sig_validation_args(make_secret_key_if_needed: bool = True) -> dict[str, Any]:
from airflow.configuration import conf
sentinel = object()
@@ -527,6 +548,9 @@ def get_sig_validation_args() -> dict[str, Any]:
if key is not None:
jwks = JWKS.from_private_key(key)
- return {"jwks": jwks}
+ return {
+ "jwks": jwks,
+ "algorithm": conf.get("api_auth", "jwt_algorithm", fallback=None) or _guess_best_algorithm(key),
+ }
- return {"secret_key": get_signing_key("api_auth", "jwt_secret")}
+ return {"secret_key": get_signing_key("api_auth", "jwt_secret", make_secret_key_if_needed)}
diff --git a/airflow/api_fastapi/execution_api/app.py b/airflow/api_fastapi/execution_api/app.py
index b6b36e08c11..317a3886434 100644
--- a/airflow/api_fastapi/execution_api/app.py
+++ b/airflow/api_fastapi/execution_api/app.py
@@ -17,15 +17,18 @@
from __future__ import annotations
-from contextlib import asynccontextmanager
+from contextlib import AsyncExitStack
from functools import cached_property
from typing import TYPE_CHECKING
import attrs
+import svcs
from fastapi import FastAPI, Request
from fastapi.openapi.utils import get_openapi
from fastapi.responses import JSONResponse
+from airflow.api_fastapi.auth.tokens import JWTValidator, get_sig_validation_args
+
if TYPE_CHECKING:
import httpx
@@ -34,10 +37,33 @@ import structlog
logger = structlog.get_logger(logger_name=__name__)
-@asynccontextmanager
-async def lifespan(app: FastAPI):
- """Context manager for the lifespan of the FastAPI app. For now does nothing."""
+def _jwt_validator():
+ from airflow.configuration import conf
+
+ required_claims = frozenset(["aud", "exp", "iat"])
+
+ if issuer := conf.get("api_auth", "jwt_issuer", fallback=None):
+ required_claims = required_claims | {"iss"}
+ validator = JWTValidator(
+ required_claims=required_claims,
+ issuer=issuer,
+ leeway=conf.getint("api_auth", "jwt_leeway"),
+ audience=conf.get_mandatory_list_value("execution_api", "jwt_audience"),
+ **get_sig_validation_args(make_secret_key_if_needed=False),
+ )
+ return validator
+
+
+@svcs.fastapi.lifespan
+async def lifespan(app: FastAPI, registry: svcs.Registry):
app.state.lifespan_called = True
+
+ # According to svcs's docs this shouldn't be needed, but something about SubApps is odd, and we need to
+ # record this here
+ app.state.svcs_registry = registry
+
+ # Create an app scoped validator, so that we don't have to fetch it every time
+ registry.register_value(JWTValidator, _jwt_validator(), ping=JWTValidator.status)
yield
@@ -129,22 +155,45 @@ class InProcessExecuctionAPI:
"""
_app: FastAPI | None = None
+ _cm: AsyncExitStack | None = None
@cached_property
def app(self):
if not self._app:
from airflow.api_fastapi.execution_api.app import create_task_execution_api_app
+ from airflow.api_fastapi.execution_api.deps import JWTBearerDep
+ from airflow.api_fastapi.execution_api.routes.connections import has_connection_access
+ from airflow.api_fastapi.execution_api.routes.variables import has_variable_access
+ from airflow.api_fastapi.execution_api.routes.xcoms import has_xcom_access
self._app = create_task_execution_api_app()
+ async def always_allow(): ...
+
+ self._app.dependency_overrides[JWTBearerDep.dependency] = always_allow
+ self._app.dependency_overrides[has_connection_access] = always_allow
+ self._app.dependency_overrides[has_variable_access] = always_allow
+ self._app.dependency_overrides[has_xcom_access] = always_allow
+
return self._app
@cached_property
def transport(self) -> httpx.WSGITransport:
+ import asyncio
+
import httpx
from a2wsgi import ASGIMiddleware
- return httpx.WSGITransport(app=ASGIMiddleware(self.app)) # type: ignore[arg-type]
+ middleware = ASGIMiddleware(self.app)
+
+ # https://github.com/abersheeran/a2wsgi/discussions/64
+ async def start_lifespan(cm: AsyncExitStack, app: FastAPI):
+ await cm.enter_async_context(app.router.lifespan_context(app))
+
+ self._cm = AsyncExitStack()
+
+ asyncio.run_coroutine_threadsafe(start_lifespan(self._cm, self.app), middleware.loop)
+ return httpx.WSGITransport(app=middleware) # type: ignore[arg-type]
@cached_property
def atransport(self) -> httpx.ASGITransport:
diff --git a/airflow/api_fastapi/execution_api/datamodels/token.py b/airflow/api_fastapi/execution_api/datamodels/token.py
index 568fdcf592a..43e52914fea 100644
--- a/airflow/api_fastapi/execution_api/datamodels/token.py
+++ b/airflow/api_fastapi/execution_api/datamodels/token.py
@@ -17,6 +17,9 @@
from __future__ import annotations
+from typing import Any
+from uuid import UUID
+
from airflow.api_fastapi.core_api.base import BaseModel
@@ -24,4 +27,6 @@ from airflow.api_fastapi.core_api.base import BaseModel
class TIToken(BaseModel):
"""Task Identity Token."""
- ti_key: str
+ id: UUID
+
+ claims: dict[str, Any]
diff --git a/airflow/api_fastapi/execution_api/deps.py b/airflow/api_fastapi/execution_api/deps.py
index 9e409bd3d6c..92c3f178d74 100644
--- a/airflow/api_fastapi/execution_api/deps.py
+++ b/airflow/api_fastapi/execution_api/deps.py
@@ -15,18 +15,80 @@
# specific language governing permissions and limitations
# under the License.
-from __future__ import annotations
+# Disable future annotations in this file to work around https://github.com/fastapi/fastapi/issues/13056
+# ruff: noqa: I002
-from typing import Annotated
+from typing import Any, Optional
-from fastapi import Depends
+import structlog
+import svcs
+from fastapi import Depends, HTTPException, Request, status
+from fastapi.security import HTTPBearer
+from airflow.api_fastapi.auth.tokens import JWTValidator
from airflow.api_fastapi.execution_api.datamodels.token import TIToken
+log = structlog.get_logger(logger_name=__name__)
-def get_task_token() -> TIToken:
- """TODO: Placeholder for task identity authentication. This should be replaced with actual JWT decoding and validation."""
- return TIToken(ti_key="test_key")
+# See https://github.com/fastapi/fastapi/issues/13056
+async def _container(request: Request):
+ async with svcs.Container(request.app.state.svcs_registry) as cont:
+ yield cont
-TokenDep = Annotated[TIToken, Depends(get_task_token)]
+
+DepContainer: svcs.Container = Depends(_container)
+
+
+class JWTBearer(HTTPBearer):
+ """
+ A FastAPI security dependency that validates JWT tokens using for the Execution API.
+
+ This will validate the tokens are signed and that the ``sub`` is a UUID, but nothing deeper than that.
+
+ The dependency result will be an `TIToken` object containing the ``id`` UUID (from the ``sub``) and other
+ validated claims.
+ """
+
+ def __init__(
+ self,
+ path_param_name: Optional[str] = None,
+ required_claims: Optional[dict[str, Any]] = None,
+ ):
+ super().__init__(auto_error=False)
+ self.path_param_name = path_param_name
+ self.required_claims = required_claims or {}
+
+ async def __call__( # type: ignore[override]
+ self,
+ request: Request,
+ services=DepContainer,
+ ) -> Optional[TIToken]:
+ creds = await super().__call__(request)
+ if not creds:
+ raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing auth token")
+
+ validator: JWTValidator = await services.aget(JWTValidator)
+
+ try:
+ # Example: Validate "task_instance_id" component of the path matches the one in the token
+ if self.path_param_name:
+ id = request.path_params[self.path_param_name]
+ validators: dict[str, Any] = {
+ **self.required_claims,
+ "sub": {"essential": True, "value": id},
+ }
+ else:
+ validators = self.required_claims
+ claims = await validator.avalidated_claims(creds.credentials, validators)
+ return TIToken(id=claims["sub"], claims=claims)
+ except Exception as err:
+ log.warning(
+ "Failed to validate JWT",
+ exc_info=True,
+ token=creds.credentials,
+ )
+ raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=f"Invalid auth token: {err}")
+
+
+JWTBearerDep: TIToken = Depends(JWTBearer())
diff --git a/airflow/api_fastapi/execution_api/routes/__init__.py b/airflow/api_fastapi/execution_api/routes/__init__.py
index 14c130b10a0..2f4dac92235 100644
--- a/airflow/api_fastapi/execution_api/routes/__init__.py
+++ b/airflow/api_fastapi/execution_api/routes/__init__.py
@@ -16,7 +16,9 @@
# under the License.
from __future__ import annotations
-from airflow.api_fastapi.common.router import AirflowRouter
+from fastapi import APIRouter
+
+from airflow.api_fastapi.execution_api.deps import JWTBearerDep
from airflow.api_fastapi.execution_api.routes import (
asset_events,
assets,
@@ -27,11 +29,17 @@ from airflow.api_fastapi.execution_api.routes import (
xcoms,
)
-execution_api_router = AirflowRouter()
-execution_api_router.include_router(assets.router, prefix="/assets", tags=["Assets"])
-execution_api_router.include_router(asset_events.router, prefix="/asset-events", tags=["Asset Events"])
-execution_api_router.include_router(connections.router, prefix="/connections", tags=["Connections"])
-execution_api_router.include_router(health.router, tags=["Health"])
-execution_api_router.include_router(task_instances.router, prefix="/task-instances", tags=["Task Instances"])
-execution_api_router.include_router(variables.router, prefix="/variables", tags=["Variables"])
-execution_api_router.include_router(xcoms.router, prefix="/xcoms", tags=["XComs"])
+execution_api_router = APIRouter()
+execution_api_router.include_router(health.router, prefix="/health", tags=["Health"])
+
+# _Every_ single endpoint under here must be authenticated. Some do further checks
+authenticated_router = APIRouter(dependencies=[JWTBearerDep]) # type: ignore[list-item]
+
+authenticated_router.include_router(assets.router, prefix="/assets", tags=["Assets"])
+authenticated_router.include_router(asset_events.router, prefix="/asset-events", tags=["Asset Events"])
+authenticated_router.include_router(connections.router, prefix="/connections", tags=["Connections"])
+authenticated_router.include_router(task_instances.router, prefix="/task-instances", tags=["Task Instances"])
+authenticated_router.include_router(variables.router, prefix="/variables", tags=["Variables"])
+authenticated_router.include_router(xcoms.router, prefix="/xcoms", tags=["XComs"])
+
+execution_api_router.include_router(authenticated_router)
diff --git a/airflow/api_fastapi/execution_api/routes/connections.py b/airflow/api_fastapi/execution_api/routes/connections.py
index ed72522ee8e..f950582e537 100644
--- a/airflow/api_fastapi/execution_api/routes/connections.py
+++ b/airflow/api_fastapi/execution_api/routes/connections.py
@@ -19,18 +19,32 @@ from __future__ import annotations
import logging
-from fastapi import HTTPException, status
+from fastapi import APIRouter, Depends, HTTPException, Path, status
-from airflow.api_fastapi.common.router import AirflowRouter
-from airflow.api_fastapi.execution_api import deps
from airflow.api_fastapi.execution_api.datamodels.connection import ConnectionResponse
-from airflow.api_fastapi.execution_api.datamodels.token import TIToken
+from airflow.api_fastapi.execution_api.deps import JWTBearerDep
from airflow.exceptions import AirflowNotFoundException
from airflow.models.connection import Connection
-# TODO: Add dependency on JWT token
-router = AirflowRouter(
+
+async def has_connection_access(
+ connection_id: str = Path(),
+ token=JWTBearerDep,
+) -> bool:
+ """Check if the task has access to the connection."""
+ # TODO: Placeholder for actual implementation
+
+ log.debug(
+ "Checking access for task instance with key '%s' to connection '%s'",
+ token.id,
+ connection_id,
+ )
+ return True
+
+
+router = APIRouter(
responses={status.HTTP_404_NOT_FOUND: {"description": "Connection not found"}},
+ dependencies=[Depends(has_connection_access)],
)
log = logging.getLogger(__name__)
@@ -43,19 +57,8 @@ log = logging.getLogger(__name__)
status.HTTP_403_FORBIDDEN: {"description": "Task does not have access to the connection"},
},
)
-def get_connection(
- connection_id: str,
- token: deps.TokenDep,
-) -> ConnectionResponse:
+def get_connection(connection_id: str) -> ConnectionResponse:
"""Get an Airflow connection."""
- if not has_connection_access(connection_id, token):
- raise HTTPException(
- status_code=status.HTTP_403_FORBIDDEN,
- detail={
- "reason": "access_denied",
- "message": f"Task does not have access to connection {connection_id}",
- },
- )
try:
connection = Connection.get_connection_from_secrets(connection_id)
except AirflowNotFoundException:
@@ -67,16 +70,3 @@ def get_connection(
},
)
return ConnectionResponse.model_validate(connection)
-
-
-def has_connection_access(connection_id: str, token: TIToken) -> bool:
- """Check if the task has access to the connection."""
- # TODO: Placeholder for actual implementation
-
- ti_key = token.ti_key
- log.debug(
- "Checking access for task instance with key '%s' to connection '%s'",
- ti_key,
- connection_id,
- )
- return True
diff --git a/airflow/api_fastapi/execution_api/routes/health.py b/airflow/api_fastapi/execution_api/routes/health.py
index 7bf4c4a2de0..bed519fdb2d 100644
--- a/airflow/api_fastapi/execution_api/routes/health.py
+++ b/airflow/api_fastapi/execution_api/routes/health.py
@@ -17,11 +17,31 @@
from __future__ import annotations
+from fastapi.responses import JSONResponse
+
from airflow.api_fastapi.common.router import AirflowRouter
+from airflow.api_fastapi.execution_api.deps import DepContainer
router = AirflowRouter()
-@router.get("/health")
+@router.get("")
def health() -> dict:
return {"status": "healthy"}
+
+
+@router.get("/ping")
+async def ping(services=DepContainer):
+ ok: list[str] = []
+ failing: dict[str, str] = {}
+ code = 200
+
+ for svc in services.get_pings():
+ try:
+ await svc.aping()
+ ok.append(svc.name)
+ except Exception as e:
+ failing[svc.name] = repr(e)
+ code = 500
+
+ return JSONResponse(content={"ok": ok, "failing": failing}, status_code=code)
diff --git a/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow/api_fastapi/execution_api/routes/task_instances.py
index e49897269cb..98b24d28d3b 100644
--- a/airflow/api_fastapi/execution_api/routes/task_instances.py
+++ b/airflow/api_fastapi/execution_api/routes/task_instances.py
@@ -22,7 +22,7 @@ import logging
from typing import Annotated
from uuid import UUID
-from fastapi import Body, HTTPException, status
+from fastapi import Body, Depends, HTTPException, status
from pydantic import JsonValue
from sqlalchemy import func, tuple_, update
from sqlalchemy.exc import NoResultFound, SQLAlchemyError
@@ -43,6 +43,7 @@ from airflow.api_fastapi.execution_api.datamodels.taskinstance import (
TISuccessStatePayload,
TITerminalStatePayload,
)
+from airflow.api_fastapi.execution_api.deps import JWTBearer
from airflow.models.dagrun import DagRun as DR
from airflow.models.taskinstance import TaskInstance as TI, _update_rtif
from airflow.models.taskreschedule import TaskReschedule
@@ -51,8 +52,12 @@ from airflow.models.xcom import XComModel
from airflow.utils import timezone
from airflow.utils.state import DagRunState, TaskInstanceState, TerminalTIState
-# TODO: Add dependency on JWT token
-router = AirflowRouter()
+router = AirflowRouter(
+ dependencies=[
+ # This checks that the UUID in the url matches the one in the token for us.
+ Depends(JWTBearer(path_param_name="task_instance_id")),
+ ]
+)
log = logging.getLogger(__name__)
diff --git a/airflow/api_fastapi/execution_api/routes/variables.py b/airflow/api_fastapi/execution_api/routes/variables.py
index 0e454f7dae0..649c89325ac 100644
--- a/airflow/api_fastapi/execution_api/routes/variables.py
+++ b/airflow/api_fastapi/execution_api/routes/variables.py
@@ -19,17 +19,33 @@ from __future__ import annotations
import logging
-from fastapi import HTTPException, status
+from fastapi import APIRouter, Depends, HTTPException, Path, Request, status
-from airflow.api_fastapi.common.router import AirflowRouter
-from airflow.api_fastapi.execution_api import deps
-from airflow.api_fastapi.execution_api.datamodels.token import TIToken
from airflow.api_fastapi.execution_api.datamodels.variable import VariablePostBody, VariableResponse
+from airflow.api_fastapi.execution_api.deps import JWTBearerDep
from airflow.models.variable import Variable
-# TODO: Add dependency on JWT token
-router = AirflowRouter(
+
+async def has_variable_access(
+ request: Request,
+ variable_key: str = Path(),
+ token=JWTBearerDep,
+):
+ """Check if the task has access to the variable."""
+ write = request.method not in {"GET", "HEAD", "OPTIONS"}
+ # TODO: Placeholder for actual implementation
+ log.debug(
+ "Checking %s access for task instance with key '%s' to variable '%s'",
+ "write" if write else "read",
+ token.id,
+ variable_key,
+ )
+ return True
+
+
+router = APIRouter(
responses={status.HTTP_404_NOT_FOUND: {"description": "Variable not found"}},
+ dependencies=[Depends(has_variable_access)],
)
log = logging.getLogger(__name__)
@@ -42,17 +58,8 @@ log = logging.getLogger(__name__)
status.HTTP_403_FORBIDDEN: {"description": "Task does not have access to the variable"},
},
)
-def get_variable(variable_key: str, token: deps.TokenDep) -> VariableResponse:
+def get_variable(variable_key: str) -> VariableResponse:
"""Get an Airflow Variable."""
- if not has_variable_access(variable_key, token):
- raise HTTPException(
- status_code=status.HTTP_403_FORBIDDEN,
- detail={
- "reason": "access_denied",
- "message": f"Task does not have access to variable {variable_key}",
- },
- )
-
try:
variable_value = Variable.get(variable_key)
except KeyError:
@@ -75,28 +82,7 @@ def get_variable(variable_key: str, token: deps.TokenDep) -> VariableResponse:
status.HTTP_403_FORBIDDEN: {"description": "Task does not have access to the variable"},
},
)
-def put_variable(variable_key: str, body: VariablePostBody, token: deps.TokenDep):
+def put_variable(variable_key: str, body: VariablePostBody):
"""Set an Airflow Variable."""
- if not has_variable_access(variable_key, token, write_access=True):
- raise HTTPException(
- status_code=status.HTTP_403_FORBIDDEN,
- detail={
- "reason": "access_denied",
- "message": f"Task does not have access to write variable {variable_key}",
- },
- )
Variable.set(key=variable_key, value=body.value, description=body.description)
return {"message": "Variable successfully set"}
-
-
-def has_variable_access(variable_key: str, token: TIToken, write_access: bool = False) -> bool:
- """Check if the task has access to the variable."""
- # TODO: Placeholder for actual implementation
-
- ti_key = token.ti_key
- log.debug(
- "Checking access for task instance with key '%s' to variable '%s'",
- ti_key,
- variable_key,
- )
- return True
diff --git a/airflow/api_fastapi/execution_api/routes/xcoms.py b/airflow/api_fastapi/execution_api/routes/xcoms.py
index 92750835b4f..3612f3615ab 100644
--- a/airflow/api_fastapi/execution_api/routes/xcoms.py
+++ b/airflow/api_fastapi/execution_api/routes/xcoms.py
@@ -20,26 +20,48 @@ from __future__ import annotations
import logging
from typing import Annotated
-from fastapi import Body, Depends, HTTPException, Query, Response, status
+from fastapi import Body, Depends, HTTPException, Path, Query, Request, Response, status
from pydantic import JsonValue
from sqlalchemy import delete
from sqlalchemy.sql.selectable import Select
from airflow.api_fastapi.common.db.common import SessionDep
from airflow.api_fastapi.common.router import AirflowRouter
-from airflow.api_fastapi.execution_api import deps
-from airflow.api_fastapi.execution_api.datamodels.token import TIToken
from airflow.api_fastapi.execution_api.datamodels.xcom import XComResponse
+from airflow.api_fastapi.execution_api.deps import JWTBearerDep
from airflow.models.taskmap import TaskMap
from airflow.models.xcom import XComModel
from airflow.utils.db import get_query_count
-# TODO: Add dependency on JWT token
+
+async def has_xcom_access(
+ dag_id: str,
+ run_id: str,
+ task_id: str,
+ xcom_key: Annotated[str, Path(alias="key")],
+ request: Request,
+ token=JWTBearerDep,
+) -> bool:
+ """Check if the task has access to the XCom."""
+ # TODO: Placeholder for actual implementation
+
+ write = request.method not in {"GET", "HEAD", "OPTIONS"}
+
+ log.debug(
+ "Checking %s XCom access for xcom from TaskInstance with key '%s' to XCom '%s'",
+ "write" if write else "read",
+ token.id,
+ xcom_key,
+ )
+ return True
+
+
router = AirflowRouter(
responses={
status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
status.HTTP_403_FORBIDDEN: {"description": "Task does not have access to the XCom"},
},
+ dependencies=[Depends(has_xcom_access)],
)
log = logging.getLogger(__name__)
@@ -51,18 +73,8 @@ async def xcom_query(
task_id: str,
key: str,
session: SessionDep,
- token: deps.TokenDep,
map_index: Annotated[int | None, Query()] = None,
) -> Select:
- if not has_xcom_access(dag_id, run_id, task_id, key, token):
- raise HTTPException(
- status_code=status.HTTP_403_FORBIDDEN,
- detail={
- "reason": "access_denied",
- "message": f"Task does not have access to XCom key '{key}'",
- },
- )
-
query = XComModel.get_many(
run_id=run_id,
key=key,
@@ -92,7 +104,6 @@ async def xcom_query(
)
def head_xcom(
response: Response,
- token: deps.TokenDep,
session: SessionDep,
xcom_query: Annotated[Select, Depends(xcom_query)],
map_index: Annotated[int | None, Query()] = None,
@@ -116,7 +127,6 @@ def head_xcom(
description="Get a single XCom Value",
)
def get_xcom(
- session: SessionDep,
dag_id: str,
run_id: str,
task_id: str,
@@ -180,7 +190,6 @@ def set_xcom(
},
),
],
- token: deps.TokenDep,
session: SessionDep,
map_index: Annotated[int, Query()] = -1,
mapped_length: Annotated[
@@ -190,15 +199,6 @@ def set_xcom(
"""Set an Airflow XCom."""
from airflow.configuration import conf
- if not has_xcom_access(dag_id, run_id, task_id, key, token, write=True):
- raise HTTPException(
- status_code=status.HTTP_403_FORBIDDEN,
- detail={
- "reason": "access_denied",
- "message": f"Task does not have access to set XCom key '{key}'",
- },
- )
-
if mapped_length is not None:
task_map = TaskMap(
dag_id=dag_id,
@@ -275,23 +275,13 @@ def set_xcom(
)
def delete_xcom(
session: SessionDep,
- token: deps.TokenDep,
- key: str,
dag_id: str,
run_id: str,
task_id: str,
+ key: str,
map_index: Annotated[int, Query()] = -1,
):
"""Delete a single XCom Value."""
- if not has_xcom_access(dag_id, run_id, task_id, key, token, write=True):
- raise HTTPException(
- status_code=status.HTTP_403_FORBIDDEN,
- detail={
- "reason": "access_denied",
- "message": f"Task does not have access to delete XCom with key '{key}'",
- },
- )
-
query = delete(XComModel).where(
XComModel.key == key,
XComModel.run_id == run_id,
@@ -302,19 +292,3 @@ def delete_xcom(
session.execute(query)
session.commit()
return {"message": f"XCom with key: {key} successfully deleted."}
-
-
-def has_xcom_access(
- dag_id: str, run_id: str, task_id: str, xcom_key: str, token: TIToken, write: bool = False
-) -> bool:
- """Check if the task has access to the XCom."""
- # TODO: Placeholder for actual implementation
-
- ti_key = token.ti_key
- log.debug(
- "Checking %s XCom access for xcom from TaskInstance with key '%s' to XCom '%s'",
- "write" if write else "read",
- ti_key,
- xcom_key,
- )
- return True
diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml
index f6804ecaf7a..e771f0ffabf 100644
--- a/airflow/config_templates/config.yml
+++ b/airflow/config_templates/config.yml
@@ -1427,24 +1427,28 @@ api_auth:
This variable can be a single value, or a comma-separated string, in which case the first value is the
one that will be used when generating, and the others are accepted at validation time.
- Not required, but strongly encouraged
+ Not required, but strongly encouraged.
+
+ See also :ref:`config:execution_api__jwt_audience`
example: "my-unique-airflow-id"
default: ~
type: string
jwt_expiration_time:
description: |
- Number in seconds until the JWT token used for authentication expires. When the token expires,
+ Number in seconds until the JWTs used for authentication expires. When the token expires,
all API calls using this token will fail on authentication.
Make sure that time on ALL the machines that you run airflow components on is synchronized
(for example using ntpd) otherwise you might get "forbidden" errors.
+
+ See also :ref:`config:execution_api__jwt_expiration_time`
version_added: 3.0.0
type: integer
example: ~
default: "86400"
jwt_cli_expiration_time:
description: |
- Number in seconds until the JWT token used for authentication expires for CLI commands.
+ Number in seconds until the JWTs used for authentication expires for CLI commands.
When the token expires, all CLI calls using this token will fail on authentication.
Make sure that time on ALL the machines that you run airflow components on is synchronized
@@ -1455,7 +1459,7 @@ api_auth:
default: "3600"
jwt_secret:
description: |
- Secret key used to encode and decode JWT tokens to authenticate to public and private APIs.
+ Secret key used to encode and decode JWTs to authenticate to public and private APIs.
It should be as random as possible. However, when running more than 1 instances of API services,
make sure all of them use the same ``jwt_secret`` otherwise calls will fail on authentication.
@@ -1522,12 +1526,42 @@ api_auth:
jwt_leeway:
version_added: 3.0.0
description: |
- Number of seconds leeway in validating expiry time of JWT tokens to account for clock skew between
+ Number of seconds leeway in validating expiry time of JWTs to account for clock skew between
client and server
default: "10"
example: ~
type: integer
+execution_api:
+ description: |
+ Settings related to the Execution API server.
+
+ The ExecutionAPI also uses a lot of settings from the :ref:`config:api_auth` section.
+ options:
+ jwt_expiration_time:
+ description: |
+ Number in seconds until the JWT used for authentication expires. When the token expires,
+ all API calls using this token will fail on authentication.
+ Make sure that time on ALL the machines that you run airflow components on is synchronized
+ (for example using ntpd) otherwise you might get "forbidden" errors.
+ version_added: 3.0.0
+ type: integer
+ example: ~
+ default: "600"
+ jwt_audience:
+ version_added: 3.0.0
+ description: |
+ The audience claim to use when generating and validating JWTs for the Execution API.
+
+ This variable can be a single value, or a comma-separated string, in which case the first value is the
+ one that will be used when generating, and the others are accepted at validation time.
+
+ Not required, but strongly encouraged
+
+ See also :ref:`config:api_auth__jwt_audience`
+ default: "urn:airflow.apache.org:task"
+ example: ~
+ type: string
lineage:
description: ~
options:
diff --git a/airflow/executors/base_executor.py b/airflow/executors/base_executor.py
index f1b6b936f9f..c9e63c6ca19 100644
--- a/airflow/executors/base_executor.py
+++ b/airflow/executors/base_executor.py
@@ -23,6 +23,7 @@ import sys
from collections import defaultdict, deque
from collections.abc import Sequence
from dataclasses import dataclass, field
+from functools import cached_property
from typing import TYPE_CHECKING, Any, Optional
import pendulum
@@ -46,6 +47,7 @@ if TYPE_CHECKING:
from sqlalchemy.orm import Session
+ from airflow.api_fastapi.auth.tokens import JWTGenerator
from airflow.callbacks.base_callback_sink import BaseCallbackSink
from airflow.callbacks.callback_requests import CallbackRequest
from airflow.cli.cli_config import GroupCommand
@@ -128,8 +130,32 @@ class BaseExecutor(LoggingMixin):
name: None | ExecutorName = None
callback_sink: BaseCallbackSink | None = None
+ @cached_property
+ def jwt_generator(self) -> JWTGenerator:
+ from airflow.api_fastapi.auth.tokens import (
+ JWTGenerator,
+ get_signing_args,
+ )
+ from airflow.configuration import conf
+
+ generator = JWTGenerator(
+ valid_for=conf.getint("execution_api", "jwt_expiration_time"),
+ audience=conf.get_mandatory_list_value("execution_api", "jwt_audience")[0],
+ issuer=conf.get("api_auth", "jwt_issuer", fallback=None),
+ # Since this one is used across components/server, there is no point trying to generate one, error
+ # instead
+ **get_signing_args(make_secret_key_if_needed=False),
+ )
+
+ return generator
+
def __init__(self, parallelism: int = PARALLELISM, team_id: str | None = None):
super().__init__()
+ # Ensure we set this now, so that each subprocess gets the same value
+ from airflow.api_fastapi.auth.tokens import get_signing_args
+
+ get_signing_args()
+
self.parallelism: int = parallelism
self.team_id: str | None = team_id
self.queued_tasks: dict[TaskInstanceKey, QueuedTaskInstanceType] = {}
diff --git a/airflow/executors/workloads.py b/airflow/executors/workloads.py
index f42a0fe062f..cd6c66c33db 100644
--- a/airflow/executors/workloads.py
+++ b/airflow/executors/workloads.py
@@ -22,9 +22,11 @@ from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, Literal, Union
+import structlog
from pydantic import BaseModel, Field
if TYPE_CHECKING:
+ from airflow.api_fastapi.auth.tokens import JWTGenerator
from airflow.models.taskinstance import TaskInstance as TIModel
from airflow.models.taskinstancekey import TaskInstanceKey
@@ -34,6 +36,8 @@ __all__ = [
"ExecuteTask",
]
+log = structlog.get_logger()
+
class BaseWorkload(BaseModel):
token: str
@@ -93,7 +97,9 @@ class ExecuteTask(BaseWorkload):
kind: Literal["ExecuteTask"] = Field(init=False, default="ExecuteTask")
@classmethod
- def make(cls, ti: TIModel, dag_rel_path: Path | None = None) -> ExecuteTask:
+ def make(
+ cls, ti: TIModel, dag_rel_path: Path | None = None, generator: JWTGenerator | None = None
+ ) -> ExecuteTask:
from pathlib import Path
from airflow.utils.helpers import log_filename_template_renderer
@@ -104,10 +110,14 @@ class ExecuteTask(BaseWorkload):
version=ti.dag_run.bundle_version,
)
fname = log_filename_template_renderer()(ti=ti)
+ token = ""
+
+ if generator:
+ token = generator.generate({"sub": str(ti.id)})
return cls(
ti=ser_ti,
dag_rel_path=dag_rel_path or Path(ti.dag_model.relative_fileloc),
- token="",
+ token=token,
log_path=fname,
bundle_info=bundle_info,
)
diff --git a/airflow/jobs/scheduler_job_runner.py b/airflow/jobs/scheduler_job_runner.py
index a8b5ccced84..b4660c8dad0 100644
--- a/airflow/jobs/scheduler_job_runner.py
+++ b/airflow/jobs/scheduler_job_runner.py
@@ -646,7 +646,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
# TODO: Task-SDK: This check is transitionary. Remove once all executors are ported over.
# Has a real queue_activity implemented
if executor.queue_workload.__func__ is not BaseExecutor.queue_workload: # type: ignore[attr-defined]
- workload = workloads.ExecuteTask.make(ti)
+ workload = workloads.ExecuteTask.make(ti, generator=executor.jwt_generator)
executor.queue_workload(workload, session=session)
continue
diff --git a/dev/breeze/src/airflow_breeze/commands/kubernetes_commands.py b/dev/breeze/src/airflow_breeze/commands/kubernetes_commands.py
index 03c191ee154..f3b29083513 100644
--- a/dev/breeze/src/airflow_breeze/commands/kubernetes_commands.py
+++ b/dev/breeze/src/airflow_breeze/commands/kubernetes_commands.py
@@ -1030,6 +1030,8 @@ def _deploy_helm_chart(
f"executor={executor}",
"--set",
f"airflowVersion={params.airflow_semver_version}",
+ "--set",
+ "config.api_auth.jwt_secret=foo",
]
if multi_namespace_mode:
helm_command.extend(["--set", "multiNamespaceMode=true"])
diff --git a/hatch_build.py b/hatch_build.py
index 11109f13935..341a96d366d 100644
--- a/hatch_build.py
+++ b/hatch_build.py
@@ -262,6 +262,7 @@ DEPENDENCIES = [
"sqlalchemy>=1.4.49,<2.0",
"sqlalchemy-jsonfield>=1.0",
"sqlalchemy-utils>=0.41.2",
+ "svcs>=25.1.0",
"tabulate>=0.7.5",
"tenacity>=8.0.0,!=8.2.0",
"termcolor>=2.5.0",
diff --git a/kubernetes_tests/test_base.py b/kubernetes_tests/test_base.py
index 50651485bbe..7d595f745ca 100644
--- a/kubernetes_tests/test_base.py
+++ b/kubernetes_tests/test_base.py
@@ -39,9 +39,6 @@ KUBERNETES_HOST_PORT = (os.environ.get("CLUSTER_HOST") or "localhost") + ":" + C
EXECUTOR = os.environ.get("EXECUTOR")
CONFIG_MAP_NAME = "airflow-config"
CONFIG_MAP_KEY = "airflow.cfg"
-AIRFLOW_API_SERVER_JWT_SECRET = os.environ.get(
- "AIRFLOW_API_SERVER_JWT_SECRET", "airflow_api_server_jwt_secret"
-)
print()
print(f"Cluster host/port used: ${KUBERNETES_HOST_PORT}")
@@ -69,9 +66,6 @@ class BaseK8STest:
if self.set_api_server_base_url_config():
self.rollout_restart_deployment("airflow-api-server")
self.ensure_deployment_health("airflow-api-server")
- if self.set_api_auth_jwt_secret_config():
- self.rollout_restart_deployment("airflow-api-server")
- self.ensure_deployment_health("airflow-api-server")
# Replacement for unittests.TestCase.id()
self.test_id = f"{request.node.cls.__name__}_{request.node.name}"
@@ -351,15 +345,6 @@ class BaseK8STest:
"api", "base_url", f"http://{KUBERNETES_HOST_PORT}"
)
- def set_api_auth_jwt_secret_config(self) -> bool:
- """Set [api_auth/jwt_secret] with AIRFLOW_API_SERVER_JWT_SECRET as env in k8s configmap."
-
- :return: True if the configmap was updated successfully, False otherwise"
- """
- return self.set_airflow_cfg_in_kubernetes_configmap(
- "api_auth", "jwt_secret", AIRFLOW_API_SERVER_JWT_SECRET
- )
-
def ensure_dag_expected_state(self, host, logical_date, dag_id, expected_final_state, timeout):
tries = 0
state = ""
diff --git a/tests/api_fastapi/execution_api/conftest.py b/tests/api_fastapi/execution_api/conftest.py
index 784cb29249a..7ff9e31c26d 100644
--- a/tests/api_fastapi/execution_api/conftest.py
+++ b/tests/api_fastapi/execution_api/conftest.py
@@ -16,12 +16,23 @@
# under the License.
from __future__ import annotations
+from unittest.mock import AsyncMock
+
import pytest
from fastapi.testclient import TestClient
from airflow.api_fastapi.app import cached_app
+from airflow.api_fastapi.auth.tokens import JWTValidator
+from airflow.api_fastapi.execution_api.app import lifespan
@pytest.fixture
-def client():
- return TestClient(cached_app(apps="execution"))
+def client(request: pytest.FixtureRequest):
+ app = cached_app(apps="execution")
+ with TestClient(app, headers={"Authorization": "Bearer fake"}) as client:
+ auth = AsyncMock(spec=JWTValidator)
+ auth.avalidated_claims.return_value = {"sub": "edb09971-4e0e-4221-ad3f-800852d38085"}
+
+ # Inject our fake JWTValidator object. Can be over-ridden by tests if they want
+ lifespan.registry.register_value(JWTValidator, auth)
+ yield client
diff --git a/tests/api_fastapi/execution_api/routes/test_connections.py b/tests/api_fastapi/execution_api/routes/test_connections.py
index 287a5bac801..a2e3cb51fab 100644
--- a/tests/api_fastapi/execution_api/routes/test_connections.py
+++ b/tests/api_fastapi/execution_api/routes/test_connections.py
@@ -20,12 +20,37 @@ from __future__ import annotations
from unittest import mock
import pytest
+from fastapi import FastAPI, HTTPException, status
from airflow.models.connection import Connection
pytestmark = pytest.mark.db_test
+@pytest.fixture
+def access_denied(client):
+ from airflow.api_fastapi.execution_api.routes.connections import has_connection_access
+
+ last_route = client.app.routes[-1]
+ assert isinstance(last_route.app, FastAPI)
+ exec_app = last_route.app
+
+ async def _(connection_id: str):
+ raise HTTPException(
+ status_code=status.HTTP_403_FORBIDDEN,
+ detail={
+ "reason": "access_denied",
+ "message": f"Task does not have access to connection {connection_id}",
+ },
+ )
+
+ exec_app.dependency_overrides[has_connection_access] = _
+
+ yield
+
+ exec_app.dependency_overrides = {}
+
+
class TestGetConnection:
def test_connection_get_from_db(self, client, session):
connection = Connection(
@@ -91,11 +116,9 @@ class TestGetConnection:
}
}
+ @pytest.mark.usefixtures("access_denied")
def test_connection_get_access_denied(self, client):
- with mock.patch(
- "airflow.api_fastapi.execution_api.routes.connections.has_connection_access", return_value=False
- ):
- response = client.get("/execution/connections/test_conn")
+ response = client.get("/execution/connections/test_conn")
# Assert response status code and detail for access denied
assert response.status_code == 403
diff --git a/tests/api_fastapi/execution_api/routes/test_task_instances.py b/tests/api_fastapi/execution_api/routes/test_task_instances.py
index e784043920d..500c730a783 100644
--- a/tests/api_fastapi/execution_api/routes/test_task_instances.py
+++ b/tests/api_fastapi/execution_api/routes/test_task_instances.py
@@ -23,9 +23,15 @@ from unittest import mock
import pytest
import uuid6
+from fastapi import FastAPI
+from fastapi.routing import Mount
from sqlalchemy import select, update
from sqlalchemy.exc import SQLAlchemyError
+from airflow.api_fastapi.app import purge_cached_app
+from airflow.api_fastapi.auth.tokens import JWTValidator
+from airflow.api_fastapi.execution_api.app import lifespan
+from airflow.api_fastapi.execution_api.routes.task_instances import router as task_instance_router
from airflow.models import RenderedTaskInstanceFields, TaskReschedule, Trigger
from airflow.models.asset import AssetActive, AssetAliasModel, AssetEvent, AssetModel
from airflow.models.taskinstance import TaskInstance
@@ -56,6 +62,46 @@ def _create_asset_aliases(session, num: int = 2) -> None:
session.commit()
+@pytest.fixture
+def add_foo_test_route(client):
+ @task_instance_router.get("/{task_instance_id}/foo")
+ def foo(task_instance_id: str):
+ return {"hi": task_instance_id}
+
+ app: FastAPI = client.app
+
+ last_route = app.routes[-1]
+ assert isinstance(last_route, Mount)
+ assert isinstance(last_route.app, FastAPI)
+ # Re-add it, so it gets the new route we've added
+ last_route.app.include_router(task_instance_router, prefix="/task-instances")
+
+ yield
+
+ purge_cached_app()
+
+
+@pytest.mark.usefixtures("add_foo_test_route")
+def test_id_matches_sub_claim(client):
+ # Test that this is validated at the router level, so we don't have to test it in each component
+ validator = mock.AsyncMock(spec=JWTValidator)
+ claims = {"sub": "edb09971-4e0e-4221-ad3f-800852d38085"}
+ validator.avalidated_claims.side_effect = [claims, RuntimeError("Fail for test")]
+
+ lifespan.registry.register_value(JWTValidator, validator)
+
+ resp = client.get(f"/execution/task-instances/{claims['sub']}/foo")
+ assert resp.status_code == 200
+
+ validator.avalidated_claims.assert_awaited_once()
+
+ resp = client.get("/execution/task-instances/9c230b40-da03-451d-8bd7-be30471be383/foo")
+ assert resp.status_code == 403
+ validator.avalidated_claims.assert_called_with(
+ mock.ANY, {"sub": {"essential": True, "value": "9c230b40-da03-451d-8bd7-be30471be383"}}
+ )
+
+
class TestTIRunState:
def setup_method(self):
clear_db_runs()
diff --git a/tests/api_fastapi/execution_api/routes/test_variables.py b/tests/api_fastapi/execution_api/routes/test_variables.py
index 20a9b43c07a..01d5a77f7be 100644
--- a/tests/api_fastapi/execution_api/routes/test_variables.py
+++ b/tests/api_fastapi/execution_api/routes/test_variables.py
@@ -17,9 +17,12 @@
from __future__ import annotations
+import logging
from unittest import mock
import pytest
+from fastapi import FastAPI, HTTPException, Request, status
+from fastapi.routing import Mount
from airflow.models.variable import Variable
@@ -35,6 +38,32 @@ def setup_method():
clear_db_variables()
+@pytest.fixture
+def access_denied(client):
+ from airflow.api_fastapi.execution_api.deps import JWTBearerDep
+ from airflow.api_fastapi.execution_api.routes.variables import has_variable_access
+
+ last_route = client.app.routes[-1]
+ assert isinstance(last_route, Mount)
+ assert isinstance(last_route.app, FastAPI)
+ exec_app = last_route.app
+
+ async def _(request: Request, variable_key: str, token=JWTBearerDep):
+ await has_variable_access(request, variable_key, token)
+ raise HTTPException(
+ status_code=status.HTTP_403_FORBIDDEN,
+ detail={
+ "reason": "access_denied",
+ },
+ )
+
+ exec_app.dependency_overrides[has_variable_access] = _
+
+ yield
+
+ exec_app.dependency_overrides = {}
+
+
class TestGetVariable:
def test_variable_get_from_db(self, client, session):
Variable.set(key="var1", value="value", session=session)
@@ -70,10 +99,9 @@ class TestGetVariable:
}
}
- def test_variable_get_access_denied(self, client):
- with mock.patch(
- "airflow.api_fastapi.execution_api.routes.variables.has_variable_access", return_value=False
- ):
+ @pytest.mark.usefixtures("access_denied")
+ def test_variable_get_access_denied(self, client, caplog):
+ with caplog.at_level(logging.DEBUG):
response = client.get("/execution/variables/key1")
# Assert response status code and detail for access denied
@@ -81,10 +109,11 @@ class TestGetVariable:
assert response.json() == {
"detail": {
"reason": "access_denied",
- "message": "Task does not have access to variable key1",
}
}
+ assert any(msg.startswith("Checking read access for task instance") for msg in caplog.messages)
+
class TestPutVariable:
@pytest.mark.parametrize(
@@ -168,10 +197,9 @@ class TestPutVariable:
assert var_from_db.key == key
assert var_from_db.val == payload["value"]
- def test_post_variable_access_denied(self, client):
- with mock.patch(
- "airflow.api_fastapi.execution_api.routes.variables.has_variable_access", return_value=False
- ):
+ @pytest.mark.usefixtures("access_denied")
+ def test_post_variable_access_denied(self, client, caplog):
+ with caplog.at_level(logging.DEBUG):
key = "var_create"
payload = {"value": "{}"}
response = client.put(
@@ -184,6 +212,6 @@ class TestPutVariable:
assert response.json() == {
"detail": {
"reason": "access_denied",
- "message": "Task does not have access to write variable var_create",
}
}
+ assert any(msg.startswith("Checking write access for task instance") for msg in caplog.messages)
diff --git a/tests/api_fastapi/execution_api/routes/test_xcoms.py b/tests/api_fastapi/execution_api/routes/test_xcoms.py
index f08c8371e28..ae21fefdbf4 100644
--- a/tests/api_fastapi/execution_api/routes/test_xcoms.py
+++ b/tests/api_fastapi/execution_api/routes/test_xcoms.py
@@ -19,10 +19,11 @@
from __future__ import annotations
import contextlib
-from unittest import mock
+import logging
import httpx
import pytest
+from fastapi import FastAPI, HTTPException, Path, Request, status
from airflow.api_fastapi.execution_api.datamodels.xcom import XComResponse
from airflow.models.dagrun import DagRun
@@ -42,6 +43,38 @@ def reset_db():
session.query(XComModel).delete()
+@pytest.fixture
+def access_denied(client):
+ from airflow.api_fastapi.execution_api.deps import JWTBearerDep
+ from airflow.api_fastapi.execution_api.routes.xcoms import has_xcom_access
+
+ last_route = client.app.routes[-1]
+ assert isinstance(last_route.app, FastAPI)
+ exec_app = last_route.app
+
+ async def _(
+ request: Request,
+ dag_id: str = Path(),
+ run_id: str = Path(),
+ task_id: str = Path(),
+ xcom_key: str = Path(alias="key"),
+ token=JWTBearerDep,
+ ):
+ await has_xcom_access(dag_id, run_id, task_id, xcom_key, request, token)
+ raise HTTPException(
+ status_code=status.HTTP_403_FORBIDDEN,
+ detail={
+ "reason": "access_denied",
+ },
+ )
+
+ exec_app.dependency_overrides[has_xcom_access] = _
+
+ yield
+
+ exec_app.dependency_overrides = {}
+
+
class TestXComsGetEndpoint:
@pytest.mark.parametrize(
("db_value"),
@@ -84,17 +117,18 @@ class TestXComsGetEndpoint:
}
}
- def test_xcom_access_denied(self, client):
- with mock.patch("airflow.api_fastapi.execution_api.routes.xcoms.has_xcom_access", return_value=False):
+ @pytest.mark.usefixtures("access_denied")
+ def test_xcom_access_denied(self, client, caplog):
+ with caplog.at_level(logging.DEBUG):
response = client.get("/execution/xcoms/dag/runid/task/xcom_perms")
- assert response.status_code == 403
+ assert response.status_code == 403, response.json()
assert response.json() == {
"detail": {
"reason": "access_denied",
- "message": "Task does not have access to XCom key 'xcom_perms'",
}
}
+ assert any(msg.startswith("Checking read XCom access") for msg in caplog.messages)
class TestXComsSetEndpoint:
@@ -193,8 +227,9 @@ class TestXComsSetEndpoint:
task_map = session.query(TaskMap).filter_by(task_id=ti.task_id, dag_id=ti.dag_id).one_or_none()
assert task_map.length == length
- def test_xcom_access_denied(self, client):
- with mock.patch("airflow.api_fastapi.execution_api.routes.xcoms.has_xcom_access", return_value=False):
+ @pytest.mark.usefixtures("access_denied")
+ def test_xcom_access_denied(self, client, caplog):
+ with caplog.at_level(logging.DEBUG):
response = client.post(
"/execution/xcoms/dag/runid/task/xcom_perms",
json='"value1"',
@@ -204,9 +239,9 @@ class TestXComsSetEndpoint:
assert response.json() == {
"detail": {
"reason": "access_denied",
- "message": "Task does not have access to set XCom key 'xcom_perms'",
}
}
+ assert any(msg.startswith("Checking write XCom access") for msg in caplog.messages)
@pytest.mark.parametrize(
("value", "expected_value"),