Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,14 @@

from __future__ import annotations

from datetime import datetime
from typing import Generic, Literal, TypeVar

from pydantic import computed_field

from airflow.api_fastapi.core_api.base import BaseModel
from airflow.utils.state import TaskInstanceState
from airflow.utils.types import DagRunType


class BaseEdgeResponse(BaseModel):
Expand Down Expand Up @@ -52,8 +57,46 @@ class BaseNodeResponse(BaseModel):
N = TypeVar("N", bound=BaseNodeResponse)


class GridNodeResponse(BaseModel):
"""Base Node serializer for responses."""

id: str
label: str
children: list[GridNodeResponse] | None = None
is_mapped: bool | None
setup_teardown_type: Literal["setup", "teardown"] | None = None


class GridRunsResponse(BaseModel):
"""Base Node serializer for responses."""

dag_id: str
run_id: str
queued_at: datetime | None
start_date: datetime | None
end_date: datetime | None
run_after: datetime
state: TaskInstanceState | None
run_type: DagRunType

@computed_field
def duration(self) -> int | None:
if self.start_date and self.end_date:
return (self.end_date - self.start_date).seconds
return None


class BaseGraphResponse(BaseModel, Generic[E, N]):
"""Base Graph serializer for responses."""

edges: list[E]
nodes: list[N]


class LatestRunResponse(BaseModel):
"""Base Node serializer for responses."""

id: int
dag_id: str
run_id: str
run_after: datetime
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@

from pydantic import BaseModel, Field

from airflow.api_fastapi.core_api.datamodels.ui.structure import StructureDataResponse
from airflow.utils.state import DagRunState, TaskInstanceState
from airflow.utils.types import DagRunType

Expand All @@ -40,6 +39,13 @@ class GridTaskInstanceSummary(BaseModel):
note: str | None


class LightGridTaskInstanceSummary(BaseModel):
"""Task Instance Summary model for the Grid UI."""

task_id: str
state: TaskInstanceState | None


class GridDAGRunwithTIs(BaseModel):
"""DAG Run model for the Grid UI."""

Expand All @@ -57,8 +63,15 @@ class GridDAGRunwithTIs(BaseModel):
task_instances: list[GridTaskInstanceSummary]


class GridTISummaries(BaseModel):
"""DAG Run model for the Grid UI."""

run_id: str
dag_id: str
task_instances: list[LightGridTaskInstanceSummary]


class GridResponse(BaseModel):
"""Response model for the Grid UI."""

dag_runs: list[GridDAGRunwithTIs]
structure: StructureDataResponse
Loading