UI Rest API authentication

airflow/api_ui/app.py

@@ -16,12 +16,18 @@
 # under the License.
 from __future__ import annotations
 
+import logging
+
 from fastapi import APIRouter, FastAPI
 
+from airflow.configuration import conf
+from airflow.www.app import create_app as create_flask_app
 from airflow.www.extensions.init_dagbag import get_dag_bag
 
 app: FastAPI | None = None
 
+log = logging.getLogger(__name__)
+
 
 def init_dag_bag(app: FastAPI) -> None:
     """
@@ -32,17 +38,35 @@ def init_dag_bag(app: FastAPI) -> None:
     app.state.dag_bag = get_dag_bag()
 
 
-def create_app() -> FastAPI:
+def init_flask_app(app: FastAPI, testing: bool = False) -> None:
+    """
+    Auth providers and permission logic are tightly coupled to Flask.
+
+    Leverage the create_app for flask to initialize the api_auth, app_builder, auth_manager.
+    Also limit the backend auth in production mode.
+    """
+    flask_app = create_flask_app(testing=testing)
+
+    app.state.flask_app = flask_app
+
+    if conf.get("api", "auth_backends") != "airflow.providers.fab.auth_manager.api.auth.backend.basic_auth":
+        log.warning(
+            "FastAPI UI API only supports 'airflow.providers.fab.auth_manager.api.auth.backend.basic_auth' other backends will be ignored",
+        )
+
+
+def create_app(testing: bool = False) -> FastAPI:
     app = FastAPI(
         description="Internal Rest API for the UI frontend. It is subject to breaking change "
         "depending on the need of the frontend. Users should not rely on this API but use the "
-        "public API instead."
+        "public API instead.",
     )
 
     init_dag_bag(app)
-
     init_views(app)
 
+    init_flask_app(app, testing)
+
     return app
 
 
@@ -57,15 +81,15 @@ def init_views(app) -> None:
     app.include_router(root_router)
 
 
-def cached_app(config=None, testing=False):
+def cached_app(config=None, testing=False) -> FastAPI:
     """Return cached instance of Airflow UI app."""
     global app
     if not app:
         app = create_app()
     return app
 
 
-def purge_cached_app():
+def purge_cached_app() -> None:
     """Remove the cached version of the app in global state."""
     global app
     app = None

airflow/api_ui/db.py

@@ -0,0 +1,43 @@
+# 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
+
+from typing import TYPE_CHECKING
+
+from airflow.utils.session import create_session
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+
+async def get_session() -> Session:
+    """
+    Dependency for providing a session.
+
+    For non route function please use the use the :class:`airflow.utils.session.provide_session` decorator.
+
+    Example usage:
+
+    .. code:: python
+
+        @router.get("/your_path")
+        def your_route(session: Annotated[Session, Depends(get_session)]):
+            pass
+    """
+    with create_session() as session:
+        yield session

airflow/api_ui/security.py

@@ -0,0 +1,104 @@
+# 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
+
+from typing import Callable, cast
+
+from fastapi import Depends, HTTPException, Request
+from fastapi.security import HTTPBasic, HTTPBasicCredentials
+from typing_extensions import Annotated
+
+from airflow.auth.managers.base_auth_manager import ResourceMethod
+from airflow.auth.managers.models.base_user import BaseUser
+from airflow.auth.managers.models.resource_details import DagAccessEntity, DagDetails, DatasetDetails
+from airflow.providers.fab.auth_manager.api.auth.backend.basic_auth import auth_current_user
+from airflow.providers.fab.auth_manager.models import User
+from airflow.www.extensions.init_auth_manager import get_auth_manager
+
+security = HTTPBasic()
+
+
+def method(request: Request) -> ResourceMethod:
+    return cast(ResourceMethod, request.method)
+
+
+def check_authentication(
+    credentials: Annotated[HTTPBasicCredentials, Depends(security)],
+) -> User | None:
+    """Check that the request has valid authorization information."""
+    # TODO: Handle other auth backends when they support FastAPI. (session, kerberos, etc.)
+    user = auth_current_user(credentials)
+    if user is not None:
+        return user
+
+    # since this handler only checks authentication, not authorization,
+    # we should always return 401
+    raise HTTPException(401, headers={"WWW-Authenticate": "Basic"})
+
+
+def _requires_access(
+    *,
+    is_authorized_callback: Callable[[], bool],
+) -> None:
+    if not is_authorized_callback():
+        raise HTTPException(403, "Forbidden")
+
+
+def requires_access_dataset(
+    method: Annotated[ResourceMethod, Depends(method)],
+    uri: str | None = None,
+    user: Annotated[BaseUser | None, Depends(check_authentication)] = None,
+) -> None:
+    _requires_access(
+        is_authorized_callback=lambda: get_auth_manager().is_authorized_dataset(
+            user=user,
+            method=method,
+            details=DatasetDetails(uri=uri),
+        )
+    )
+
+
+def requires_access_dag(access_entity: DagAccessEntity | None = None) -> Callable:
+    def inner(
+        method: Annotated[ResourceMethod, Depends(method)],
+        dag_id: str | None = None,
+        user: Annotated[BaseUser | None, Depends(check_authentication)] = None,
+    ) -> None:
+        def callback():
+            access = get_auth_manager().is_authorized_dag(
+                method=method, access_entity=access_entity, details=DagDetails(id=dag_id), user=user
+            )
+
+            # ``access`` means here:
+            # - if a DAG id is provided (``dag_id`` not None): is the user authorized to access this DAG
+            # - if no DAG id is provided: is the user authorized to access all DAGs
+            if dag_id or access or access_entity:
+                return access
+
+            # No DAG id is provided, the user is not authorized to access all DAGs and authorization is done
+            # on DAG level
+            # If method is "GET", return whether the user has read access to any DAGs
+            # If method is "PUT", return whether the user has edit access to any DAGs
+            return (method == "GET" and any(get_auth_manager().get_permitted_dag_ids(methods=["GET"]))) or (
+                method == "PUT" and any(get_auth_manager().get_permitted_dag_ids(methods=["PUT"]))
+            )
+
+        _requires_access(
+            is_authorized_callback=callback,
+        )
+
+    return inner

airflow/api_ui/views/datasets.py

@@ -14,69 +14,80 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-
 from __future__ import annotations
 
-from fastapi import APIRouter, HTTPException, Request
+from fastapi import APIRouter, Depends, HTTPException, Request
 from sqlalchemy import and_, func, select
+from sqlalchemy.orm import Session
+from typing_extensions import Annotated
 
+from airflow.api_ui.db import get_session
+from airflow.api_ui.security import (
+    requires_access_dag,
+    requires_access_dataset,
+)
+from airflow.auth.managers.models.resource_details import DagAccessEntity
 from airflow.models import DagModel
 from airflow.models.dataset import DagScheduleDatasetReference, DatasetDagRunQueue, DatasetEvent, DatasetModel
-from airflow.utils.session import create_session
 
 dataset_router = APIRouter(tags=["Dataset"])
 
 
-# Ultimately we want async routes, with async sqlalchemy session / context manager.
-# Additional effort to make airflow utility code async, not handled for now and most likely part of the AIP-70
-@dataset_router.get("/next_run_datasets/{dag_id}")
-async def next_run_datasets(dag_id: str, request: Request) -> dict:
+@dataset_router.get(
+    "/next_run_datasets/{dag_id}",
+    dependencies=[
+        Depends(requires_access_dataset),
+        Depends(requires_access_dag(access_entity=DagAccessEntity.RUN)),
+    ],
+)
+async def next_run_datasets(
+    dag_id: str, request: Request, session: Annotated[Session, Depends(get_session)]
+) -> dict:
     dag = request.app.state.dag_bag.get_dag(dag_id)
 
     if not dag:
         raise HTTPException(404, f"can't find dag {dag_id}")
 
-    with create_session() as session:
-        dag_model = DagModel.get_dagmodel(dag_id, session=session)
+    dag_model = DagModel.get_dagmodel(dag_id, session=session)
 
-        if dag_model is None:
-            raise HTTPException(404, f"can't find associated dag_model {dag_id}")
+    if dag_model is None:
+        raise HTTPException(404, f"can't find associated dag_model {dag_id}")
 
-        latest_run = dag_model.get_last_dagrun(session=session)
+    latest_run = dag_model.get_last_dagrun(session=session)
 
-        events = [
-            dict(info._mapping)
-            for info in session.execute(
-                select(
-                    DatasetModel.id,
-                    DatasetModel.uri,
-                    func.max(DatasetEvent.timestamp).label("lastUpdate"),
-                )
-                .join(DagScheduleDatasetReference, DagScheduleDatasetReference.dataset_id == DatasetModel.id)
-                .join(
-                    DatasetDagRunQueue,
-                    and_(
-                        DatasetDagRunQueue.dataset_id == DatasetModel.id,
-                        DatasetDagRunQueue.target_dag_id == DagScheduleDatasetReference.dag_id,
-                    ),
-                    isouter=True,
-                )
-                .join(
-                    DatasetEvent,
-                    and_(
-                        DatasetEvent.dataset_id == DatasetModel.id,
-                        (
-                            DatasetEvent.timestamp >= latest_run.execution_date
-                            if latest_run and latest_run.execution_date
-                            else True
-                        ),
+    events = [
+        dict(info._mapping)
+        for info in session.execute(
+            select(
+                DatasetModel.id,
+                DatasetModel.uri,
+                func.max(DatasetEvent.timestamp).label("lastUpdate"),
+            )
+            .join(DagScheduleDatasetReference, DagScheduleDatasetReference.dataset_id == DatasetModel.id)
+            .join(
+                DatasetDagRunQueue,
+                and_(
+                    DatasetDagRunQueue.dataset_id == DatasetModel.id,
+                    DatasetDagRunQueue.target_dag_id == DagScheduleDatasetReference.dag_id,
+                ),
+                isouter=True,
+            )
+            .join(
+                DatasetEvent,
+                and_(
+                    DatasetEvent.dataset_id == DatasetModel.id,
+                    (
+                        DatasetEvent.timestamp >= latest_run.execution_date
+                        if latest_run and latest_run.execution_date
+                        else True
                     ),
-                    isouter=True,
-                )
-                .where(DagScheduleDatasetReference.dag_id == dag_id, ~DatasetModel.is_orphaned)
-                .group_by(DatasetModel.id, DatasetModel.uri)
-                .order_by(DatasetModel.uri)
+                ),
+                isouter=True,
             )
-        ]
-        data = {"dataset_expression": dag_model.dataset_expression, "events": events}
-        return data
+            .where(DagScheduleDatasetReference.dag_id == dag_id, ~DatasetModel.is_orphaned)
+            .group_by(DatasetModel.id, DatasetModel.uri)
+            .order_by(DatasetModel.uri)
+        )
+    ]
+    data = {"dataset_expression": dag_model.dataset_expression, "events": events}
+    return data