Repository: apache/airflow
Pull request: Add DAG run abort endpoint
Fixes #15888 

The new endpoint will set the DAG run state and all of the task instances' state to FAILED
I couldn't find any tests for the endpoints directly, so did not add a unit test... Please let me know if I missed their location and I'll happily add them!

Problem statement:
Abort a DAG Run
**Description**

It would be great having a option to abort a DAG Runs through the REST API.

**Use case / motivation**

The proposed input params would be:
 - DAG_ID
 - DAG_RUN_ID

The DAG Run should abort all its tasks running and mark them as "failed".

**Are you willing to submit a PR?**

**Related Issues**

Hints from issue discussion:
I still can't believe this hasn't been implemented yet, seems like a no brainer.

I am currently in need of this feature, and the only workaround i found is modifying the DB and marking the relevant task instances as failed. Seems like overkill for such a small task.
Hi @alexing,

We tried a different approach through the REST API to avoid touching the DB, this a code example:


```
    run_id = str(sys.argv[2])

    # We get first the task instance of the DAG run
    r = requests.get('http://localhost:8081/api/v1/dags/' + dag_id + '/dagRuns/' + str(run_id) + '/taskInstances?state=running', 
        auth=HTTPBasicAuth('admin', 'admin'))

    print(r.json()['task_instances'][0])
    task_id = r.json()['task_instances'][0]['task_id']
    execution_date = r.json()['task_instances'][0]['execution_date']

    # We set the tasks to failed
    r = requests.post('http://localhost:8081/api/v1/dags/' + dag_id + '/updateTaskInstancesState', 
        json={"dry_run": False,
            "task_id": str(task_id), 
            "execution_date": str(execution_date),
            "include_upstream": False,
            "include_downstream": False,
            "include_future": False,
            "include_past": False,
            "new_state": "failed" 
        },
        auth=HTTPBasicAuth('admin', 'admin'))

    if r.status_code == 200:
        print("Dag aborted")

```

I hope it helps.

Well, we thought similar solutions for sure. Yours is a bit better designed! Thanks! I'll take into account.

Still, this is a basic feature that should have been implemented by apache a long time ago!
> I still can't believe this hasn't been implemented yet, seems like a no brainer.

> Still, this is a basic feature that should have been implemented by apache a long time ago!

Many needed features and capabilities hasn't been implemented yet.
It's a community based project so what we have is what the community contribute.

We welcome PRs :)

> > I still can't believe this hasn't been implemented yet, seems like a no brainer.
> 
> > Still, this is a basic feature that should have been implemented by apache a long time ago!
> 
> Many needed features and capabilities hasn't been implemented yet.
> It's a community based project so what we have is what the community contribute.
> 
> We welcome PRs :)
> 

Indeed. We even will run workshop for first-time contributors at the Airflow Summit in July. More info soon. @alexing - would you be interested in taking part and contributing? It will be very cheap (all proceedings go to Apache Software Foundation). Happy to keep place for you. Contributing back is the best way to thank all the contributors who make the software available for free for you!

Commit message: fix tests and remove schema load

Patch to review:
diff --git a/airflow/api_connexion/endpoints/dag_run_endpoint.py b/airflow/api_connexion/endpoints/dag_run_endpoint.py
index feb6d24cf3247..9d4e554313c88 100644
--- a/airflow/api_connexion/endpoints/dag_run_endpoint.py
+++ b/airflow/api_connexion/endpoints/dag_run_endpoint.py
@@ -17,6 +17,7 @@
 from flask import current_app, g, request
 from marshmallow import ValidationError
 from sqlalchemy import or_
+from typing import Optional
 
 from airflow._vendor.connexion import NoContent
 from airflow.api_connexion import security
@@ -31,8 +32,9 @@
 from airflow.models import DagModel, DagRun
 from airflow.security import permissions
 from airflow.utils.session import provide_session
-from airflow.utils.state import State
+from airflow.utils.state import State, DagRunState
 from airflow.utils.types import DagRunType
+from airflow.exceptions import DagRunNotFound
 
 
 @security.requires_access(
@@ -271,3 +273,24 @@ def post_dag_run(dag_id, session):
         )
 
     raise AlreadyExists(detail=f"DAGRun with DAG ID: '{dag_id}' and DAGRun ID: '{run_id}' already exists")
+
+
+@security.requires_access(
+    [
+        (permissions.ACTION_CAN_READ, permissions.RESOURCE_DAG),
+        (permissions.ACTION_CAN_EDIT, permissions.RESOURCE_DAG_RUN)
+    ]
+)
+@provide_session
+def post_set_dag_run_state(dag_id: str, dag_run_id: str, session) -> dict:
+    """Set a state of a dag run."""
+    dag_run: Optional[DagRun] = session.query(DagRun) \
+        .filter(DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id).one_or_none()
+    if dag_run is None:
+        error_message = f'Dag Run id {dag_run_id} not found in dag {dag_id}'
+        raise NotFound(error_message)
+
+    state = request.json['state']
+    dag_run.set_state(state=DagRunState(state.lower()).value)
+    session.merge(dag_run)
+    return dagrun_schema.dump(dag_run)
diff --git a/airflow/api_connexion/openapi/v1.yaml b/airflow/api_connexion/openapi/v1.yaml
index 2ce6804bbd888..525b712badeea 100644
--- a/airflow/api_connexion/openapi/v1.yaml
+++ b/airflow/api_connexion/openapi/v1.yaml
@@ -604,6 +604,41 @@ paths:
         '404':
           $ref: '#/components/responses/NotFound'
 
+  /dags/{dag_id}/dagRuns/{dag_run_id}/state:
+    parameters:
+      - $ref: '#/components/parameters/DAGID'
+      - $ref: '#/components/parameters/DAGRunID'
+
+    post:
+      summary: Set a state of DAG run
+      description: >
+        Set a state of DAG run
+      x-openapi-router-controller: airflow.api_connexion.endpoints.dag_run_endpoint
+      operationId: post_set_dag_run_state
+      tags: [UpdateDagRunState]
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/UpdateDagRunState'
+
+      responses:
+        '200':
+          description: Success.
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/DAGRun'
+        '400':
+          $ref: '#/components/responses/BadRequest'
+        '401':
+          $ref: '#/components/responses/Unauthenticated'
+        '403':
+          $ref: '#/components/responses/PermissionDenied'
+        '404':
+          $ref: '#/components/responses/NotFound'
+
   /dags/~/dagRuns/list:
     post:
       summary: List DAG runs (batch)
@@ -1996,6 +2031,14 @@ components:
       required:
         - dag_id
 
+    UpdateDagRunState:
+      type: object
+      properties:
+        state:
+          $ref: '#/components/schemas/DagState'
+      required:
+        - state
+
     DAGRunCollection:
       type: object
       description: Collection of DAG runs.
diff --git a/airflow/api_connexion/schemas/dag_run_schema.py b/airflow/api_connexion/schemas/dag_run_schema.py
index 62ac172ce5525..a5c6dc4d9e7a6 100644
--- a/airflow/api_connexion/schemas/dag_run_schema.py
+++ b/airflow/api_connexion/schemas/dag_run_schema.py
@@ -104,6 +104,14 @@ def autofill(self, data, **kwargs):
         return data
 
 
+class SetDagRunStateFormSchema(Schema):
+    """Schema for handling the request of setting state of DAG run"""
+
+    run_id = auto_field(data_key='dag_run_id')
+    dag_id = auto_field(dump_only=True)
+    state = DagStateField(dump_only=True)
+
+
 class DAGRunCollection(NamedTuple):
     """List of DAGRuns with metadata"""
 
diff --git a/tests/api_connexion/endpoints/test_dag_run_endpoint.py b/tests/api_connexion/endpoints/test_dag_run_endpoint.py
index 71b52747dcff7..618ee8b677b81 100644
--- a/tests/api_connexion/endpoints/test_dag_run_endpoint.py
+++ b/tests/api_connexion/endpoints/test_dag_run_endpoint.py
@@ -19,12 +19,14 @@
 
 import pytest
 from parameterized import parameterized
+from freezegun import freeze_time
 
 from airflow.api_connexion.exceptions import EXCEPTIONS_LINK_MAP
 from airflow.models import DAG, DagModel, DagRun
 from airflow.security import permissions
 from airflow.utils import timezone
 from airflow.utils.session import create_session, provide_session
+from airflow.utils.state import DagRunState
 from airflow.utils.types import DagRunType
 from tests.test_utils.api_connexion_utils import assert_401, create_user, delete_user
 from tests.test_utils.config import conf_vars
@@ -1154,3 +1156,115 @@ def test_should_raises_403_unauthorized(self):
             environ_overrides={'REMOTE_USER': "test_view_dags"},
         )
         assert response.status_code == 403
+
+
+class TestPostSetDagRunState(TestDagRunEndpoint):
+
+    @parameterized.expand([
+        ("TEST_DAG_ID", "TEST_DAG_RUN_ID_1", "failed"),
+        ("TEST_DAG_ID", "TEST_DAG_RUN_ID_1", "success")
+    ])
+    @freeze_time(TestDagRunEndpoint.default_time)
+    def test_should_respond_200(self, dag_id, dag_run_id, state):
+        test_time = timezone.parse(self.default_time)
+        with create_session() as session:
+            dag = DagModel(dag_id=dag_id)
+            dag_run = DagRun(
+                dag_id=dag_id,
+                run_id=dag_run_id,
+                state=DagRunState.RUNNING,
+                run_type=DagRunType.MANUAL,
+                execution_date=test_time,
+                start_date=test_time,
+                external_trigger=True
+            )
+            session.add(dag)
+            session.add(dag_run)
+
+        request_json = {
+            "state": state
+        }
+
+        response = self.client.post(
+            "api/v1/dags/TEST_DAG_ID/dagRuns/TEST_DAG_RUN_ID_1/state",
+            json=request_json,
+            environ_overrides={"REMOTE_USER": "test"},
+        )
+
+        assert response.status_code == 200
+        assert response.json == {
+            'conf': {},
+            'dag_id': 'TEST_DAG_ID',
+            'dag_run_id': 'TEST_DAG_RUN_ID_1',
+            'end_date': self.default_time,
+            'execution_date': self.default_time,
+            'external_trigger': True,
+            'logical_date': self.default_time,
+            'start_date': self.default_time,
+            'state': state
+        }
+
+    def test_should_response_400_for_non_existing_dag_run_state(self):
+        test_time = timezone.parse(self.default_time)
+        with create_session() as session, freeze_time(self.default_time):
+            dag_id = "TEST_DAG_ID"
+            dag_run_id = "TEST_DAG_RUN_ID_1"
+            dag = DagModel(dag_id=dag_id)
+            dag_run = DagRun(
+                dag_id=dag_id,
+                run_id=dag_run_id,
+                state=DagRunState.RUNNING,
+                run_type=DagRunType.MANUAL,
+                execution_date=test_time,
+                start_date=test_time,
+                external_trigger=True
+            )
+            session.add(dag)
+            session.add(dag_run)
+
+            request_json = {
+                "state": "madeUpState"
+            }
+
+            response = self.client.post(
+                "api/v1/dags/TEST_DAG_ID/dagRuns/TEST_DAG_RUN_ID_1/state",
+                json=request_json,
+                environ_overrides={"REMOTE_USER": "test"},
+            )
+            assert response.status_code == 400
+            assert response.json == {
+                'detail': "'madeUpState' is not one of ['queued', 'running', 'success', 'failed'] - 'state'",
+                'status': 400,
+                'title': 'Bad Request',
+                'type': EXCEPTIONS_LINK_MAP[400]
+            }
+
+    def test_should_raises_401_unauthenticated(self, session):
+        response = self.client.post(
+            "api/v1/dags/TEST_DAG_ID/dagRuns/TEST_DAG_RUN_ID_1/state",
+            json={
+                "state": 'success',
+            },
+        )
+
+        assert_401(response)
+
+    def test_should_raise_403_forbidden(self):
+        response = self.client.post(
+            "api/v1/dags/TEST_DAG_ID/dagRuns/TEST_DAG_RUN_ID_1/state",
+            json={
+                "state": 'success',
+            },
+            environ_overrides={'REMOTE_USER': "test_no_permissions"},
+        )
+        assert response.status_code == 403
+
+    def test_should_respond_404(self):
+        response = self.client.post(
+            "api/v1/dags/INVALID_DAG_ID/dagRuns/TEST_DAG_RUN_ID_1/state",
+            json={
+                "state": 'success',
+            },
+            environ_overrides={"REMOTE_USER": "test"}
+        )
+        assert response.status_code == 404


Code review comments:
diff_hunk: @@ -104,6 +104,14 @@ def autofill(self, data, **kwargs):
         return data
 
 
+class SetDagRunStateFormSchema(Schema):
+    """Schema for handling the request of setting state of DAG run"""
+
+    run_id = auto_field(data_key='dag_run_id')
+    dag_id = auto_field(dump_only=True)
+    state = DagStateField(dump_only=True)
path: airflow/api_connexion/schemas/dag_run_schema.py
text: @user1:
```suggestion
class SetDagRunStateFormSchema(Schema):
    """Schema for handling the request of setting state of DAG run"""

    state = DagStateField()
```

We don’t have these fields in the schema now. Also `dump_only` means this field cannot be include in a POST form, which is opposite to what you want. Technically you can add `load_only=True`, but since this entire schema is only used to load data anyway, the flag would not practically change anything.

@author:
> Also you will still need to satisfy the linter. Again, please consider setting up `pre-commit`.

ran the static checks locally, I think it should be good now (for some reason they don't run automatically here in the GitHub actions), thanks!