Skip to content

Commit 5e52684

Browse files
dolfinusdolfinus
committed
[DOP-22578] Return statistics in GET /locations
1 parent e2f8346 commit 5e52684

File tree

12 files changed

+486
-83
lines changed

12 files changed

+486
-83
lines changed

‎data_rentgen/db/repositories/dataset.py‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from sqlalchemy import (
66
ColumnElement,
77
CompoundSelect,
8+
Row,
89
Select,
910
SQLColumnExpression,
1011
any_,
@@ -102,6 +103,24 @@ async def list_by_ids(self, dataset_ids: Sequence[int]) -> list[Dataset]:
102103
result = await self._session.scalars(query)
103104
return list(result.all())
104105

106+
async def get_stats_by_location_ids(self, location_ids: Sequence[int]) -> dict[int, Row]:
107+
if not location_ids:
108+
return {}
109+
110+
query = (
111+
select(
112+
Dataset.location_id.label("location_id"),
113+
func.count(Dataset.id.distinct()).label("total_datasets"),
114+
)
115+
.where(
116+
Dataset.location_id == any_(location_ids), # type: ignore[arg-type]
117+
)
118+
.group_by(Dataset.location_id)
119+
)
120+
121+
query_result = await self._session.execute(query)
122+
return {row.location_id: row for row in query_result.all()}
123+
105124
async def _get(self, dataset: DatasetDTO) -> Dataset | None:
106125
statement = select(Dataset).where(Dataset.location_id == dataset.location.id, Dataset.name == dataset.name)
107126
return await self._session.scalar(statement)

‎data_rentgen/db/repositories/job.py‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from sqlalchemy import (
66
ColumnElement,
77
CompoundSelect,
8+
Row,
89
Select,
910
SQLColumnExpression,
1011
any_,
@@ -101,6 +102,24 @@ async def list_by_ids(self, job_ids: Sequence[int]) -> list[Job]:
101102
result = await self._session.scalars(query)
102103
return list(result.all())
103104

105+
async def get_stats_by_location_ids(self, location_ids: Sequence[int]) -> dict[int, Row]:
106+
if not location_ids:
107+
return {}
108+
109+
query = (
110+
select(
111+
Job.location_id.label("location_id"),
112+
func.count(Job.id.distinct()).label("total_jobs"),
113+
)
114+
.where(
115+
Job.location_id == any_(location_ids), # type: ignore[arg-type]
116+
)
117+
.group_by(Job.location_id)
118+
)
119+
120+
query_result = await self._session.execute(query)
121+
return {row.location_id: row for row in query_result.all()}
122+
104123
async def _get(self, job: JobDTO) -> Job | None:
105124
statement = select(Job).where(Job.location_id == job.location.id, Job.name == job.name)
106125
return await self._session.scalar(statement)

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

Lines changed: 15 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -8,43 +8,45 @@
88
from data_rentgen.server.errors import get_error_responses
99
from data_rentgen.server.errors.schemas import InvalidRequestSchema, NotFoundSchema
1010
from data_rentgen.server.schemas.v1 import (
11+
LocationDetailedResponseV1,
1112
LocationPaginateQueryV1,
12-
LocationResponseV1,
1313
PageResponseV1,
1414
UpdateLocationRequestV1,
1515
)
16-
from data_rentgen.server.services import get_user
17-
from data_rentgen.services import UnitOfWork
16+
from data_rentgen.server.services import LocationService, get_user
1817

1918
router = APIRouter(
2019
prefix="/locations",
2120
tags=["Locations"],
22-
responses=get_error_responses(include={InvalidRequestSchema, NotFoundSchema}),
21+
responses=get_error_responses(include={InvalidRequestSchema}),
2322
)
2423

2524

2625
@router.get("", summary="Paginated list of Locations")
2726
async def paginate_locations(
2827
query_args: Annotated[LocationPaginateQueryV1, Depends()],
29-
unit_of_work: Annotated[UnitOfWork, Depends()],
28+
location_service: Annotated[LocationService, Depends()],
3029
current_user: User = Depends(get_user()),
31-
) -> PageResponseV1[LocationResponseV1]:
32-
pagination = await unit_of_work.location.paginate(
30+
) -> PageResponseV1[LocationDetailedResponseV1]:
31+
pagination = await location_service.paginate(
3332
page=query_args.page,
3433
page_size=query_args.page_size,
3534
location_ids=query_args.location_id,
3635
location_type=query_args.location_type,
3736
search_query=query_args.search_query,
3837
)
39-
return PageResponseV1[LocationResponseV1].from_pagination(pagination)
38+
return PageResponseV1[LocationDetailedResponseV1].from_pagination(pagination)
4039

4140

42-
@router.patch("/{location_id}")
41+
@router.patch(
42+
"/{location_id}",
43+
responses=get_error_responses(include={InvalidRequestSchema, NotFoundSchema}),
44+
)
4345
async def update_location(
4446
location_id: int,
4547
location_data: UpdateLocationRequestV1,
46-
unit_of_work: Annotated[UnitOfWork, Depends()],
48+
location_service: Annotated[LocationService, Depends()],
4749
current_user: User = Depends(get_user()),
48-
) -> LocationResponseV1:
49-
location = await unit_of_work.location.update_external_id(location_id, location_data.external_id)
50-
return LocationResponseV1.model_validate(location)
50+
) -> LocationDetailedResponseV1:
51+
location = await location_service.update_external_id(location_id, location_data.external_id)
52+
return LocationDetailedResponseV1.model_validate(location)

‎data_rentgen/server/schemas/v1/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
RunLineageQueryV1,
2222
)
2323
from data_rentgen.server.schemas.v1.location import (
24+
LocationDetailedResponseV1,
2425
LocationPaginateQueryV1,
2526
LocationResponseV1,
2627
UpdateLocationRequestV1,
@@ -68,6 +69,7 @@
6869
"LineageOutputRelationV1",
6970
"LineageParentRelationV1",
7071
"LineageSymlinkRelationV1",
72+
"LocationDetailedResponseV1",
7173
"LocationPaginateQueryV1",
7274
"LocationResponseV1",
7375
"UpdateLocationRequestV1",

‎data_rentgen/server/schemas/v1/location.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,38 @@ class LocationResponseV1(BaseModel):
1717
model_config = ConfigDict(from_attributes=True)
1818

1919

20+
class LocationDatasetStatisticsReponseV1(BaseModel):
21+
"""Location dataset statistics response."""
22+
23+
total_datasets: int = Field(description="Total number of datasets bound to this location")
24+
25+
model_config = ConfigDict(from_attributes=True)
26+
27+
28+
class LocationJobStatisticsReponseV1(BaseModel):
29+
"""Location job statistics response."""
30+
31+
total_jobs: int = Field(description="Total number of jobs bound to this location")
32+
33+
model_config = ConfigDict(from_attributes=True)
34+
35+
36+
class LocationStatisticsReponseV1(BaseModel):
37+
"""Location statistics response."""
38+
39+
datasets: LocationDatasetStatisticsReponseV1 = Field(description="Dataset statistics")
40+
jobs: LocationJobStatisticsReponseV1 = Field(description="Dataset statistics")
41+
42+
model_config = ConfigDict(from_attributes=True)
43+
44+
45+
class LocationDetailedResponseV1(BaseModel):
46+
data: LocationResponseV1 = Field(description="Location data")
47+
statistics: LocationStatisticsReponseV1 = Field(description="Location statistics")
48+
49+
model_config = ConfigDict(from_attributes=True)
50+
51+
2052
class LocationPaginateQueryV1(PaginateQueryV1):
2153
"""Query params for Location paginate request."""
2254

‎data_rentgen/server/services/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,14 @@
22
# SPDX-License-Identifier: Apache-2.0
33
from data_rentgen.server.services.get_user import get_user
44
from data_rentgen.server.services.lineage import LineageService
5+
from data_rentgen.server.services.location import LocationService
56
from data_rentgen.server.services.operation import OperationService
67
from data_rentgen.server.services.run import RunService
78

89
__all__ = [
910
"get_user",
1011
"LineageService",
12+
"LocationService",
1113
"OperationService",
1214
"RunService",
1315
]

‎data_rentgen/server/services/location.py‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
# SPDX-FileCopyrightText: 2024-2025 MTS PJSC
2+
# SPDX-License-Identifier: Apache-2.0
3+
from dataclasses import dataclass
4+
from typing import Annotated
5+
6+
from fastapi import Depends
7+
from sqlalchemy import Row
8+
9+
from data_rentgen.db.models.location import Location
10+
from data_rentgen.dto.pagination import PaginationDTO
11+
from data_rentgen.services.uow import UnitOfWork
12+
13+
14+
@dataclass
15+
class LocationServiceDatasetStatistics:
16+
total_datasets: int = 0
17+
18+
@classmethod
19+
def from_row(cls, row: Row | None):
20+
if not row:
21+
return cls()
22+
23+
return cls(total_datasets=row.total_datasets)
24+
25+
26+
@dataclass
27+
class LocationServiceJobStatistics:
28+
total_jobs: int = 0
29+
30+
@classmethod
31+
def from_row(cls, row: Row | None):
32+
if not row:
33+
return cls()
34+
35+
return cls(total_jobs=row.total_jobs)
36+
37+
38+
@dataclass
39+
class LocationServiceStatistics:
40+
datasets: LocationServiceDatasetStatistics
41+
jobs: LocationServiceJobStatistics
42+
43+
44+
@dataclass
45+
class LocationServiceResult:
46+
data: Location
47+
statistics: LocationServiceStatistics
48+
49+
50+
class LocationServicePaginatedResult(PaginationDTO[LocationServiceResult]):
51+
pass
52+
53+
54+
class LocationService:
55+
def __init__(self, uow: Annotated[UnitOfWork, Depends()]):
56+
self._uow = uow
57+
58+
async def paginate(
59+
self,
60+
page: int,
61+
page_size: int,
62+
location_ids: list[int],
63+
location_type: str | None,
64+
search_query: str | None,
65+
) -> LocationServicePaginatedResult:
66+
pagination = await self._uow.location.paginate(
67+
page=page,
68+
page_size=page_size,
69+
location_ids=location_ids,
70+
location_type=location_type,
71+
search_query=search_query,
72+
)
73+
location_ids = [item.id for item in pagination.items]
74+
dataset_stats = await self._uow.dataset.get_stats_by_location_ids(location_ids)
75+
job_stats = await self._uow.job.get_stats_by_location_ids(location_ids)
76+
77+
return LocationServicePaginatedResult(
78+
page=pagination.page,
79+
page_size=pagination.page_size,
80+
total_count=pagination.total_count,
81+
items=[
82+
LocationServiceResult(
83+
data=location,
84+
statistics=LocationServiceStatistics(
85+
datasets=LocationServiceDatasetStatistics.from_row(dataset_stats.get(location.id)),
86+
jobs=LocationServiceJobStatistics.from_row(job_stats.get(location.id)),
87+
),
88+
)
89+
for location in pagination.items
90+
],
91+
)
92+
93+
async def update_external_id(
94+
self,
95+
location_id: int,
96+
external_id: str | None,
97+
) -> LocationServiceResult:
98+
location = await self._uow.location.update_external_id(location_id, external_id)
99+
dataset_stats = await self._uow.dataset.get_stats_by_location_ids([location.id])
100+
job_stats = await self._uow.job.get_stats_by_location_ids([location.id])
101+
return LocationServiceResult(
102+
data=location,
103+
statistics=LocationServiceStatistics(
104+
datasets=LocationServiceDatasetStatistics.from_row(dataset_stats.get(location.id)),
105+
jobs=LocationServiceJobStatistics.from_row(job_stats.get(location.id)),
106+
),
107+
)

‎docs/changelog/next_release/160.breaking.rst‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
Change response schema of ``GET /locations`` from:
2+
3+
.. code:: python
4+
5+
{
6+
"meta": {...},
7+
"items": [
8+
{
9+
"kind": "LOCATION",
10+
"id": ...,
11+
# ...
12+
}
13+
],
14+
}
15+
16+
to:
17+
18+
.. code:: python
19+
20+
{
21+
"meta": {...},
22+
"items": [
23+
{
24+
"data": {
25+
"kind": "LOCATION",
26+
"id": ...,
27+
# ...
28+
},
29+
"statistics": {
30+
"datasets": {"total_datasets": 2},
31+
"jobs": {"total_jobs": 0},
32+
},
33+
}
34+
],
35+
}
36+
37+
Same for ``PATCH /datasets/:id`` - before:
38+
39+
.. code:: python
40+
41+
{
42+
"kind": "LOCATION",
43+
"id": ...,
44+
# ...
45+
}
46+
47+
after:
48+
49+
.. code:: python
50+
51+
{
52+
"data": {
53+
"kind": "LOCATION",
54+
"id": ...,
55+
# ...
56+
},
57+
"statistics": {
58+
"datasets": {"total_datasets": 2},
59+
"jobs": {"total_jobs": 0},
60+
},
61+
}
62+
63+
This allows to show location statistics in UI.

0 commit comments

Comments
 (0)