[DOP-22531] Add location_id filter for dataset (#294)

TiGrib · web-flow · commit 4c858f085f74 · 2025-08-21T12:51:03.000+03:00
* [DOP-22531] Add location_id filter for dataset, fix type for pagination query

* [DOP-22531] Add changelog
diff --git a/data_rentgen/db/repositories/dataset.py b/data_rentgen/db/repositories/dataset.py
@@ -44,12 +44,16 @@ async def paginate(
         page_size: int,
         dataset_ids: Collection[int],
         tag_value_ids: Collection[int],
+        location_id: int | None,
         search_query: str | None,
     ) -> PaginationDTO[Dataset]:
         where = []
         if dataset_ids:
             where.append(Dataset.id == any_(list(dataset_ids)))  # type: ignore[arg-type]
 
+        if location_id:
+            where.append(Dataset.location_id == location_id)
+
         if tag_value_ids:
             tv_ids = list(tag_value_ids)
             dataset_ids_subq = (
diff --git a/data_rentgen/server/api/v1/router/dataset.py b/data_rentgen/server/api/v1/router/dataset.py
@@ -29,7 +29,7 @@
 
 @router.get("", summary="Paginated list of Datasets")
 async def paginate_datasets(
-    query_args: Annotated[DatasetPaginateQueryV1, Depends()],
+    query_args: Annotated[DatasetPaginateQueryV1, Query()],
     dataset_service: Annotated[DatasetService, Depends()],
     current_user: Annotated[User, Depends(get_user())],
 ) -> PageResponseV1[DatasetDetailedResponseV1]:
@@ -38,6 +38,7 @@ async def paginate_datasets(
         page_size=query_args.page_size,
         dataset_ids=query_args.dataset_id,
         tag_value_ids=query_args.tag_value_id,
+        location_id=query_args.location_id,
         search_query=query_args.search_query,
     )
     return PageResponseV1[DatasetDetailedResponseV1].from_pagination(pagination)
diff --git a/data_rentgen/server/api/v1/router/job.py b/data_rentgen/server/api/v1/router/job.py
@@ -27,7 +27,7 @@
 
 @router.get("", summary="Paginated list of Jobs")
 async def paginate_jobs(
-    query_args: Annotated[JobPaginateQueryV1, Depends()],
+    query_args: Annotated[JobPaginateQueryV1, Query()],
     job_service: Annotated[JobService, Depends()],
     current_user: Annotated[User, Depends(get_user())],
 ) -> PageResponseV1[JobDetailedResponseV1]:
diff --git a/data_rentgen/server/api/v1/router/location.py b/data_rentgen/server/api/v1/router/location.py
@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from typing import Annotated
 
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, Query
 
 from data_rentgen.db.models import User
 from data_rentgen.server.errors import get_error_responses
@@ -28,7 +28,7 @@
     responses=get_error_responses(include={NotAuthorizedSchema, NotAuthorizedRedirectSchema, InvalidRequestSchema}),
 )
 async def paginate_locations(
-    query_args: Annotated[LocationPaginateQueryV1, Depends()],
+    query_args: Annotated[LocationPaginateQueryV1, Query()],
     location_service: Annotated[LocationService, Depends()],
     current_user: Annotated[User, Depends(get_user())],
 ) -> PageResponseV1[LocationDetailedResponseV1]:
diff --git a/data_rentgen/server/api/v1/router/operation.py b/data_rentgen/server/api/v1/router/operation.py
@@ -27,7 +27,7 @@
 
 @router.get("", summary="Paginated list of Operations")
 async def operations(
-    query_args: Annotated[OperationQueryV1, Depends()],
+    query_args: Annotated[OperationQueryV1, Query()],
     unit_of_work: Annotated[UnitOfWork, Depends()],
     operation_service: Annotated[OperationService, Depends()],
     current_user: Annotated[User, Depends(get_user())],
diff --git a/data_rentgen/server/api/v1/router/run.py b/data_rentgen/server/api/v1/router/run.py
@@ -26,7 +26,7 @@
 
 @router.get("", summary="Paginated list of Runs")
 async def runs(
-    query_args: Annotated[RunsQueryV1, Depends()],
+    query_args: Annotated[RunsQueryV1, Query()],
     run_service: Annotated[RunService, Depends()],
     current_user: Annotated[User, Depends(get_user())],
 ) -> PageResponseV1[RunDetailedResponseV1]:
diff --git a/data_rentgen/server/api/v1/router/tag.py b/data_rentgen/server/api/v1/router/tag.py
@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from typing import Annotated
 
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, Query
 
 from data_rentgen.db.models.user import User
 from data_rentgen.server.errors import get_error_responses
@@ -24,7 +24,7 @@
 
 @router.get("", summary="Paginated list of Tags")
 async def paginate_tags(
-    query_args: Annotated[TagPaginateQueryV1, Depends()],
+    query_args: Annotated[TagPaginateQueryV1, Query()],
     tag_service: Annotated[TagService, Depends()],
     current_user: Annotated[User, Depends(get_user())],
 ) -> PageResponseV1[TagDetailedResponseV1]:
diff --git a/data_rentgen/server/schemas/v1/dataset.py b/data_rentgen/server/schemas/v1/dataset.py
@@ -4,7 +4,6 @@
 
 from typing import Literal
 
-from fastapi import Query
 from pydantic import BaseModel, ConfigDict, Field
 
 from data_rentgen.server.schemas.v1.location import LocationResponseV1
@@ -56,14 +55,13 @@ class DatasetDetailedResponseV1(BaseModel):
 class DatasetPaginateQueryV1(PaginateQueryV1):
     """Query params for Dataset paginate request."""
 
-    dataset_id: list[int] = Field(Query(default_factory=list), description="Dataset id")
-    tag_value_id: list[int] = Field(Query(default_factory=list), description="Tag value id")
+    dataset_id: list[int] = Field(default_factory=list, description="Dataset id")
+    tag_value_id: list[int] = Field(default_factory=list, description="Tag value id")
+    location_id: int | None = Field(default=None, description="Location id to filter dataset")
     search_query: str | None = Field(
-        Query(
-            default=None,
-            min_length=3,
-            description="Search query",
-        ),
+        default=None,
+        min_length=3,
+        description="Search query",
     )
 
     model_config = ConfigDict(extra="forbid")
diff --git a/data_rentgen/server/schemas/v1/job.py b/data_rentgen/server/schemas/v1/job.py
@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-from fastapi import Query
 from pydantic import BaseModel, ConfigDict, Field
 
 from data_rentgen.server.schemas.v1.location import LocationResponseV1
@@ -30,13 +29,11 @@ class JobDetailedResponseV1(BaseModel):
 class JobPaginateQueryV1(PaginateQueryV1):
     """Query params for Jobs paginate request."""
 
-    job_id: list[int] = Field(Query(default_factory=list), description="Job id")
+    job_id: list[int] = Field(default_factory=list, description="Job id")
     search_query: str | None = Field(
-        Query(
-            default=None,
-            min_length=3,
-            description="Search query",
-        ),
+        default=None,
+        min_length=3,
+        description="Search query",
     )
 
     model_config = ConfigDict(extra="forbid")
diff --git a/data_rentgen/server/schemas/v1/location.py b/data_rentgen/server/schemas/v1/location.py
@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-from fastapi import Query
 from pydantic import BaseModel, ConfigDict, Field
 
 from data_rentgen.server.schemas.v1.address import AddressResponseV1
@@ -55,14 +54,12 @@ class LocationDetailedResponseV1(BaseModel):
 class LocationPaginateQueryV1(PaginateQueryV1):
     """Query params for Location paginate request."""
 
-    location_id: list[int] = Field(Query(default_factory=list), description="Location id")
-    location_type: str | None = Field(Query(default=None), description="Location type")
+    location_id: list[int] = Field(default_factory=list, description="Location id")
+    location_type: str | None = Field(default=None, description="Location type")
     search_query: str | None = Field(
-        Query(
-            default=None,
-            min_length=3,
-            description="Search query",
-        ),
+        default=None,
+        min_length=3,
+        description="Search query",
     )
 
     model_config = ConfigDict(extra="forbid")
diff --git a/data_rentgen/server/schemas/v1/operation.py b/data_rentgen/server/schemas/v1/operation.py
@@ -6,7 +6,6 @@
 from enum import IntEnum
 from uuid import UUID
 
-from fastapi import Query
 from pydantic import (
     UUID7,
     BaseModel,
@@ -102,30 +101,22 @@ class OperationQueryV1(PaginateQueryV1):
     """Query params for Operations paginate request."""
 
     since: datetime | None = Field(
-        Query(
-            default=None,
-            description="Minimum value of Operation 'created_at' field, in ISO 8601 format",
-            examples=["2008-09-15T15:53:00+05:00"],
-        ),
+        default=None,
+        description="Minimum value of Operation 'created_at' field, in ISO 8601 format",
+        examples=["2008-09-15T15:53:00+05:00"],
     )
     until: datetime | None = Field(
-        Query(
-            default=None,
-            description="Maximum value of Operation 'created_at' field, in ISO 8601 format",
-            examples=["2008-09-15T15:53:00+05:00"],
-        ),
+        default=None,
+        description="Maximum value of Operation 'created_at' field, in ISO 8601 format",
+        examples=["2008-09-15T15:53:00+05:00"],
     )
     operation_id: list[UUID7] = Field(
-        Query(
-            default_factory=list,
-            description="Operation ids, for exact match",
-        ),
+        default_factory=list,
+        description="Operation ids, for exact match",
     )
     run_id: UUID7 | None = Field(
-        Query(
-            default=None,
-            description="Run id, can be used only with 'since'",
-        ),
+        default=None,
+        description="Run id, can be used only with 'since'",
     )
 
     model_config = ConfigDict(extra="forbid")
diff --git a/data_rentgen/server/schemas/v1/run.py b/data_rentgen/server/schemas/v1/run.py
@@ -6,7 +6,6 @@
 from enum import IntEnum
 from uuid import UUID
 
-from fastapi import Query
 from pydantic import (
     UUID7,
     BaseModel,
@@ -116,46 +115,34 @@ class RunsQueryV1(PaginateQueryV1):
     """Query params for Runs paginate request."""
 
     since: datetime | None = Field(
-        Query(
-            default=None,
-            description="Minimum value of Run 'created_at' field, in ISO 8601 format",
-            examples=["2008-09-15T15:53:00+05:00"],
-        ),
+        default=None,
+        description="Minimum value of Run 'created_at' field, in ISO 8601 format",
+        examples=["2008-09-15T15:53:00+05:00"],
     )
     until: datetime | None = Field(
-        Query(
-            default=None,
-            description="Maximum value of Run 'created_at' field, in ISO 8601 format",
-            examples=["2008-09-15T15:53:00+05:00"],
-        ),
+        default=None,
+        description="Maximum value of Run 'created_at' field, in ISO 8601 format",
+        examples=["2008-09-15T15:53:00+05:00"],
     )
     run_id: list[UUID7] = Field(
-        Query(
-            default_factory=list,
-            description="Run ids, for exact match",
-        ),
+        default_factory=list,
+        description="Run ids, for exact match",
     )
     job_id: int | None = Field(
-        Query(
-            default=None,
-            description="Job id, can be used only with 'since'",
-        ),
+        default=None,
+        description="Job id, can be used only with 'since'",
     )
 
     parent_run_id: UUID7 | None = Field(
-        Query(
-            default=None,
-            description="Parent run id, can be used only with 'since' and 'until'",
-            examples=["01913217-b761-7b1a-bb52-489da9c8b9c8"],
-        ),
+        default=None,
+        description="Parent run id, can be used only with 'since' and 'until'",
+        examples=["01913217-b761-7b1a-bb52-489da9c8b9c8"],
     )
 
     search_query: str | None = Field(
-        Query(
-            default=None,
-            min_length=3,
-            description="Search query",
-        ),
+        default=None,
+        min_length=3,
+        description="Search query",
     )
 
     model_config = ConfigDict(extra="forbid")
diff --git a/data_rentgen/server/schemas/v1/tag.py b/data_rentgen/server/schemas/v1/tag.py
@@ -1,6 +1,5 @@
 # SPDX-FileCopyrightText: 2024-2025 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
-from fastapi import Query
 from pydantic import BaseModel, ConfigDict, Field
 
 from data_rentgen.server.schemas.v1.pagination import PaginateQueryV1
@@ -31,13 +30,11 @@ class TagDetailedResponseV1(BaseModel):
 class TagPaginateQueryV1(PaginateQueryV1):
     """Query params for Tag paginate request."""
 
-    tag_id: list[int] = Field(Query(default_factory=list), description="Tag id")
+    tag_id: list[int] = Field(default_factory=list, description="Tag id")
     search_query: str | None = Field(
-        Query(
-            default=None,
-            min_length=3,
-            description="Search query",
-        ),
+        default=None,
+        min_length=3,
+        description="Search query",
     )
 
     model_config = ConfigDict(extra="forbid")
diff --git a/data_rentgen/server/services/dataset.py b/data_rentgen/server/services/dataset.py
@@ -42,13 +42,15 @@ async def paginate(
         page_size: int,
         dataset_ids: Collection[int],
         tag_value_ids: Collection[int],
+        location_id: int | None,
         search_query: str | None,
     ) -> DatasetServicePaginatedResult:
         pagination = await self._uow.dataset.paginate(
             page=page,
             page_size=page_size,
             dataset_ids=dataset_ids,
             tag_value_ids=tag_value_ids,
+            location_id=location_id,
             search_query=search_query,
         )
 
diff --git a/docs/changelog/next_release/294.improvement.rst b/docs/changelog/next_release/294.improvement.rst
@@ -0,0 +1 @@
+Add `location_id` for dataset paginate query. Now you can get all datasets which are related to specific location.
diff --git a/tests/test_server/test_datasets/test_get_datasets_by_id.py b/tests/test_server/test_datasets/test_get_datasets_by_id.py
@@ -229,3 +229,42 @@ async def test_get_datasets_by_multiple_tag_value_ids(
             },
         ],
     }
+
+
+@pytest.mark.asyncio
+async def test_get_datasets_by_location_id(
+    test_client: AsyncClient,
+    datasets: list[Dataset],
+    async_session: AsyncSession,
+    mocked_user: MockedUser,
+):
+    # Each dataset at this ficture has it's own location
+    dataset = datasets[0]
+    [dataset] = await enrich_datasets([dataset], async_session)
+
+    response = await test_client.get(
+        "v1/datasets",
+        headers={"Authorization": f"Bearer {mocked_user.access_token}"},
+        params={"location_id": dataset.location.id},
+    )
+
+    assert response.status_code == HTTPStatus.OK, response.json()
+    assert response.json() == {
+        "meta": {
+            "page": 1,
+            "page_size": 20,
+            "total_count": 1,
+            "pages_count": 1,
+            "has_next": False,
+            "has_previous": False,
+            "next_page": None,
+            "previous_page": None,
+        },
+        "items": [
+            {
+                "id": str(dataset.id),
+                "data": dataset_to_json(dataset),
+                "tags": tag_values_to_json(dataset.tag_values) if dataset.tag_values else [],
+            },
+        ],
+    }
diff --git a/tests/test_server/test_operations/test_get_operations.py b/tests/test_server/test_operations/test_get_operations.py
diff --git a/tests/test_server/test_operations/test_get_operations_by_run_id.py b/tests/test_server/test_operations/test_get_operations_by_run_id.py
diff --git a/tests/test_server/test_runs/test_get_runs.py b/tests/test_server/test_runs/test_get_runs.py
diff --git a/tests/test_server/test_runs/test_get_runs_by_job_id.py b/tests/test_server/test_runs/test_get_runs_by_job_id.py
diff --git a/tests/test_server/test_runs/test_get_runs_by_parent_run_id.py b/tests/test_server/test_runs/test_get_runs_by_parent_run_id.py
diff --git a/tests/test_server/test_runs/test_search_runs.py b/tests/test_server/test_runs/test_search_runs.py
