[DOP-29593] Add filter by Run.started_at and Run.ended_at

Author: dolfinus

data_rentgen/db/repositories/run.py

@@ -84,7 +84,7 @@
 
 
 class RunRepository(Repository[Run]):
-    async def paginate(
+    async def paginate(  # noqa: PLR0912, C901
         self,
         page: int,
         page_size: int,
@@ -95,7 +95,11 @@ async def paginate(
         parent_run_id: UUID | None,
         search_query: str | None,
         job_type: Collection[str],
-        status: Collection[str],
+        status: Collection[RunStatus],
+        started_since: datetime | None,
+        started_until: datetime | None,
+        ended_since: datetime | None,
+        ended_until: datetime | None,
     ) -> PaginationDTO[Run]:
         # do not use `tuple_(Run.created_at, Run.id).in_(...),
         # as this is too complex filter for Postgres to make an optimal query plan
@@ -133,8 +137,15 @@ async def paginate(
         if parent_run_id:
             where.append(Run.parent_run_id == parent_run_id)
         if status:
-            serialize_status: Collection[RunStatus] = [RunStatus[status] for status in status]
-            where.append(Run.status == any_(serialize_status))  # type: ignore[arg-type]
+            where.append(Run.status == any_(status))  # type: ignore[arg-type]
+        if started_since:
+            where.append(Run.started_at >= started_since)
+        if started_until:
+            where.append(Run.started_at <= started_until)
+        if ended_since:
+            where.append(Run.ended_at >= ended_since)
+        if ended_until:
+            where.append(Run.ended_at <= ended_until)
 
         query: Select | CompoundSelect
         order_by: list[ColumnElement | SQLColumnExpression]

data_rentgen/server/api/v1/router/run.py

@@ -41,6 +41,10 @@ async def runs(
         search_query=query_args.search_query,
         job_type=query_args.job_type,
         status=query_args.status,
+        started_since=query_args.started_since,
+        started_until=query_args.started_until,
+        ended_since=query_args.ended_since,
+        ended_until=query_args.ended_until,
     )
     return PageResponseV1[RunDetailedResponseV1].from_pagination(pagination)
 

data_rentgen/server/schemas/v1/run.py

@@ -43,10 +43,19 @@ def __str__(self) -> str:
 
 class RunStatusForQueryV1(StrEnum):
     UNKNOWN = "UNKNOWN"
+    """No data about status"""
+
     STARTED = "STARTED"
+    """Received START event"""
+
     SUCCEEDED = "SUCCEEDED"
+    """Finished successfully"""
+
     FAILED = "FAILED"
+    """Internal failure"""
+
     KILLED = "KILLED"
+    """Killed externally, e.g. by user request or in case of OOM"""
 
 
 class RunResponseV1(BaseModel):
@@ -150,16 +159,42 @@ class RunsQueryV1(PaginateQueryV1):
     job_type: list[str] = Field(
         default_factory=list,
         description="Filter runs by type of a Job",
-        examples=["SPARK_APPLICATION", "AIRFLOW_TASK"],
+        examples=[["SPARK_APPLICATION", "AIRFLOW_TASK"]],
+    )
+    status: list[RunStatusForQueryV1] = Field(
+        default_factory=list,
+        description="Filter runs by status",
+        examples=[["SUCCEEDED", "FAILED"]],
     )
-    status: list[RunStatusForQueryV1] = Field(default_factory=list, description="Filter runs by status")
 
     search_query: str | None = Field(
         default=None,
         min_length=3,
         description="Search query",
     )
 
+    started_since: datetime | None = Field(
+        default=None,
+        description="Minimum value of Run 'started_at' field, in ISO 8601 format",
+        examples=["2008-09-15T15:53:00+05:00"],
+    )
+    started_until: datetime | None = Field(
+        default=None,
+        description="Maximum value of Run 'started_at' field, in ISO 8601 format",
+        examples=["2008-09-15T15:53:00+05:00"],
+    )
+
+    ended_since: datetime | None = Field(
+        default=None,
+        description="Minimum value of Run 'ended_at' field, in ISO 8601 format",
+        examples=["2008-09-15T15:53:00+05:00"],
+    )
+    ended_until: datetime | None = Field(
+        default=None,
+        description="Maximum value of Run 'ended_at' field, in ISO 8601 format",
+        examples=["2008-09-15T15:53:00+05:00"],
+    )
+
     model_config = ConfigDict(extra="forbid")
 
     @field_validator("until", mode="after")

data_rentgen/server/services/run.py

@@ -9,7 +9,7 @@
 from fastapi import Depends
 from sqlalchemy import Row
 
-from data_rentgen.db.models.run import Run
+from data_rentgen.db.models import Run, RunStatus
 from data_rentgen.dto.pagination import PaginationDTO
 from data_rentgen.services.uow import UnitOfWork
 
@@ -82,6 +82,10 @@ async def paginate(
         search_query: str | None,
         job_type: Collection[str],
         status: Collection[str],
+        started_since: datetime | None,
+        started_until: datetime | None,
+        ended_since: datetime | None,
+        ended_until: datetime | None,
     ) -> RunServicePaginatedResult:
         pagination = await self._uow.run.paginate(
             page=page,
@@ -93,7 +97,11 @@ async def paginate(
             parent_run_id=parent_run_id,
             search_query=search_query,
             job_type=job_type,
-            status=status,
+            status=[RunStatus[s] for s in status],
+            started_since=started_since,
+            started_until=started_until,
+            ended_since=ended_since,
+            ended_until=ended_until,
         )
         run_ids = [item.id for item in pagination.items]
         input_stats = await self._uow.input.get_stats_by_run_ids(run_ids)

tests/test_server/fixtures/factories/run.py

@@ -194,16 +194,38 @@ async def runs_search(
         {"name": "spark_application_name", "type": "SPARK_APPLICATION"},
         {"name": "airflow_dag_name", "type": "AIRFLOW_DAG"},
     ]
+
+    created_at = datetime.now(tz=UTC)
     runs_kwargs = [
-        {"external_id": "application_1638922609021_0001", "status": RunStatus.KILLED},
+        {
+            "external_id": "application_1638922609021_0001",
+            "status": RunStatus.KILLED,
+            "created_at": created_at + timedelta(seconds=0.1),
+            "started_at": created_at + timedelta(seconds=1),
+            "ended_at": created_at + timedelta(seconds=60),
+        },
         {
             "external_id": "application_1638922609021_0002",
             "status": RunStatus.SUCCEEDED,
+            "created_at": created_at + timedelta(seconds=0.2),
+            "started_at": created_at + timedelta(seconds=2),
+            "ended_at": created_at + timedelta(seconds=120),
+        },
+        {
+            "external_id": "extract_task_0001",
+            "status": RunStatus.STARTED,
+            "created_at": created_at + timedelta(seconds=0.3),
+            "started_at": created_at + timedelta(seconds=3),
+            "ended_at": None,
+        },
+        {
+            "external_id": "extract_task_0002",
+            "status": RunStatus.FAILED,
+            "created_at": created_at + timedelta(seconds=0.4),
+            "started_at": created_at + timedelta(seconds=4),
+            "ended_at": created_at + timedelta(seconds=240),
         },
-        {"external_id": "extract_task_0001", "status": RunStatus.STARTED},
-        {"external_id": "extract_task_0002", "status": RunStatus.FAILED},
     ]
-    started_at = datetime.now(tz=UTC)
     async with async_session_maker() as async_session:
         jobs = []
         for kwargs in job_kwargs:
@@ -221,7 +243,6 @@ async def runs_search(
             await create_run(
                 async_session,
                 run_kwargs={
-                    "created_at": started_at + timedelta(seconds=0.1 * i),
                     "job_id": job.id,
                     "started_by_user_id": user.id,
                     **kwargs,