omkar-foss commented on code in PR #44220:
URL: https://github.com/apache/airflow/pull/44220#discussion_r1852796424
##########
airflow/api_fastapi/core_api/datamodels/task_instances.py:
##########
@@ -150,3 +152,54 @@ class TaskInstanceHistoryCollectionResponse(BaseModel):
task_instances: list[TaskInstanceHistoryResponse]
total_entries: int
+
+
+class ClearTaskInstancesBody(BaseModel):
+ """Request body for Clear Task Instances endpoint."""
+
+ dry_run: bool = True
+ start_date: AwareDatetime | None = None
+ end_date: AwareDatetime | None = None
+ only_failed: bool = True
+ only_running: bool = False
+ reset_dag_runs: bool = False
+ task_ids: list[str] | None = None
+ dag_run_id: str | None = None
+ include_upstream: bool = False
+ include_downstream: bool = False
+ include_future: bool = False
+ include_past: bool = False
+
+ @model_validator(mode="before")
+ @classmethod
+ def validate_model(cls, data: Any) -> Any:
+ """Validate clear task instance form."""
+ if data.get("only_failed") and data.get("only_running"):
+ raise ValidationError("only_failed and only_running both are set
to True")
+ if data.get("start_date") and data.get("end_date"):
+ if data.get("start_date") > data.get("end_date"):
+ raise ValidationError("end_date is sooner than start_date")
+ if data.get("start_date") and data.get("end_date") and
data.get("dag_run_id"):
+ raise ValidationError("Exactly one of dag_run_id or (start_date
and end_date) must be provided")
+ if data.get("start_date") and data.get("dag_run_id"):
+ raise ValidationError("Exactly one of dag_run_id or start_date
must be provided")
+ if data.get("end_date") and data.get("dag_run_id"):
+ raise ValidationError("Exactly one of dag_run_id or end_date must
be provided")
+ if isinstance(data.get("task_ids"), list) and
len(data.get("task_ids")) < 1:
+ raise ValidationError("task_ids list should have at least 1
element.")
+ return data
+
+
+class TaskInstanceReferenceResponse(BaseModel):
+ """Task Instance Reference serializer for responses."""
+
+ task_id: str
+ dag_run_id: str = Field(validation_alias="run_id")
+ dag_id: str
+ logical_date: datetime
+
+
+class TaskInstanceReferenceCollectionResponse(BaseModel):
+ """Task Instance Reference collection serializer for responses."""
+
+ task_instances: list[TaskInstanceReferenceResponse]
Review Comment:
Done, added @pierrejeambrun. `total_entries` will just be the length of the
`task_instances`, so it'll serve the purpose of consistency.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
