This is an automated email from the ASF dual-hosted git repository.
pierrejeambrun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new cb4590a7486 Enforce to forbid extra fields under execution_api (#44986)
cb4590a7486 is described below
commit cb4590a748659d0c1e48aa0c9b80da5f9f2a66f9
Author: jj.lee <[email protected]>
AuthorDate: Wed Feb 5 21:53:41 2025 +0900
Enforce to forbid extra fields under execution_api (#44986)
* extra forbidden
* test
* gen_file
* fix gen_file
* rm clear_db_*
* rm testcase
* rm testcase # 2
---
.../api_fastapi/execution_api/datamodels/asset.py | 4 +--
.../execution_api/datamodels/taskinstance.py | 2 +-
.../src/airflow/sdk/api/datamodels/_generated.py | 34 ++++++++++++++++++++++
3 files changed, 37 insertions(+), 3 deletions(-)
diff --git a/airflow/api_fastapi/execution_api/datamodels/asset.py
b/airflow/api_fastapi/execution_api/datamodels/asset.py
index 29b260c291c..28d352aa231 100644
--- a/airflow/api_fastapi/execution_api/datamodels/asset.py
+++ b/airflow/api_fastapi/execution_api/datamodels/asset.py
@@ -17,7 +17,7 @@
from __future__ import annotations
-from airflow.api_fastapi.core_api.base import BaseModel
+from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel
class AssetResponse(BaseModel):
@@ -36,7 +36,7 @@ class AssetAliasResponse(BaseModel):
group: str
-class AssetProfile(BaseModel):
+class AssetProfile(StrictBaseModel):
"""
Profile of an Asset.
diff --git a/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
b/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
index c3ec79c9ddd..a4e3b1ff25f 100644
--- a/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
+++ b/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
@@ -252,7 +252,7 @@ class PrevSuccessfulDagRunResponse(BaseModel):
end_date: UtcDateTime | None = None
-class TIRuntimeCheckPayload(BaseModel):
+class TIRuntimeCheckPayload(StrictBaseModel):
"""Payload for performing Runtime checks on the TaskInstance model as
requested by the SDK."""
inlets: list[AssetProfile] | None = None
diff --git a/task_sdk/src/airflow/sdk/api/datamodels/_generated.py
b/task_sdk/src/airflow/sdk/api/datamodels/_generated.py
index 1d6d0eb4156..8395b072f9b 100644
--- a/task_sdk/src/airflow/sdk/api/datamodels/_generated.py
+++ b/task_sdk/src/airflow/sdk/api/datamodels/_generated.py
@@ -38,6 +38,9 @@ class AssetProfile(BaseModel):
AssetUriRef will have uri and asset_type defined.
"""
+ model_config = ConfigDict(
+ extra="forbid",
+ )
name: Annotated[str | None, Field(title="Name")] = None
uri: Annotated[str | None, Field(title="Uri")] = None
asset_type: Annotated[str, Field(title="Asset Type")]
@@ -110,6 +113,9 @@ class TIDeferredStatePayload(BaseModel):
Schema for updating TaskInstance to a deferred state.
"""
+ model_config = ConfigDict(
+ extra="forbid",
+ )
state: Annotated[Literal["deferred"] | None, Field(title="State")] =
"deferred"
classpath: Annotated[str, Field(title="Classpath")]
trigger_kwargs: Annotated[dict[str, Any] | None, Field(title="Trigger
Kwargs")] = None
@@ -122,6 +128,9 @@ class TIEnterRunningPayload(BaseModel):
Schema for updating TaskInstance to 'RUNNING' state with minimal required
fields.
"""
+ model_config = ConfigDict(
+ extra="forbid",
+ )
state: Annotated[Literal["running"] | None, Field(title="State")] =
"running"
hostname: Annotated[str, Field(title="Hostname")]
unixname: Annotated[str, Field(title="Unixname")]
@@ -134,6 +143,9 @@ class TIHeartbeatInfo(BaseModel):
Schema for TaskInstance heartbeat endpoint.
"""
+ model_config = ConfigDict(
+ extra="forbid",
+ )
hostname: Annotated[str, Field(title="Hostname")]
pid: Annotated[int, Field(title="Pid")]
@@ -143,6 +155,9 @@ class TIRescheduleStatePayload(BaseModel):
Schema for updating TaskInstance to a up_for_reschedule state.
"""
+ model_config = ConfigDict(
+ extra="forbid",
+ )
state: Annotated[Literal["up_for_reschedule"] | None,
Field(title="State")] = "up_for_reschedule"
reschedule_date: Annotated[datetime, Field(title="Reschedule Date")]
end_date: Annotated[datetime, Field(title="End Date")]
@@ -153,6 +168,9 @@ class TIRuntimeCheckPayload(BaseModel):
Payload for performing Runtime checks on the TaskInstance model as
requested by the SDK.
"""
+ model_config = ConfigDict(
+ extra="forbid",
+ )
inlets: Annotated[list[AssetProfile] | None, Field(title="Inlets")] = None
outlets: Annotated[list[AssetProfile] | None, Field(title="Outlets")] =
None
@@ -162,6 +180,9 @@ class TISuccessStatePayload(BaseModel):
Schema for updating TaskInstance to success state.
"""
+ model_config = ConfigDict(
+ extra="forbid",
+ )
state: Annotated[Literal["success"] | None, Field(title="State")] =
"success"
end_date: Annotated[datetime, Field(title="End Date")]
task_outlets: Annotated[list[AssetProfile] | None, Field(title="Task
Outlets")] = None
@@ -173,6 +194,9 @@ class TITargetStatePayload(BaseModel):
Schema for updating TaskInstance to a target state, excluding terminal and
running states.
"""
+ model_config = ConfigDict(
+ extra="forbid",
+ )
state: IntermediateTIState
@@ -234,6 +258,9 @@ class TaskInstance(BaseModel):
Schema for TaskInstance model with minimal required fields needed for
Runtime.
"""
+ model_config = ConfigDict(
+ extra="forbid",
+ )
id: Annotated[UUID, Field(title="Id")]
task_id: Annotated[str, Field(title="Task Id")]
dag_id: Annotated[str, Field(title="Dag Id")]
@@ -248,6 +275,9 @@ class DagRun(BaseModel):
Schema for DagRun model with minimal required fields needed for Runtime.
"""
+ model_config = ConfigDict(
+ extra="forbid",
+ )
dag_id: Annotated[str, Field(title="Dag Id")]
run_id: Annotated[str, Field(title="Run Id")]
logical_date: Annotated[datetime, Field(title="Logical Date")]
@@ -280,5 +310,9 @@ class TITerminalStatePayload(BaseModel):
Schema for updating TaskInstance to a terminal state except SUCCESS state.
"""
+ model_config = ConfigDict(
+ extra="forbid",
+ )
+
state: TerminalTIState
end_date: Annotated[datetime, Field(title="End Date")]
