Add GetDagState endpoint to execution_api

Author: gopidesupavan

airflow-core/src/airflow/api_fastapi/execution_api/datamodels/dags.py

@@ -0,0 +1,26 @@
+# 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 airflow.api_fastapi.core_api.base import BaseModel
+
+
+class DagStateResponse(BaseModel):
+    """Schema for DAG State response."""
+
+    is_paused: bool

airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py

@@ -25,6 +25,7 @@
     assets,
     connections,
     dag_runs,
+    dags,
     health,
     hitl,
     task_instances,
@@ -43,6 +44,7 @@
 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(dag_runs.router, prefix="/dag-runs", tags=["Dag Runs"])
+authenticated_router.include_router(dags.router, prefix="/dags", tags=["Dags"])
 authenticated_router.include_router(task_instances.router, prefix="/task-instances", tags=["Task Instances"])
 authenticated_router.include_router(
     task_reschedules.router, prefix="/task-reschedules", tags=["Task Reschedules"]

airflow-core/src/airflow/api_fastapi/execution_api/routes/dags.py

@@ -0,0 +1,57 @@
+# 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
+
+import logging
+
+from fastapi import APIRouter, HTTPException, status
+
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.datamodels.dags import DagStateResponse
+from airflow.models.dag import DagModel
+
+router = APIRouter()
+
+
+log = logging.getLogger(__name__)
+
+
+@router.get(
+    "/{dag_id}/state",
+    responses={
+        status.HTTP_404_NOT_FOUND: {"description": "DAG not found for the given dag_id"},
+    },
+)
+def get_dag_state(
+    dag_id: str,
+    session: SessionDep,
+) -> DagStateResponse:
+    """Get a DAG Run State."""
+    dag_model: DagModel = session.get(DagModel, dag_id)
+    if not dag_model:
+        raise HTTPException(
+            status.HTTP_404_NOT_FOUND,
+            detail={
+                "reason": "not_found",
+                "message": f"The Dag with dag_id: `{dag_id}` was not found",
+            },
+        )
+
+    is_paused = False if dag_model.is_paused is None else dag_model.is_paused
+
+    return DagStateResponse(is_paused=is_paused)

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dags.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
+
+import pytest
+
+from airflow.models import DagModel
+from airflow.providers.standard.operators.empty import EmptyOperator
+
+from tests_common.test_utils.db import clear_db_runs
+
+pytestmark = pytest.mark.db_test
+
+
+class TestDagState:
+    def setup_method(self):
+        clear_db_runs()
+
+    def teardown_method(self):
+        clear_db_runs()
+
+    @pytest.mark.parametrize(
+        "state, expected",
+        [
+            (True, True),
+            (False, False),
+            (None, False),
+        ],
+    )
+    def test_dag_is_paused(self, state, expected, client, session, dag_maker):
+        """Test DagState is active or paused"""
+
+        dag_id = "test_dag_is_paused"
+
+        with dag_maker(dag_id=dag_id, session=session, serialized=True):
+            EmptyOperator(task_id="test_task")
+
+        session.query(DagModel).filter(DagModel.dag_id == dag_id).update({"is_paused": state})
+
+        session.commit()
+
+        response = client.get(
+            f"/execution/dags/{dag_id}/state",
+        )
+
+        assert response.status_code == 200
+        assert response.json() == {"is_paused": expected}
+
+    def test_dag_not_found(self, client, session, dag_maker):
+        """Test Dag not found"""
+
+        dag_id = "test_dag_is_paused"
+
+        response = client.get(
+            f"/execution/dags/{dag_id}/state",
+        )
+
+        assert response.status_code == 404
+        assert response.json() == {
+            "detail": {
+                "message": "The Dag with dag_id: `test_dag_is_paused` was not found",
+                "reason": "not_found",
+            }
+        }

task-sdk/src/airflow/sdk/api/client.py

@@ -43,6 +43,7 @@
     ConnectionResponse,
     DagRunStateResponse,
     DagRunType,
+    DagStateResponse,
     HITLDetailResponse,
     HITLUser,
     InactiveAssetsResponse,
@@ -721,6 +722,18 @@ def get_previous(
         return PreviousDagRunResult(dag_run=resp.json())
 
 
+class DagsOperations:
+    __slots__ = ("client",)
+
+    def __init__(self, client: Client):
+        self.client = client
+
+    def get_state(self, dag_id: str) -> DagStateResponse:
+        """Get the state of a Dag via the API server."""
+        resp = self.client.get(f"dags/{dag_id}/state")
+        return DagStateResponse.model_validate_json(resp.read())
+
+
 class HITLOperations:
     """
     Operations related to Human in the loop. Require Airflow 3.1+.
@@ -931,6 +944,11 @@ def hitl(self):
         """Operations related to HITL Responses."""
         return HITLOperations(self)
 
+    @lru_cache()  # type: ignore[misc]
+    @property
+    def dags(self):
+        return DagsOperations(self)
+
 
 # This is only used for parsing. ServerResponseError is raised instead
 class _ErrorBody(BaseModel):

task-sdk/src/airflow/sdk/api/datamodels/_generated.py

@@ -154,6 +154,14 @@ class DagRunType(str, Enum):
     ASSET_TRIGGERED = "asset_triggered"
 
 
+class DagStateResponse(BaseModel):
+    """
+    Schema for DAG State response.
+    """
+
+    is_paused: Annotated[bool, Field(title="Is Paused")]
+
+
 class HITLUser(BaseModel):
     """
     Schema for a Human-in-the-loop users.

task-sdk/src/airflow/sdk/execution_time/comms.py

@@ -585,12 +585,18 @@ class HITLDetailRequestResult(HITLDetailRequest):
     type: Literal["HITLDetailRequestResult"] = "HITLDetailRequestResult"
 
 
+class DagStateResult(BaseModel):
+    is_paused: bool
+    type: Literal["DagStateResult"] = "DagStateResult"
+
+
 ToTask = Annotated[
     AssetResult
     | AssetEventsResult
     | ConnectionResult
     | DagRunStateResult
     | DRCount
+    | DagStateResult
     | ErrorResponse
     | PrevSuccessfulDagRunResult
     | SentFDs
@@ -910,6 +916,11 @@ class MaskSecret(BaseModel):
     type: Literal["MaskSecret"] = "MaskSecret"
 
 
+class GetDagState(BaseModel):
+    dag_id: str
+    type: Literal["GetDagState"] = "GetDagState"
+
+
 ToSupervisor = Annotated[
     DeferTask
     | DeleteXCom
@@ -920,6 +931,7 @@ class MaskSecret(BaseModel):
     | GetConnection
     | GetDagRunState
     | GetDRCount
+    | GetDagState
     | GetPrevSuccessfulDagRun
     | GetPreviousDagRun
     | GetTaskRescheduleStartDate

task-sdk/src/airflow/sdk/execution_time/supervisor.py

@@ -83,6 +83,7 @@
     GetAssetEventByAssetAlias,
     GetConnection,
     GetDagRunState,
+    GetDagState,
     GetDRCount,
     GetPreviousDagRun,
     GetPrevSuccessfulDagRun,
@@ -1382,6 +1383,10 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id:
             self.send_msg(resp, request_id=req_id, error=None, **dump_opts)
         elif isinstance(msg, MaskSecret):
             mask_secret(msg.value, msg.name)
+        elif isinstance(msg, GetDagState):
+            resp = self.client.dags.get_state(
+                dag_id=msg.dag_id,
+            )
         else:
             log.error("Unhandled request", msg=msg)
             self.send_msg(

task-sdk/src/airflow/sdk/execution_time/task_runner.py

@@ -64,10 +64,12 @@
     AssetEventDagRunReferenceResult,
     CommsDecoder,
     DagRunStateResult,
+    DagStateResult,
     DeferTask,
     DRCount,
     ErrorResponse,
     GetDagRunState,
+    GetDagState,
     GetDRCount,
     GetPreviousDagRun,
     GetTaskRescheduleStartDate,
@@ -550,6 +552,16 @@ def get_dagrun_state(dag_id: str, run_id: str) -> str:
 
         return response.state
 
+    @staticmethod
+    def get_dag_state(dag_id: str) -> DagStateResult:
+        """Return the state of the Dag run with the given Run ID."""
+        response = SUPERVISOR_COMMS.send(msg=GetDagState(dag_id=dag_id))
+
+        if TYPE_CHECKING:
+            assert isinstance(response, DagStateResult)
+
+        return response
+
     @property
     def log_url(self) -> str:
         run_id = quote(self.run_id)

task-sdk/tests/task_sdk/api/test_client.py

@@ -37,6 +37,7 @@
     ConnectionResponse,
     DagRunState,
     DagRunStateResponse,
+    DagStateResponse,
     HITLDetailResponse,
     HITLUser,
     VariableResponse,
@@ -1349,3 +1350,21 @@ def handle_request(request: httpx.Request) -> httpx.Response:
         assert result.params_input == {}
         assert result.responded_by_user == HITLUser(id="admin", name="admin")
         assert result.responded_at == timezone.datetime(2025, 7, 3, 0, 0, 0)
+
+
+class TestDagsOperations:
+    def test_get_state(self):
+        """Test that the client can get the state of a dag run"""
+
+        def handle_request(request: httpx.Request) -> httpx.Response:
+            if request.url.path == "/dags/test_dag/state":
+                return httpx.Response(
+                    status_code=200,
+                    json={"is_paused": False},
+                )
+            return httpx.Response(status_code=200)
+
+        client = make_client(transport=httpx.MockTransport(handle_request))
+        result = client.dags.get_state(dag_id="test_dag")
+
+        assert result == DagStateResponse(is_paused=False)