[DOP-22578] Return statistics in GET /locations

data_rentgen/db/repositories/dataset.py

@@ -5,6 +5,7 @@
 from sqlalchemy import (
     ColumnElement,
     CompoundSelect,
+    Row,
     Select,
     SQLColumnExpression,
     any_,
@@ -102,6 +103,24 @@ async def list_by_ids(self, dataset_ids: Sequence[int]) -> list[Dataset]:
         result = await self._session.scalars(query)
         return list(result.all())
 
+    async def get_stats_by_location_ids(self, location_ids: Sequence[int]) -> dict[int, Row]:
+        if not location_ids:
+            return {}
+
+        query = (
+            select(
+                Dataset.location_id.label("location_id"),
+                func.count(Dataset.id.distinct()).label("total_datasets"),
+            )
+            .where(
+                Dataset.location_id == any_(location_ids),  # type: ignore[arg-type]
+            )
+            .group_by(Dataset.location_id)
+        )
+
+        query_result = await self._session.execute(query)
+        return {row.location_id: row for row in query_result.all()}
+
     async def _get(self, dataset: DatasetDTO) -> Dataset | None:
         statement = select(Dataset).where(Dataset.location_id == dataset.location.id, Dataset.name == dataset.name)
         return await self._session.scalar(statement)

data_rentgen/db/repositories/job.py

@@ -5,6 +5,7 @@
 from sqlalchemy import (
     ColumnElement,
     CompoundSelect,
+    Row,
     Select,
     SQLColumnExpression,
     any_,
@@ -101,6 +102,24 @@ async def list_by_ids(self, job_ids: Sequence[int]) -> list[Job]:
         result = await self._session.scalars(query)
         return list(result.all())
 
+    async def get_stats_by_location_ids(self, location_ids: Sequence[int]) -> dict[int, Row]:
+        if not location_ids:
+            return {}
+
+        query = (
+            select(
+                Job.location_id.label("location_id"),
+                func.count(Job.id.distinct()).label("total_jobs"),
+            )
+            .where(
+                Job.location_id == any_(location_ids),  # type: ignore[arg-type]
+            )
+            .group_by(Job.location_id)
+        )
+
+        query_result = await self._session.execute(query)
+        return {row.location_id: row for row in query_result.all()}
+
     async def _get(self, job: JobDTO) -> Job | None:
         statement = select(Job).where(Job.location_id == job.location.id, Job.name == job.name)
         return await self._session.scalar(statement)

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

@@ -8,43 +8,45 @@
 from data_rentgen.server.errors import get_error_responses
 from data_rentgen.server.errors.schemas import InvalidRequestSchema, NotFoundSchema
 from data_rentgen.server.schemas.v1 import (
+    LocationDetailedResponseV1,
     LocationPaginateQueryV1,
-    LocationResponseV1,
     PageResponseV1,
     UpdateLocationRequestV1,
 )
-from data_rentgen.server.services import get_user
-from data_rentgen.services import UnitOfWork
+from data_rentgen.server.services import LocationService, get_user
 
 router = APIRouter(
     prefix="/locations",
     tags=["Locations"],
-    responses=get_error_responses(include={InvalidRequestSchema, NotFoundSchema}),
+    responses=get_error_responses(include={InvalidRequestSchema}),
 )
 
 
 @router.get("", summary="Paginated list of Locations")
 async def paginate_locations(
     query_args: Annotated[LocationPaginateQueryV1, Depends()],
-    unit_of_work: Annotated[UnitOfWork, Depends()],
+    location_service: Annotated[LocationService, Depends()],
     current_user: User = Depends(get_user()),
-) -> PageResponseV1[LocationResponseV1]:
-    pagination = await unit_of_work.location.paginate(
+) -> PageResponseV1[LocationDetailedResponseV1]:
+    pagination = await location_service.paginate(
         page=query_args.page,
         page_size=query_args.page_size,
         location_ids=query_args.location_id,
         location_type=query_args.location_type,
         search_query=query_args.search_query,
     )
-    return PageResponseV1[LocationResponseV1].from_pagination(pagination)
+    return PageResponseV1[LocationDetailedResponseV1].from_pagination(pagination)
 
 
-@router.patch("/{location_id}")
+@router.patch(
+    "/{location_id}",
+    responses=get_error_responses(include={InvalidRequestSchema, NotFoundSchema}),
+)
 async def update_location(
     location_id: int,
     location_data: UpdateLocationRequestV1,
-    unit_of_work: Annotated[UnitOfWork, Depends()],
+    location_service: Annotated[LocationService, Depends()],
     current_user: User = Depends(get_user()),
-) -> LocationResponseV1:
-    location = await unit_of_work.location.update_external_id(location_id, location_data.external_id)
-    return LocationResponseV1.model_validate(location)
+) -> LocationDetailedResponseV1:
+    location = await location_service.update_external_id(location_id, location_data.external_id)
+    return LocationDetailedResponseV1.model_validate(location)

data_rentgen/server/schemas/v1/__init__.py

@@ -21,6 +21,7 @@
     RunLineageQueryV1,
 )
 from data_rentgen.server.schemas.v1.location import (
+    LocationDetailedResponseV1,
     LocationPaginateQueryV1,
     LocationResponseV1,
     UpdateLocationRequestV1,
@@ -68,6 +69,7 @@
     "LineageOutputRelationV1",
     "LineageParentRelationV1",
     "LineageSymlinkRelationV1",
+    "LocationDetailedResponseV1",
     "LocationPaginateQueryV1",
     "LocationResponseV1",
     "UpdateLocationRequestV1",

data_rentgen/server/schemas/v1/location.py

@@ -17,6 +17,38 @@ class LocationResponseV1(BaseModel):
     model_config = ConfigDict(from_attributes=True)
 
 
+class LocationDatasetStatisticsReponseV1(BaseModel):
+    """Location dataset statistics response."""
+
+    total_datasets: int = Field(description="Total number of datasets bound to this location")
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class LocationJobStatisticsReponseV1(BaseModel):
+    """Location job statistics response."""
+
+    total_jobs: int = Field(description="Total number of jobs bound to this location")
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class LocationStatisticsReponseV1(BaseModel):
+    """Location statistics response."""
+
+    datasets: LocationDatasetStatisticsReponseV1 = Field(description="Dataset statistics")
+    jobs: LocationJobStatisticsReponseV1 = Field(description="Dataset statistics")
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class LocationDetailedResponseV1(BaseModel):
+    data: LocationResponseV1 = Field(description="Location data")
+    statistics: LocationStatisticsReponseV1 = Field(description="Location statistics")
+
+    model_config = ConfigDict(from_attributes=True)
+
+
 class LocationPaginateQueryV1(PaginateQueryV1):
     """Query params for Location paginate request."""
 

data_rentgen/server/services/__init__.py

@@ -2,12 +2,14 @@
 # SPDX-License-Identifier: Apache-2.0
 from data_rentgen.server.services.get_user import get_user
 from data_rentgen.server.services.lineage import LineageService
+from data_rentgen.server.services.location import LocationService
 from data_rentgen.server.services.operation import OperationService
 from data_rentgen.server.services.run import RunService
 
 __all__ = [
     "get_user",
     "LineageService",
+    "LocationService",
     "OperationService",
     "RunService",
 ]

data_rentgen/server/services/location.py

@@ -0,0 +1,107 @@
+# SPDX-FileCopyrightText: 2024-2025 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+from typing import Annotated
+
+from fastapi import Depends
+from sqlalchemy import Row
+
+from data_rentgen.db.models.location import Location
+from data_rentgen.dto.pagination import PaginationDTO
+from data_rentgen.services.uow import UnitOfWork
+
+
+@dataclass
+class LocationServiceDatasetStatistics:
+    total_datasets: int = 0
+
+    @classmethod
+    def from_row(cls, row: Row | None):
+        if not row:
+            return cls()
+
+        return cls(total_datasets=row.total_datasets)
+
+
+@dataclass
+class LocationServiceJobStatistics:
+    total_jobs: int = 0
+
+    @classmethod
+    def from_row(cls, row: Row | None):
+        if not row:
+            return cls()
+
+        return cls(total_jobs=row.total_jobs)
+
+
+@dataclass
+class LocationServiceStatistics:
+    datasets: LocationServiceDatasetStatistics
+    jobs: LocationServiceJobStatistics
+
+
+@dataclass
+class LocationServiceResult:
+    data: Location
+    statistics: LocationServiceStatistics
+
+
+class LocationServicePaginatedResult(PaginationDTO[LocationServiceResult]):
+    pass
+
+
+class LocationService:
+    def __init__(self, uow: Annotated[UnitOfWork, Depends()]):
+        self._uow = uow
+
+    async def paginate(
+        self,
+        page: int,
+        page_size: int,
+        location_ids: list[int],
+        location_type: str | None,
+        search_query: str | None,
+    ) -> LocationServicePaginatedResult:
+        pagination = await self._uow.location.paginate(
+            page=page,
+            page_size=page_size,
+            location_ids=location_ids,
+            location_type=location_type,
+            search_query=search_query,
+        )
+        location_ids = [item.id for item in pagination.items]
+        dataset_stats = await self._uow.dataset.get_stats_by_location_ids(location_ids)
+        job_stats = await self._uow.job.get_stats_by_location_ids(location_ids)
+
+        return LocationServicePaginatedResult(
+            page=pagination.page,
+            page_size=pagination.page_size,
+            total_count=pagination.total_count,
+            items=[
+                LocationServiceResult(
+                    data=location,
+                    statistics=LocationServiceStatistics(
+                        datasets=LocationServiceDatasetStatistics.from_row(dataset_stats.get(location.id)),
+                        jobs=LocationServiceJobStatistics.from_row(job_stats.get(location.id)),
+                    ),
+                )
+                for location in pagination.items
+            ],
+        )
+
+    async def update_external_id(
+        self,
+        location_id: int,
+        external_id: str | None,
+    ) -> LocationServiceResult:
+        location = await self._uow.location.update_external_id(location_id, external_id)
+        dataset_stats = await self._uow.dataset.get_stats_by_location_ids([location.id])
+        job_stats = await self._uow.job.get_stats_by_location_ids([location.id])
+        return LocationServiceResult(
+            data=location,
+            statistics=LocationServiceStatistics(
+                datasets=LocationServiceDatasetStatistics.from_row(dataset_stats.get(location.id)),
+                jobs=LocationServiceJobStatistics.from_row(job_stats.get(location.id)),
+            ),
+        )

docs/changelog/next_release/160.breaking.rst

@@ -0,0 +1,63 @@
+Change response schema of ``GET /locations`` from:
+
+.. code:: python
+
+    {
+        "meta": {...},
+        "items": [
+            {
+                "kind": "LOCATION",
+                "id": ...,
+                # ...
+            }
+        ],
+    }
+
+to:
+
+.. code:: python
+
+    {
+        "meta": {...},
+        "items": [
+            {
+                "data": {
+                    "kind": "LOCATION",
+                    "id": ...,
+                    # ...
+                },
+                "statistics": {
+                    "datasets": {"total_datasets": 2},
+                    "jobs": {"total_jobs": 0},
+                },
+            }
+        ],
+    }
+
+Same for ``PATCH /datasets/:id`` - before:
+
+.. code:: python
+
+    {
+        "kind": "LOCATION",
+        "id": ...,
+        # ...
+    }
+
+after:
+
+.. code:: python
+
+    {
+        "data": {
+            "kind": "LOCATION",
+            "id": ...,
+            # ...
+        },
+        "statistics": {
+            "datasets": {"total_datasets": 2},
+            "jobs": {"total_jobs": 0},
+        },
+    }
+
+This allows to show location statistics in UI.