AIP-84 | Add Auth for Dags (apache#47433)

* AIP-84 | Add Auth for Dag

Refactor conftest for api_fastapi and test_dags

Add unauthorized 403 test cases

Remove PATCH in requires_access

Fix unauthorized_test_client, requires_access_dag

Add EditableDagsFilterDep, ReadableDagsFilterDep

Add permitted_dag_filter for dags API

Fix test_security

Add OrmFilterClause

Fix mypy error

* fix(api_fastapi): rename methods argument to method

* Fix kubernetes_tests

* Fix api_fastapi/test_dags

* Add dags_reserialize for k8s tests

Refactor _get_jwt_token

* Increase threshold of test_integration_run_dag_with_scheduler_failure

* test: raise if we cannot get jwt_token not due to connection error

* Fix _get_jwt_token after dynamic patching k8s configMap

* Remove dags_reserialize setup in BaseK8STest

* Fix test_docker_compose_quick_start

* Ensure scheduler health in test_integration_run_dag_with_scheduler_failure

* Increase timeout threshold

* Add HTTP retry for _get_jwt_token

* Add JWTRefreshAdapter and restart api-server if needed

---------

Co-authored-by: Wei Lee <[email protected]>

Author: 2 people

airflow/api_fastapi/common/db/common.py

@@ -35,7 +35,7 @@
 if TYPE_CHECKING:
     from sqlalchemy.sql import Select
 
-    from airflow.api_fastapi.common.parameters import BaseParam
+    from airflow.api_fastapi.core_api.base import OrmClause
 
 
 def _get_session() -> Session:
@@ -47,7 +47,7 @@ def _get_session() -> Session:
 
 
 def apply_filters_to_select(
-    *, statement: Select, filters: Sequence[BaseParam | None] | None = None
+    *, statement: Select, filters: Sequence[OrmClause | None] | None = None
 ) -> Select:
     if filters is None:
         return statement
@@ -71,10 +71,10 @@ async def _get_async_session() -> AsyncSession:
 async def paginated_select_async(
     *,
     statement: Select,
-    filters: Sequence[BaseParam] | None = None,
-    order_by: BaseParam | None = None,
-    offset: BaseParam | None = None,
-    limit: BaseParam | None = None,
+    filters: Sequence[OrmClause] | None = None,
+    order_by: OrmClause | None = None,
+    offset: OrmClause | None = None,
+    limit: OrmClause | None = None,
     session: AsyncSession,
     return_total_entries: Literal[True] = True,
 ) -> tuple[Select, int]: ...
@@ -84,10 +84,10 @@ async def paginated_select_async(
 async def paginated_select_async(
     *,
     statement: Select,
-    filters: Sequence[BaseParam] | None = None,
-    order_by: BaseParam | None = None,
-    offset: BaseParam | None = None,
-    limit: BaseParam | None = None,
+    filters: Sequence[OrmClause] | None = None,
+    order_by: OrmClause | None = None,
+    offset: OrmClause | None = None,
+    limit: OrmClause | None = None,
     session: AsyncSession,
     return_total_entries: Literal[False],
 ) -> tuple[Select, None]: ...
@@ -96,10 +96,10 @@ async def paginated_select_async(
 async def paginated_select_async(
     *,
     statement: Select,
-    filters: Sequence[BaseParam | None] | None = None,
-    order_by: BaseParam | None = None,
-    offset: BaseParam | None = None,
-    limit: BaseParam | None = None,
+    filters: Sequence[OrmClause | None] | None = None,
+    order_by: OrmClause | None = None,
+    offset: OrmClause | None = None,
+    limit: OrmClause | None = None,
     session: AsyncSession,
     return_total_entries: bool = True,
 ) -> tuple[Select, int | None]:
@@ -129,10 +129,10 @@ async def paginated_select_async(
 def paginated_select(
     *,
     statement: Select,
-    filters: Sequence[BaseParam] | None = None,
-    order_by: BaseParam | None = None,
-    offset: BaseParam | None = None,
-    limit: BaseParam | None = None,
+    filters: Sequence[OrmClause] | None = None,
+    order_by: OrmClause | None = None,
+    offset: OrmClause | None = None,
+    limit: OrmClause | None = None,
     session: Session = NEW_SESSION,
     return_total_entries: Literal[True] = True,
 ) -> tuple[Select, int]: ...
@@ -142,10 +142,10 @@ def paginated_select(
 def paginated_select(
     *,
     statement: Select,
-    filters: Sequence[BaseParam] | None = None,
-    order_by: BaseParam | None = None,
-    offset: BaseParam | None = None,
-    limit: BaseParam | None = None,
+    filters: Sequence[OrmClause] | None = None,
+    order_by: OrmClause | None = None,
+    offset: OrmClause | None = None,
+    limit: OrmClause | None = None,
     session: Session = NEW_SESSION,
     return_total_entries: Literal[False],
 ) -> tuple[Select, None]: ...
@@ -155,10 +155,10 @@ def paginated_select(
 def paginated_select(
     *,
     statement: Select,
-    filters: Sequence[BaseParam] | None = None,
-    order_by: BaseParam | None = None,
-    offset: BaseParam | None = None,
-    limit: BaseParam | None = None,
+    filters: Sequence[OrmClause] | None = None,
+    order_by: OrmClause | None = None,
+    offset: OrmClause | None = None,
+    limit: OrmClause | None = None,
     session: Session = NEW_SESSION,
     return_total_entries: bool = True,
 ) -> tuple[Select, int | None]:

airflow/api_fastapi/common/parameters.py

@@ -40,6 +40,7 @@
 from sqlalchemy import Column, and_, case, or_
 from sqlalchemy.inspection import inspect
 
+from airflow.api_fastapi.core_api.base import OrmClause
 from airflow.models import Base
 from airflow.models.asset import (
     AssetAliasModel,
@@ -65,18 +66,14 @@
 T = TypeVar("T")
 
 
-class BaseParam(Generic[T], ABC):
-    """Base class for filters."""
+class BaseParam(OrmClause[T], ABC):
+    """Base class for path or query parameters with ORM transformation."""
 
     def __init__(self, value: T | None = None, skip_none: bool = True) -> None:
-        self.value = value
+        super().__init__(value)
         self.attribute: ColumnElement | None = None
         self.skip_none = skip_none
 
-    @abstractmethod
-    def to_orm(self, select: Select) -> Select:
-        pass
-
     def set_value(self, value: T | None) -> Self:
         self.value = value
         return self

airflow/api_fastapi/core_api/base.py

@@ -16,8 +16,16 @@
 # under the License.
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Generic, TypeVar
+
 from pydantic import BaseModel as PydanticBaseModel, ConfigDict
 
+if TYPE_CHECKING:
+    from sqlalchemy.sql import Select
+
+T = TypeVar("T")
+
 
 class BaseModel(PydanticBaseModel):
     """
@@ -39,3 +47,18 @@ class StrictBaseModel(BaseModel):
     """
 
     model_config = ConfigDict(from_attributes=True, populate_by_name=True, extra="forbid")
+
+
+class OrmClause(Generic[T], ABC):
+    """
+    Base class for filtering clauses with paginated_select.
+
+    The subclasses should implement the `to_orm` method and set the `value` attribute.
+    """
+
+    def __init__(self, value: T | None = None):
+        self.value = value
+
+    @abstractmethod
+    def to_orm(self, select: Select) -> Select:
+        pass

airflow/api_fastapi/core_api/openapi/v1-generated.yaml

@@ -3058,6 +3058,8 @@ paths:
       summary: Get Dags
       description: Get all DAGs.
       operationId: get_dags
+      security:
+      - OAuth2PasswordBearer: []
       parameters:
       - name: limit
         in: query
@@ -3223,6 +3225,8 @@ paths:
       summary: Patch Dags
       description: Patch multiple DAGs.
       operationId: patch_dags
+      security:
+      - OAuth2PasswordBearer: []
       parameters:
       - name: update_mask
         in: query
@@ -3358,6 +3362,8 @@ paths:
       summary: Get Dag
       description: Get basic information about a DAG.
       operationId: get_dag
+      security:
+      - OAuth2PasswordBearer: []
       parameters:
       - name: dag_id
         in: path
@@ -3408,6 +3414,8 @@ paths:
       summary: Patch Dag
       description: Patch the specific DAG.
       operationId: patch_dag
+      security:
+      - OAuth2PasswordBearer: []
       parameters:
       - name: dag_id
         in: path
@@ -3474,6 +3482,8 @@ paths:
       summary: Delete Dag
       description: Delete the specific DAG.
       operationId: delete_dag
+      security:
+      - OAuth2PasswordBearer: []
       parameters:
       - name: dag_id
         in: path
@@ -3524,6 +3534,8 @@ paths:
       summary: Get Dag Details
       description: Get details of DAG.
       operationId: get_dag_details
+      security:
+      - OAuth2PasswordBearer: []
       parameters:
       - name: dag_id
         in: path

airflow/api_fastapi/core_api/routes/public/dags.py

@@ -57,6 +57,11 @@
     DAGResponse,
 )
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
+from airflow.api_fastapi.core_api.security import (
+    EditableDagsFilterDep,
+    ReadableDagsFilterDep,
+    requires_access_dag,
+)
 from airflow.api_fastapi.logging.decorators import action_logging
 from airflow.exceptions import AirflowException, DagNotFound
 from airflow.models import DAG, DagModel
@@ -65,7 +70,7 @@
 dags_router = AirflowRouter(tags=["DAG"], prefix="/dags")
 
 
-@dags_router.get("")
+@dags_router.get("", dependencies=[Depends(requires_access_dag(method="GET"))])
 def get_dags(
     limit: QueryLimit,
     offset: QueryOffset,
@@ -105,6 +110,7 @@ def get_dags(
             ).dynamic_depends()
         ),
     ],
+    readable_dags_filter: ReadableDagsFilterDep,
     session: SessionDep,
 ) -> DAGCollectionResponse:
     """Get all DAGs."""
@@ -132,6 +138,7 @@ def get_dags(
             tags,
             owners,
             last_dag_run_state,
+            readable_dags_filter,
         ],
         order_by=order_by,
         offset=offset,
@@ -156,6 +163,7 @@ def get_dags(
             status.HTTP_422_UNPROCESSABLE_ENTITY,
         ]
     ),
+    dependencies=[Depends(requires_access_dag(method="GET"))],
 )
 def get_dag(dag_id: str, session: SessionDep, request: Request) -> DAGResponse:
     """Get basic information about a DAG."""
@@ -182,6 +190,7 @@ def get_dag(dag_id: str, session: SessionDep, request: Request) -> DAGResponse:
             status.HTTP_404_NOT_FOUND,
         ]
     ),
+    dependencies=[Depends(requires_access_dag(method="GET"))],
 )
 def get_dag_details(dag_id: str, session: SessionDep, request: Request) -> DAGDetailsResponse:
     """Get details of DAG."""
@@ -208,7 +217,7 @@ def get_dag_details(dag_id: str, session: SessionDep, request: Request) -> DAGDe
             status.HTTP_404_NOT_FOUND,
         ]
     ),
-    dependencies=[Depends(action_logging())],
+    dependencies=[Depends(requires_access_dag(method="PUT")), Depends(action_logging())],
 )
 def patch_dag(
     dag_id: str,
@@ -251,7 +260,7 @@ def patch_dag(
             status.HTTP_404_NOT_FOUND,
         ]
     ),
-    dependencies=[Depends(action_logging())],
+    dependencies=[Depends(requires_access_dag(method="PUT")), Depends(action_logging())],
 )
 def patch_dags(
     patch_body: DAGPatchBody,
@@ -263,6 +272,7 @@ def patch_dags(
     only_active: QueryOnlyActiveFilter,
     paused: QueryPausedFilter,
     last_dag_run_state: QueryLastDagRunStateFilter,
+    editable_dags_filter: EditableDagsFilterDep,
     session: SessionDep,
     update_mask: list[str] | None = Query(None),
 ) -> DAGCollectionResponse:
@@ -283,7 +293,7 @@ def patch_dags(
 
     dags_select, total_entries = paginated_select(
         statement=generate_dag_with_latest_run_query(),
-        filters=[only_active, paused, dag_id_pattern, tags, owners, last_dag_run_state],
+        filters=[only_active, paused, dag_id_pattern, tags, owners, last_dag_run_state, editable_dags_filter],
         order_by=None,
         offset=offset,
         limit=limit,
@@ -313,7 +323,7 @@ def patch_dags(
             status.HTTP_422_UNPROCESSABLE_ENTITY,
         ]
     ),
-    dependencies=[Depends(action_logging())],
+    dependencies=[Depends(requires_access_dag(method="DELETE")), Depends(action_logging())],
 )
 def delete_dag(
     dag_id: str,