From d472ee4060732e92a20683e95b11fa20ec4cff7c Mon Sep 17 00:00:00 2001
From: Omkar P <45419097+omkar-foss@users.noreply.github.com>
Date: Tue, 5 Nov 2024 18:08:06 +0530
Subject: [PATCH 1/3] Migrate public endpoint Get DAG Stats to FastAPI, with
 main resynced

---
 .../endpoints/dag_stats_endpoint.py           |   2 +
 airflow/api_fastapi/common/db/common.py       |   5 +-
 airflow/api_fastapi/common/db/dag_runs.py     |  32 ++
 airflow/api_fastapi/common/parameters.py      |  15 +
 .../core_api/routes/public/__init__.py        |   2 +
 .../core_api/routes/public/dag_stats.py       |  79 ++++
 .../core_api/serializers/dag_stats.py         |  43 ++
 .../core_api/routes/public/test_dag_stats.py  | 416 ++++++++++++++++++
 8 files changed, 593 insertions(+), 1 deletion(-)
 create mode 100644 airflow/api_fastapi/common/db/dag_runs.py
 create mode 100644 airflow/api_fastapi/core_api/routes/public/dag_stats.py
 create mode 100644 airflow/api_fastapi/core_api/serializers/dag_stats.py
 create mode 100644 tests/api_fastapi/core_api/routes/public/test_dag_stats.py

diff --git a/airflow/api_connexion/endpoints/dag_stats_endpoint.py b/airflow/api_connexion/endpoints/dag_stats_endpoint.py
index 3b6c6ab8e0df3..c4d8701f8d3c0 100644
--- a/airflow/api_connexion/endpoints/dag_stats_endpoint.py
+++ b/airflow/api_connexion/endpoints/dag_stats_endpoint.py
@@ -27,6 +27,7 @@
 )
 from airflow.auth.managers.models.resource_details import DagAccessEntity
 from airflow.models.dag import DagRun
+from airflow.utils.api_migration import mark_fastapi_migration_done
 from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.utils.state import DagRunState
 from airflow.www.extensions.init_auth_manager import get_auth_manager
@@ -37,6 +38,7 @@
     from airflow.api_connexion.types import APIResponse
 
 
+@mark_fastapi_migration_done
 @security.requires_access_dag("GET", DagAccessEntity.RUN)
 @provide_session
 def get_dag_stats(
diff --git a/airflow/api_fastapi/common/db/common.py b/airflow/api_fastapi/common/db/common.py
index 3feb3ba599504..01e1fe532bf60 100644
--- a/airflow/api_fastapi/common/db/common.py
+++ b/airflow/api_fastapi/common/db/common.py
@@ -65,13 +65,16 @@ def paginated_select(
     offset: BaseParam | None = None,
     limit: BaseParam | None = None,
     session: Session = NEW_SESSION,
+    return_total_entries: bool = True,
 ) -> Select:
     base_select = apply_filters_to_select(
         base_select,
         filters,
     )
 
-    total_entries = get_query_count(base_select, session=session)
+    total_entries = None
+    if return_total_entries:
+        total_entries = get_query_count(base_select, session=session)
 
     # TODO: Re-enable when permissions are handled. Readable / writable entities,
     # for instance:
diff --git a/airflow/api_fastapi/common/db/dag_runs.py b/airflow/api_fastapi/common/db/dag_runs.py
new file mode 100644
index 0000000000000..8f4b02a067fb3
--- /dev/null
+++ b/airflow/api_fastapi/common/db/dag_runs.py
@@ -0,0 +1,32 @@
+# 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 sqlalchemy import func, select
+
+from airflow.models.dagrun import DagRun
+
+dagruns_select_with_state_count = (
+    select(
+        DagRun.dag_id,
+        DagRun.state,
+        func.count(DagRun.state),
+    )
+    .group_by(DagRun.dag_id, DagRun.state)
+    .order_by(DagRun.dag_id)
+)
diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py
index bd65017637227..894cc8bc09ff1 100644
--- a/airflow/api_fastapi/common/parameters.py
+++ b/airflow/api_fastapi/common/parameters.py
@@ -111,6 +111,18 @@ def depends(self, only_active: bool = True) -> _OnlyActiveFilter:
         return self.set_value(only_active)
 
 
+class _DagIdsFilter(BaseParam[list[str]]):
+    """Filter on multi-valued dag_ids param for DagRun."""
+
+    def to_orm(self, select: Select) -> Select:
+        if self.value and self.skip_none:
+            return select.where(DagRun.dag_id.in_(self.value))
+        return select
+
+    def depends(self, dag_ids: list[str] = Query(None)) -> _DagIdsFilter:
+        return self.set_value(dag_ids)
+
+
 class _SearchParam(BaseParam[str]):
     """Search on attribute."""
 
@@ -323,6 +335,7 @@ def depends(self, dag_id: str | None = None) -> _DagIdFilter:
 
 # Common Safe DateTime
 DateTimeQuery = Annotated[str, AfterValidator(_safe_parse_datetime)]
+
 # DAG
 QueryLimit = Annotated[_LimitFilter, Depends(_LimitFilter().depends)]
 QueryOffset = Annotated[_OffsetFilter, Depends(_OffsetFilter().depends)]
@@ -337,8 +350,10 @@ def depends(self, dag_id: str | None = None) -> _DagIdFilter:
 ]
 QueryTagsFilter = Annotated[_TagsFilter, Depends(_TagsFilter().depends)]
 QueryOwnersFilter = Annotated[_OwnersFilter, Depends(_OwnersFilter().depends)]
+
 # DagRun
 QueryLastDagRunStateFilter = Annotated[_LastDagRunStateFilter, Depends(_LastDagRunStateFilter().depends)]
+QueryDagIdsFilter = Annotated[_DagIdsFilter, Depends(_DagIdsFilter().depends)]
 # DAGWarning
 QueryDagIdInDagWarningFilter = Annotated[_DagIdFilter, Depends(_DagIdFilter(DagWarning.dag_id).depends)]
 QueryWarningTypeFilter = Annotated[_WarningTypeFilter, Depends(_WarningTypeFilter().depends)]
diff --git a/airflow/api_fastapi/core_api/routes/public/__init__.py b/airflow/api_fastapi/core_api/routes/public/__init__.py
index a443f5a28ae8c..f8ca6e9d08c48 100644
--- a/airflow/api_fastapi/core_api/routes/public/__init__.py
+++ b/airflow/api_fastapi/core_api/routes/public/__init__.py
@@ -22,6 +22,7 @@
 from airflow.api_fastapi.core_api.routes.public.connections import connections_router
 from airflow.api_fastapi.core_api.routes.public.dag_run import dag_run_router
 from airflow.api_fastapi.core_api.routes.public.dag_sources import dag_sources_router
+from airflow.api_fastapi.core_api.routes.public.dag_stats import dag_stats_router
 from airflow.api_fastapi.core_api.routes.public.dag_warning import dag_warning_router
 from airflow.api_fastapi.core_api.routes.public.dags import dags_router
 from airflow.api_fastapi.core_api.routes.public.event_logs import event_logs_router
@@ -52,3 +53,4 @@
 public_router.include_router(variables_router)
 public_router.include_router(variables_router)
 public_router.include_router(version_router)
+public_router.include_router(dag_stats_router)
diff --git a/airflow/api_fastapi/core_api/routes/public/dag_stats.py b/airflow/api_fastapi/core_api/routes/public/dag_stats.py
new file mode 100644
index 0000000000000..deed40f0116d5
--- /dev/null
+++ b/airflow/api_fastapi/core_api/routes/public/dag_stats.py
@@ -0,0 +1,79 @@
+# 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 fastapi import Depends
+from sqlalchemy.orm import Session
+from typing_extensions import Annotated
+
+from airflow.api_fastapi.common.db.common import (
+    get_session,
+    paginated_select,
+)
+from airflow.api_fastapi.common.db.dag_runs import dagruns_select_with_state_count
+from airflow.api_fastapi.common.parameters import QueryDagIdsFilter
+from airflow.api_fastapi.common.router import AirflowRouter
+from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
+from airflow.api_fastapi.core_api.serializers.dag_stats import (
+    DagStatsCollectionResponse,
+    DagStatsResponse,
+    DagStatsStateResponse,
+)
+from airflow.utils.state import DagRunState
+
+dag_stats_router = AirflowRouter(tags=["DagStats"], prefix="/dagStats")
+
+
+@dag_stats_router.get(
+    "/",
+    responses=create_openapi_http_exception_doc([400, 401, 403, 404]),
+)
+async def get_dag_stats(
+    session: Annotated[Session, Depends(get_session)],
+    dag_ids: QueryDagIdsFilter,
+) -> DagStatsCollectionResponse:
+    """Get Dag statistics."""
+    dagruns_select, _ = paginated_select(
+        base_select=dagruns_select_with_state_count,
+        filters=[dag_ids],
+        session=session,
+        return_total_entries=False,
+    )
+    query_result = session.execute(dagruns_select)
+
+    result_dag_ids = []
+    dag_state_data = {}
+    for dag_id, state, count in query_result:
+        dag_state_data[(dag_id, state)] = count
+        if dag_id not in result_dag_ids:
+            result_dag_ids.append(dag_id)
+
+    dags = [
+        DagStatsResponse(
+            dag_id=dag_id,
+            stats=[
+                DagStatsStateResponse(
+                    state=state,
+                    count=dag_state_data.get((dag_id, state), 0),
+                )
+                for state in DagRunState
+            ],
+        )
+        for dag_id in result_dag_ids
+    ]
+    return DagStatsCollectionResponse(dags=dags, total_entries=len(dags))
diff --git a/airflow/api_fastapi/core_api/serializers/dag_stats.py b/airflow/api_fastapi/core_api/serializers/dag_stats.py
new file mode 100644
index 0000000000000..0d768c2cbac07
--- /dev/null
+++ b/airflow/api_fastapi/core_api/serializers/dag_stats.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 pydantic import BaseModel
+
+from airflow.utils.state import DagRunState
+
+
+class DagStatsStateResponse(BaseModel):
+    """DagStatsState serializer for responses."""
+
+    state: DagRunState
+    count: int
+
+
+class DagStatsResponse(BaseModel):
+    """DAG Stats serializer for responses."""
+
+    dag_id: str
+    stats: list[DagStatsStateResponse]
+
+
+class DagStatsCollectionResponse(BaseModel):
+    """DAG Stats Collection serializer for responses."""
+
+    dags: list[DagStatsResponse]
+    total_entries: int
diff --git a/tests/api_fastapi/core_api/routes/public/test_dag_stats.py b/tests/api_fastapi/core_api/routes/public/test_dag_stats.py
new file mode 100644
index 0000000000000..e2611addd70c0
--- /dev/null
+++ b/tests/api_fastapi/core_api/routes/public/test_dag_stats.py
@@ -0,0 +1,416 @@
+# 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 datetime import datetime, timedelta
+
+import pytest
+
+from airflow.models.dag import DagModel
+from airflow.models.dagrun import DagRun
+from airflow.utils import timezone
+from airflow.utils.state import DagRunState
+from airflow.utils.types import DagRunType
+
+from tests_common.test_utils.db import clear_db_dags, clear_db_runs, clear_db_serialized_dags
+
+pytestmark = pytest.mark.db_test
+
+DAG1_ID = "test_dag1"
+DAG2_ID = "test_dag2"
+DAG3_ID = "test_dag3"
+TASK_ID = "op1"
+API_PREFIX = "/public/dagStats"
+
+
+class TestDagStatsEndpoint:
+    default_time = "2020-06-11T18:00:00+00:00"
+
+    @staticmethod
+    def _clear_db():
+        clear_db_runs()
+        clear_db_dags()
+        clear_db_serialized_dags()
+
+    def _create_dag_and_runs(self, session=None):
+        dag_1 = DagModel(
+            dag_id=DAG1_ID,
+            fileloc="/tmp/dag_stats_1.py",
+            timetable_summary="2 2 * * *",
+            is_active=False,
+            is_paused=True,
+            owners="test_owner,another_test_owner",
+            next_dagrun=datetime(2021, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+        )
+        dag_1_run_1 = DagRun(
+            dag_id=DAG1_ID,
+            run_id="test_dag_run_id_1",
+            run_type=DagRunType.MANUAL,
+            execution_date=timezone.parse(self.default_time),
+            start_date=timezone.parse(self.default_time),
+            external_trigger=True,
+            state="running",
+        )
+        dag_1_run_2 = DagRun(
+            dag_id=dag_1.dag_id,
+            run_id="test_dag_run_id_2",
+            run_type=DagRunType.MANUAL,
+            execution_date=timezone.parse(self.default_time) + timedelta(days=1),
+            start_date=timezone.parse(self.default_time),
+            external_trigger=True,
+            state="failed",
+        )
+        dag_2 = DagModel(
+            dag_id=DAG2_ID,
+            fileloc="/tmp/dag_stats_2.py",
+            timetable_summary="2 2 * * *",
+            is_active=False,
+            is_paused=True,
+            owners="test_owner,another_test_owner",
+            next_dagrun=datetime(2021, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+        )
+        dag_2_run_1 = DagRun(
+            dag_id=dag_2.dag_id,
+            run_id="test_dag_2_run_id_1",
+            run_type=DagRunType.MANUAL,
+            execution_date=timezone.parse(self.default_time),
+            start_date=timezone.parse(self.default_time),
+            external_trigger=True,
+            state="queued",
+        )
+        dag_3 = DagModel(
+            dag_id=DAG3_ID,
+            fileloc="/tmp/dag_stats_3.py",
+            timetable_summary="2 2 * * *",
+            is_active=False,
+            is_paused=True,
+            owners="test_owner,another_test_owner",
+            next_dagrun=datetime(2021, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+        )
+        dag_3_run_1 = DagRun(
+            dag_id=dag_3.dag_id,
+            run_id="test_dag_3_run_id_1",
+            run_type=DagRunType.MANUAL,
+            execution_date=timezone.parse(self.default_time),
+            start_date=timezone.parse(self.default_time),
+            external_trigger=True,
+            state="success",
+        )
+        entities = (
+            dag_1,
+            dag_1_run_1,
+            dag_1_run_2,
+            dag_2,
+            dag_2_run_1,
+            dag_3,
+            dag_3_run_1,
+        )
+        session.add_all(entities)
+        session.commit()
+
+    @pytest.fixture(autouse=True)
+    def setup(self) -> None:
+        self._clear_db()
+
+    def teardown_method(self) -> None:
+        self._clear_db()
+
+
+class TestGetDagStats(TestDagStatsEndpoint):
+    """Unit tests for Get DAG Stats."""
+
+    def test_should_respond_200(self, client, session):
+        self._create_dag_and_runs(session)
+        exp_payload = {
+            "dags": [
+                {
+                    "dag_id": DAG1_ID,
+                    "stats": [
+                        {
+                            "state": DagRunState.QUEUED,
+                            "count": 0,
+                        },
+                        {
+                            "state": DagRunState.RUNNING,
+                            "count": 1,
+                        },
+                        {
+                            "state": DagRunState.SUCCESS,
+                            "count": 0,
+                        },
+                        {
+                            "state": DagRunState.FAILED,
+                            "count": 1,
+                        },
+                    ],
+                },
+                {
+                    "dag_id": DAG2_ID,
+                    "stats": [
+                        {
+                            "state": DagRunState.QUEUED,
+                            "count": 1,
+                        },
+                        {
+                            "state": DagRunState.RUNNING,
+                            "count": 0,
+                        },
+                        {
+                            "state": DagRunState.SUCCESS,
+                            "count": 0,
+                        },
+                        {
+                            "state": DagRunState.FAILED,
+                            "count": 0,
+                        },
+                    ],
+                },
+            ],
+            "total_entries": 2,
+        }
+
+        response = client().get(f"{API_PREFIX}?dag_ids={DAG1_ID}&dag_ids={DAG2_ID}")
+        assert response.status_code == 200
+        res_json = response.json()
+        assert res_json["total_entries"] == len(res_json["dags"])
+        assert res_json == exp_payload
+
+    def test_all_dags_should_respond_200(self, client, session):
+        self._create_dag_and_runs(session)
+        exp_payload = {
+            "dags": [
+                {
+                    "dag_id": DAG1_ID,
+                    "stats": [
+                        {
+                            "state": DagRunState.QUEUED,
+                            "count": 0,
+                        },
+                        {
+                            "state": DagRunState.RUNNING,
+                            "count": 1,
+                        },
+                        {
+                            "state": DagRunState.SUCCESS,
+                            "count": 0,
+                        },
+                        {
+                            "state": DagRunState.FAILED,
+                            "count": 1,
+                        },
+                    ],
+                },
+                {
+                    "dag_id": DAG2_ID,
+                    "stats": [
+                        {
+                            "state": DagRunState.QUEUED,
+                            "count": 1,
+                        },
+                        {
+                            "state": DagRunState.RUNNING,
+                            "count": 0,
+                        },
+                        {
+                            "state": DagRunState.SUCCESS,
+                            "count": 0,
+                        },
+                        {
+                            "state": DagRunState.FAILED,
+                            "count": 0,
+                        },
+                    ],
+                },
+                {
+                    "dag_id": DAG3_ID,
+                    "stats": [
+                        {
+                            "state": DagRunState.QUEUED,
+                            "count": 0,
+                        },
+                        {
+                            "state": DagRunState.RUNNING,
+                            "count": 0,
+                        },
+                        {
+                            "state": DagRunState.SUCCESS,
+                            "count": 1,
+                        },
+                        {
+                            "state": DagRunState.FAILED,
+                            "count": 0,
+                        },
+                    ],
+                },
+            ],
+            "total_entries": 3,
+        }
+
+        response = client().get(API_PREFIX)
+        assert response.status_code == 200
+        res_json = response.json()
+        assert res_json["total_entries"] == len(res_json["dags"])
+        assert res_json == exp_payload
+
+    @pytest.mark.parametrize(
+        "url, params, exp_payload",
+        [
+            (
+                API_PREFIX,
+                [
+                    ("dag_ids", DAG1_ID),
+                    ("dag_ids", DAG3_ID),
+                    ("dag_ids", DAG2_ID),
+                ],
+                {
+                    "dags": [
+                        {
+                            "dag_id": DAG1_ID,
+                            "stats": [
+                                {
+                                    "state": DagRunState.QUEUED,
+                                    "count": 0,
+                                },
+                                {
+                                    "state": DagRunState.RUNNING,
+                                    "count": 1,
+                                },
+                                {
+                                    "state": DagRunState.SUCCESS,
+                                    "count": 0,
+                                },
+                                {
+                                    "state": DagRunState.FAILED,
+                                    "count": 1,
+                                },
+                            ],
+                        },
+                        {
+                            "dag_id": DAG2_ID,
+                            "stats": [
+                                {
+                                    "state": DagRunState.QUEUED,
+                                    "count": 1,
+                                },
+                                {
+                                    "state": DagRunState.RUNNING,
+                                    "count": 0,
+                                },
+                                {
+                                    "state": DagRunState.SUCCESS,
+                                    "count": 0,
+                                },
+                                {
+                                    "state": DagRunState.FAILED,
+                                    "count": 0,
+                                },
+                            ],
+                        },
+                        {
+                            "dag_id": DAG3_ID,
+                            "stats": [
+                                {
+                                    "state": DagRunState.QUEUED,
+                                    "count": 0,
+                                },
+                                {
+                                    "state": DagRunState.RUNNING,
+                                    "count": 0,
+                                },
+                                {
+                                    "state": DagRunState.SUCCESS,
+                                    "count": 1,
+                                },
+                                {
+                                    "state": DagRunState.FAILED,
+                                    "count": 0,
+                                },
+                            ],
+                        },
+                    ],
+                    "total_entries": 3,
+                },
+            ),
+            (
+                API_PREFIX,
+                [("dag_ids", DAG1_ID)],
+                {
+                    "dags": [
+                        {
+                            "dag_id": DAG1_ID,
+                            "stats": [
+                                {
+                                    "state": DagRunState.QUEUED,
+                                    "count": 0,
+                                },
+                                {
+                                    "state": DagRunState.RUNNING,
+                                    "count": 1,
+                                },
+                                {
+                                    "state": DagRunState.SUCCESS,
+                                    "count": 0,
+                                },
+                                {
+                                    "state": DagRunState.FAILED,
+                                    "count": 1,
+                                },
+                            ],
+                        }
+                    ],
+                    "total_entries": 1,
+                },
+            ),
+            (
+                API_PREFIX,
+                [("dag_ids", DAG3_ID)],
+                {
+                    "dags": [
+                        {
+                            "dag_id": DAG3_ID,
+                            "stats": [
+                                {
+                                    "state": DagRunState.QUEUED,
+                                    "count": 0,
+                                },
+                                {
+                                    "state": DagRunState.RUNNING,
+                                    "count": 0,
+                                },
+                                {
+                                    "state": DagRunState.SUCCESS,
+                                    "count": 1,
+                                },
+                                {
+                                    "state": DagRunState.FAILED,
+                                    "count": 0,
+                                },
+                            ],
+                        },
+                    ],
+                    "total_entries": 1,
+                },
+            ),
+        ],
+    )
+    def test_single_dag_in_dag_ids(self, client, session, url, params, exp_payload):
+        self._create_dag_and_runs(session)
+        response = client().get(url, params=params)
+        assert response.status_code == 200
+        res_json = response.json()
+        assert res_json["total_entries"] == len(res_json["dags"])
+        assert res_json == exp_payload

From 47dd7d51c85bfae7d656da45e6770912ea179911 Mon Sep 17 00:00:00 2001
From: Omkar P <45419097+omkar-foss@users.noreply.github.com>
Date: Tue, 5 Nov 2024 18:08:19 +0530
Subject: [PATCH 2/3] Re-run static checks

---
 .../core_api/openapi/v1-generated.yaml        | 98 +++++++++++++++++++
 airflow/ui/openapi-gen/queries/common.ts      | 17 ++++
 airflow/ui/openapi-gen/queries/prefetch.ts    | 21 ++++
 airflow/ui/openapi-gen/queries/queries.ts     | 27 +++++
 airflow/ui/openapi-gen/queries/suspense.ts    | 27 +++++
 .../ui/openapi-gen/requests/schemas.gen.ts    | 56 +++++++++++
 .../ui/openapi-gen/requests/services.gen.ts   | 31 ++++++
 airflow/ui/openapi-gen/requests/types.gen.ts  | 61 ++++++++++++
 8 files changed, 338 insertions(+)

diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml
index e844cbceeb474..d2531ec296bf5 100644
--- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml
+++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml
@@ -2320,6 +2320,59 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/VersionInfo'
+  /public/dagStats/:
+    get:
+      tags:
+      - DagStats
+      summary: Get Dag Stats
+      description: Get Dag statistics.
+      operationId: get_dag_stats
+      parameters:
+      - name: dag_ids
+        in: query
+        required: false
+        schema:
+          type: array
+          items:
+            type: string
+          title: Dag Ids
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/DagStatsCollectionResponse'
+        '400':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Bad Request
+        '401':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Unauthorized
+        '403':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Forbidden
+        '404':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Not Found
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
 components:
   schemas:
     AppBuilderMenuItemResponse:
@@ -3270,6 +3323,51 @@ components:
       - asset_triggered
       title: DagRunType
       description: Class with DagRun types.
+    DagStatsCollectionResponse:
+      properties:
+        dags:
+          items:
+            $ref: '#/components/schemas/DagStatsResponse'
+          type: array
+          title: Dags
+        total_entries:
+          type: integer
+          title: Total Entries
+      type: object
+      required:
+      - dags
+      - total_entries
+      title: DagStatsCollectionResponse
+      description: DAG Stats Collection serializer for responses.
+    DagStatsResponse:
+      properties:
+        dag_id:
+          type: string
+          title: Dag Id
+        stats:
+          items:
+            $ref: '#/components/schemas/DagStatsStateResponse'
+          type: array
+          title: Stats
+      type: object
+      required:
+      - dag_id
+      - stats
+      title: DagStatsResponse
+      description: DAG Stats serializer for responses.
+    DagStatsStateResponse:
+      properties:
+        state:
+          $ref: '#/components/schemas/DagRunState'
+        count:
+          type: integer
+          title: Count
+      type: object
+      required:
+      - state
+      - count
+      title: DagStatsStateResponse
+      description: DagStatsState serializer for responses.
     DagTagPydantic:
       properties:
         name:
diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts
index 1248a77ce188d..a0bf4d42ca3e7 100644
--- a/airflow/ui/openapi-gen/queries/common.ts
+++ b/airflow/ui/openapi-gen/queries/common.ts
@@ -8,6 +8,7 @@ import {
   DagRunService,
   DagService,
   DagSourceService,
+  DagStatsService,
   DagWarningService,
   DagsService,
   DashboardService,
@@ -634,6 +635,22 @@ export const UseVersionServiceGetVersionKeyFn = (queryKey?: Array<unknown>) => [
   useVersionServiceGetVersionKey,
   ...(queryKey ?? []),
 ];
+export type DagStatsServiceGetDagStatsDefaultResponse = Awaited<
+  ReturnType<typeof DagStatsService.getDagStats>
+>;
+export type DagStatsServiceGetDagStatsQueryResult<
+  TData = DagStatsServiceGetDagStatsDefaultResponse,
+  TError = unknown,
+> = UseQueryResult<TData, TError>;
+export const useDagStatsServiceGetDagStatsKey = "DagStatsServiceGetDagStats";
+export const UseDagStatsServiceGetDagStatsKeyFn = (
+  {
+    dagIds,
+  }: {
+    dagIds?: string[];
+  } = {},
+  queryKey?: Array<unknown>,
+) => [useDagStatsServiceGetDagStatsKey, ...(queryKey ?? [{ dagIds }])];
 export type BackfillServiceCreateBackfillMutationResult = Awaited<
   ReturnType<typeof BackfillService.createBackfill>
 >;
diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts
index bf6ad800be012..6df529656acad 100644
--- a/airflow/ui/openapi-gen/queries/prefetch.ts
+++ b/airflow/ui/openapi-gen/queries/prefetch.ts
@@ -8,6 +8,7 @@ import {
   DagRunService,
   DagService,
   DagSourceService,
+  DagStatsService,
   DagWarningService,
   DagsService,
   DashboardService,
@@ -821,3 +822,23 @@ export const prefetchUseVersionServiceGetVersion = (queryClient: QueryClient) =>
     queryKey: Common.UseVersionServiceGetVersionKeyFn(),
     queryFn: () => VersionService.getVersion(),
   });
+/**
+ * Get Dag Stats
+ * Get Dag statistics.
+ * @param data The data for the request.
+ * @param data.dagIds
+ * @returns DagStatsCollectionResponse Successful Response
+ * @throws ApiError
+ */
+export const prefetchUseDagStatsServiceGetDagStats = (
+  queryClient: QueryClient,
+  {
+    dagIds,
+  }: {
+    dagIds?: string[];
+  } = {},
+) =>
+  queryClient.prefetchQuery({
+    queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagIds }),
+    queryFn: () => DagStatsService.getDagStats({ dagIds }),
+  });
diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts
index 70796be401719..86215f2bf8a69 100644
--- a/airflow/ui/openapi-gen/queries/queries.ts
+++ b/airflow/ui/openapi-gen/queries/queries.ts
@@ -13,6 +13,7 @@ import {
   DagRunService,
   DagService,
   DagSourceService,
+  DagStatsService,
   DagWarningService,
   DagsService,
   DashboardService,
@@ -1025,6 +1026,32 @@ export const useVersionServiceGetVersion = <
     queryFn: () => VersionService.getVersion() as TData,
     ...options,
   });
+/**
+ * Get Dag Stats
+ * Get Dag statistics.
+ * @param data The data for the request.
+ * @param data.dagIds
+ * @returns DagStatsCollectionResponse Successful Response
+ * @throws ApiError
+ */
+export const useDagStatsServiceGetDagStats = <
+  TData = Common.DagStatsServiceGetDagStatsDefaultResponse,
+  TError = unknown,
+  TQueryKey extends Array<unknown> = unknown[],
+>(
+  {
+    dagIds,
+  }: {
+    dagIds?: string[];
+  } = {},
+  queryKey?: TQueryKey,
+  options?: Omit<UseQueryOptions<TData, TError>, "queryKey" | "queryFn">,
+) =>
+  useQuery<TData, TError>({
+    queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagIds }, queryKey),
+    queryFn: () => DagStatsService.getDagStats({ dagIds }) as TData,
+    ...options,
+  });
 /**
  * Create Backfill
  * @param data The data for the request.
diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts
index 4f75c2ba0ce47..128110ccce451 100644
--- a/airflow/ui/openapi-gen/queries/suspense.ts
+++ b/airflow/ui/openapi-gen/queries/suspense.ts
@@ -8,6 +8,7 @@ import {
   DagRunService,
   DagService,
   DagSourceService,
+  DagStatsService,
   DagWarningService,
   DagsService,
   DashboardService,
@@ -1011,3 +1012,29 @@ export const useVersionServiceGetVersionSuspense = <
     queryFn: () => VersionService.getVersion() as TData,
     ...options,
   });
+/**
+ * Get Dag Stats
+ * Get Dag statistics.
+ * @param data The data for the request.
+ * @param data.dagIds
+ * @returns DagStatsCollectionResponse Successful Response
+ * @throws ApiError
+ */
+export const useDagStatsServiceGetDagStatsSuspense = <
+  TData = Common.DagStatsServiceGetDagStatsDefaultResponse,
+  TError = unknown,
+  TQueryKey extends Array<unknown> = unknown[],
+>(
+  {
+    dagIds,
+  }: {
+    dagIds?: string[];
+  } = {},
+  queryKey?: TQueryKey,
+  options?: Omit<UseQueryOptions<TData, TError>, "queryKey" | "queryFn">,
+) =>
+  useSuspenseQuery<TData, TError>({
+    queryKey: Common.UseDagStatsServiceGetDagStatsKeyFn({ dagIds }, queryKey),
+    queryFn: () => DagStatsService.getDagStats({ dagIds }) as TData,
+    ...options,
+  });
diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts
index 53272eae2e779..f483a8b93a9f9 100644
--- a/airflow/ui/openapi-gen/requests/schemas.gen.ts
+++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts
@@ -1474,6 +1474,62 @@ export const $DagRunType = {
   description: "Class with DagRun types.",
 } as const;
 
+export const $DagStatsCollectionResponse = {
+  properties: {
+    dags: {
+      items: {
+        $ref: "#/components/schemas/DagStatsResponse",
+      },
+      type: "array",
+      title: "Dags",
+    },
+    total_entries: {
+      type: "integer",
+      title: "Total Entries",
+    },
+  },
+  type: "object",
+  required: ["dags", "total_entries"],
+  title: "DagStatsCollectionResponse",
+  description: "DAG Stats Collection serializer for responses.",
+} as const;
+
+export const $DagStatsResponse = {
+  properties: {
+    dag_id: {
+      type: "string",
+      title: "Dag Id",
+    },
+    stats: {
+      items: {
+        $ref: "#/components/schemas/DagStatsStateResponse",
+      },
+      type: "array",
+      title: "Stats",
+    },
+  },
+  type: "object",
+  required: ["dag_id", "stats"],
+  title: "DagStatsResponse",
+  description: "DAG Stats serializer for responses.",
+} as const;
+
+export const $DagStatsStateResponse = {
+  properties: {
+    state: {
+      $ref: "#/components/schemas/DagRunState",
+    },
+    count: {
+      type: "integer",
+      title: "Count",
+    },
+  },
+  type: "object",
+  required: ["state", "count"],
+  title: "DagStatsStateResponse",
+  description: "DagStatsState serializer for responses.",
+} as const;
+
 export const $DagTagPydantic = {
   properties: {
     name: {
diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts
index fa5c7739c909b..21f3fbeae6ad7 100644
--- a/airflow/ui/openapi-gen/requests/services.gen.ts
+++ b/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -85,6 +85,8 @@ import type {
   PostVariableData,
   PostVariableResponse,
   GetVersionResponse,
+  GetDagStatsData,
+  GetDagStatsResponse,
 } from "./types.gen";
 
 export class AssetService {
@@ -1327,3 +1329,32 @@ export class VersionService {
     });
   }
 }
+
+export class DagStatsService {
+  /**
+   * Get Dag Stats
+   * Get Dag statistics.
+   * @param data The data for the request.
+   * @param data.dagIds
+   * @returns DagStatsCollectionResponse Successful Response
+   * @throws ApiError
+   */
+  public static getDagStats(
+    data: GetDagStatsData = {},
+  ): CancelablePromise<GetDagStatsResponse> {
+    return __request(OpenAPI, {
+      method: "GET",
+      url: "/public/dagStats/",
+      query: {
+        dag_ids: data.dagIds,
+      },
+      errors: {
+        400: "Bad Request",
+        401: "Unauthorized",
+        403: "Forbidden",
+        404: "Not Found",
+        422: "Validation Error",
+      },
+    });
+  }
+}
diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts
index 909b78dd627e1..ec136658838d1 100644
--- a/airflow/ui/openapi-gen/requests/types.gen.ts
+++ b/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -331,6 +331,30 @@ export type DagRunType =
   | "manual"
   | "asset_triggered";
 
+/**
+ * DAG Stats Collection serializer for responses.
+ */
+export type DagStatsCollectionResponse = {
+  dags: Array<DagStatsResponse>;
+  total_entries: number;
+};
+
+/**
+ * DAG Stats serializer for responses.
+ */
+export type DagStatsResponse = {
+  dag_id: string;
+  stats: Array<DagStatsStateResponse>;
+};
+
+/**
+ * DagStatsState serializer for responses.
+ */
+export type DagStatsStateResponse = {
+  state: DagRunState;
+  count: number;
+};
+
 /**
  * Serializable representation of the DagTag ORM SqlAlchemyModel used by internal API.
  */
@@ -991,6 +1015,12 @@ export type PostVariableResponse = VariableResponse;
 
 export type GetVersionResponse = VersionInfo;
 
+export type GetDagStatsData = {
+  dagIds?: Array<string>;
+};
+
+export type GetDagStatsResponse = DagStatsCollectionResponse;
+
 export type $OpenApiTs = {
   "/ui/next_run_assets/{dag_id}": {
     get: {
@@ -2036,4 +2066,35 @@ export type $OpenApiTs = {
       };
     };
   };
+  "/public/dagStats/": {
+    get: {
+      req: GetDagStatsData;
+      res: {
+        /**
+         * Successful Response
+         */
+        200: DagStatsCollectionResponse;
+        /**
+         * Bad Request
+         */
+        400: HTTPExceptionResponse;
+        /**
+         * Unauthorized
+         */
+        401: HTTPExceptionResponse;
+        /**
+         * Forbidden
+         */
+        403: HTTPExceptionResponse;
+        /**
+         * Not Found
+         */
+        404: HTTPExceptionResponse;
+        /**
+         * Validation Error
+         */
+        422: HTTPValidationError;
+      };
+    };
+  };
 };

From a812695c38123e1aa60e0330306131599337f6ed Mon Sep 17 00:00:00 2001
From: Omkar P <45419097+omkar-foss@users.noreply.github.com>
Date: Tue, 5 Nov 2024 18:11:20 +0530
Subject: [PATCH 3/3] Add newlines to separate entities

---
 airflow/api_fastapi/common/parameters.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py
index 894cc8bc09ff1..59f485cd88a01 100644
--- a/airflow/api_fastapi/common/parameters.py
+++ b/airflow/api_fastapi/common/parameters.py
@@ -354,8 +354,10 @@ def depends(self, dag_id: str | None = None) -> _DagIdFilter:
 # DagRun
 QueryLastDagRunStateFilter = Annotated[_LastDagRunStateFilter, Depends(_LastDagRunStateFilter().depends)]
 QueryDagIdsFilter = Annotated[_DagIdsFilter, Depends(_DagIdsFilter().depends)]
+
 # DAGWarning
 QueryDagIdInDagWarningFilter = Annotated[_DagIdFilter, Depends(_DagIdFilter(DagWarning.dag_id).depends)]
 QueryWarningTypeFilter = Annotated[_WarningTypeFilter, Depends(_WarningTypeFilter().depends)]
+
 # DAGTags
 QueryDagTagPatternSearch = Annotated[_DagTagNamePatternSearch, Depends(_DagTagNamePatternSearch().depends)]
