Skip to content

Commit 0131897

Browse files
guan404mingshahar1roykoandjason810496
authored
Add allowed_run_types to whitelist specific dag run types (#61833)
* Add deny_dag_run_types option to limit specific dag run types Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com> * Add missing deny_dag_run_types to test expectations and CLI helper * Support single value Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com> * Update to use accept list Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com> * Update tests Co-authored-by: Shahar Epstein <60007259+shahar1@users.noreply.github.com> Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com> * Refactor run type checks to use enum values directly Co-authored-by: Andrii Roiko <57030016+roykoand@users.noreply.github.com> Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com> * Update to use the generated API model Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com> * Update create_openapi_http_exception_doc Co-authored-by: Jason(Zhe-You) Liu <68415893+jason810496@users.noreply.github.com> Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com> --------- Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com> Co-authored-by: Shahar Epstein <60007259+shahar1@users.noreply.github.com> Co-authored-by: Andrii Roiko <57030016+roykoand@users.noreply.github.com> Co-authored-by: Jason(Zhe-You) Liu <68415893+jason810496@users.noreply.github.com>
1 parent b1a3d3b commit 0131897

File tree

41 files changed

+2257
-1661
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

41 files changed

+2257
-1661
lines changed

‎airflow-core/docs/img/airflow_erd.sha256‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
8901cbe69c2bc259020d9cde171128b719388d28aa3440c73a5100d23531b6ea
1+
2eeb343138698d327ec6c3efe1fbea620bbc47112279f9184590cc29d8bab033

‎airflow-core/docs/img/airflow_erd.svg‎

Lines changed: 1634 additions & 1630 deletions
Loading

‎airflow-core/docs/migrations-ref.rst‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
3939
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
4040
| Revision ID | Revises ID | Airflow Version | Description |
4141
+=========================+==================+===================+==============================================================+
42-
| ``f8c9d7e6b5a4`` (head) | ``53ff648b8a26`` | ``3.2.0`` | Standardize UUID column format for non-PostgreSQL databases. |
42+
| ``e42d9fcd10d9`` (head) | ``f8c9d7e6b5a4`` | ``3.2.0`` | Add allowed_run_types to dag. |
43+
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
44+
| ``f8c9d7e6b5a4`` | ``53ff648b8a26`` | ``3.2.0`` | Standardize UUID column format for non-PostgreSQL databases. |
4345
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
4446
| ``53ff648b8a26`` | ``a5a3e5eb9b8d`` | ``3.2.0`` | Add revoked_token table. |
4547
+-------------------------+------------------+-------------------+--------------------------------------------------------------+

‎airflow-core/src/airflow/api/common/trigger_dag.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -161,6 +161,9 @@ def trigger_dag(
161161
if dag_model is None:
162162
raise DagNotFound(f"Dag id {dag_id} not found in DagModel")
163163

164+
if dag_model.allowed_run_types is not None and DagRunType.MANUAL not in dag_model.allowed_run_types:
165+
raise ValueError(f"Dag with dag_id: '{dag_id}' does not allow manual runs")
166+
164167
dagbag = DBDagBag()
165168
dr = _trigger_dag(
166169
dag_id=dag_id,

‎airflow-core/src/airflow/api_fastapi/core_api/datamodels/dags.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
from airflow.api_fastapi.core_api.datamodels.dag_versions import DagVersionResponse
3838
from airflow.configuration import conf
3939
from airflow.models.dag_version import DagVersion
40+
from airflow.utils.types import DagRunType
4041

4142
if TYPE_CHECKING:
4243
from airflow.serialization.definitions.param import SerializedParamsDict
@@ -83,6 +84,7 @@ class DAGResponse(BaseModel):
8384
next_dagrun_data_interval_start: datetime | None
8485
next_dagrun_data_interval_end: datetime | None
8586
next_dagrun_run_after: datetime | None
87+
allowed_run_types: list[DagRunType] | None
8688
owners: list[str]
8789

8890
@field_serializer("tags")

‎airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1743,6 +1743,13 @@ components:
17431743
format: date-time
17441744
- type: 'null'
17451745
title: Next Dagrun Run After
1746+
allowed_run_types:
1747+
anyOf:
1748+
- items:
1749+
$ref: '#/components/schemas/DagRunType'
1750+
type: array
1751+
- type: 'null'
1752+
title: Allowed Run Types
17461753
owners:
17471754
items:
17481755
type: string
@@ -1798,6 +1805,7 @@ components:
17981805
- next_dagrun_data_interval_start
17991806
- next_dagrun_data_interval_end
18001807
- next_dagrun_run_after
1808+
- allowed_run_types
18011809
- owners
18021810
- asset_expression
18031811
- latest_dag_runs

‎airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -503,6 +503,12 @@ paths:
503503
schema:
504504
$ref: '#/components/schemas/HTTPExceptionResponse'
505505
description: Forbidden
506+
'400':
507+
content:
508+
application/json:
509+
schema:
510+
$ref: '#/components/schemas/HTTPExceptionResponse'
511+
description: Bad Request
506512
'404':
507513
content:
508514
application/json:
@@ -1017,6 +1023,12 @@ paths:
10171023
schema:
10181024
$ref: '#/components/schemas/HTTPExceptionResponse'
10191025
description: Forbidden
1026+
'400':
1027+
content:
1028+
application/json:
1029+
schema:
1030+
$ref: '#/components/schemas/HTTPExceptionResponse'
1031+
description: Bad Request
10201032
'404':
10211033
content:
10221034
application/json:
@@ -10240,6 +10252,13 @@ components:
1024010252
format: date-time
1024110253
- type: 'null'
1024210254
title: Next Dagrun Run After
10255+
allowed_run_types:
10256+
anyOf:
10257+
- items:
10258+
$ref: '#/components/schemas/DagRunType'
10259+
type: array
10260+
- type: 'null'
10261+
title: Allowed Run Types
1024310262
owners:
1024410263
items:
1024510264
type: string
@@ -10376,6 +10395,7 @@ components:
1037610395
- next_dagrun_data_interval_start
1037710396
- next_dagrun_data_interval_end
1037810397
- next_dagrun_run_after
10398+
- allowed_run_types
1037910399
- owners
1038010400
- catchup
1038110401
- dag_run_timeout
@@ -10516,6 +10536,13 @@ components:
1051610536
format: date-time
1051710537
- type: 'null'
1051810538
title: Next Dagrun Run After
10539+
allowed_run_types:
10540+
anyOf:
10541+
- items:
10542+
$ref: '#/components/schemas/DagRunType'
10543+
type: array
10544+
- type: 'null'
10545+
title: Allowed Run Types
1051910546
owners:
1052010547
items:
1052110548
type: string
@@ -10552,6 +10579,7 @@ components:
1055210579
- next_dagrun_data_interval_start
1055310580
- next_dagrun_data_interval_end
1055410581
- next_dagrun_run_after
10582+
- allowed_run_types
1055510583
- owners
1055610584
- file_token
1055710585
title: DAGResponse

‎airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -375,7 +375,9 @@ def create_asset_event(
375375

376376
@assets_router.post(
377377
"/assets/{asset_id}/materialize",
378-
responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND, status.HTTP_409_CONFLICT]),
378+
responses=create_openapi_http_exception_doc(
379+
[status.HTTP_400_BAD_REQUEST, status.HTTP_404_NOT_FOUND, status.HTTP_409_CONFLICT]
380+
),
379381
dependencies=[Depends(requires_access_asset(method="POST")), Depends(action_logging())],
380382
)
381383
def materialize_asset(
@@ -403,6 +405,13 @@ def materialize_asset(
403405
)
404406

405407
dag = get_latest_version_of_dag(dag_bag, dag_id, session)
408+
409+
if dag.allowed_run_types is not None and DagRunType.MANUAL not in dag.allowed_run_types:
410+
raise HTTPException(
411+
status.HTTP_400_BAD_REQUEST,
412+
f"Dag with dag_id: '{dag_id}' does not allow manual runs",
413+
)
414+
406415
return dag.create_dagrun(
407416
run_id=dag.timetable.generate_run_id(
408417
run_type=DagRunType.MANUAL,

‎airflow-core/src/airflow/api_fastapi/core_api/routes/public/backfills.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@
4343
)
4444
from airflow.api_fastapi.core_api.security import GetUserDep, requires_access_backfill
4545
from airflow.api_fastapi.logging.decorators import action_logging
46-
from airflow.exceptions import DagNotFound
46+
from airflow.exceptions import DagNotFound, DagRunTypeNotAllowed
4747
from airflow.models import DagRun
4848
from airflow.models.backfill import (
4949
AlreadyRunningBackfill,
@@ -215,7 +215,9 @@ def cancel_backfill(backfill_id: NonNegativeInt, session: SessionDep) -> Backfil
215215

216216
@backfills_router.post(
217217
path="",
218-
responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND, status.HTTP_409_CONFLICT]),
218+
responses=create_openapi_http_exception_doc(
219+
[status.HTTP_400_BAD_REQUEST, status.HTTP_404_NOT_FOUND, status.HTTP_409_CONFLICT]
220+
),
219221
dependencies=[
220222
Depends(action_logging()),
221223
Depends(requires_access_backfill(method="POST")),
@@ -252,6 +254,11 @@ def create_backfill(
252254
status_code=status.HTTP_404_NOT_FOUND,
253255
detail=f"Could not find dag {backfill_request.dag_id}",
254256
)
257+
except DagRunTypeNotAllowed as e:
258+
raise HTTPException(
259+
status_code=status.HTTP_400_BAD_REQUEST,
260+
detail=str(e),
261+
)
255262
except (
256263
InvalidReprocessBehavior,
257264
InvalidBackfillDirection,

‎airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -459,6 +459,12 @@ def trigger_dag_run(
459459
f"DAG with dag_id: '{dag_id}' has import errors and cannot be triggered",
460460
)
461461

462+
if dm.allowed_run_types is not None and DagRunType.MANUAL not in dm.allowed_run_types:
463+
raise HTTPException(
464+
status.HTTP_400_BAD_REQUEST,
465+
f"Dag with dag_id: '{dag_id}' does not allow manual runs",
466+
)
467+
462468
referer = request.headers.get("referer")
463469
if referer:
464470
triggered_by = DagRunTriggeredByType.UI

0 commit comments

Comments
 (0)