diff --git a/spoolman/api/v1/filament.py b/spoolman/api/v1/filament.py index 46b4310..50a70c9 100644 --- a/spoolman/api/v1/filament.py +++ b/spoolman/api/v1/filament.py @@ -4,6 +4,7 @@ import logging from typing import Annotated, Optional from fastapi import APIRouter, Depends, Query +from fastapi.encoders import jsonable_encoder from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse from pydantic import BaseModel, Field, validator @@ -109,6 +110,7 @@ class FilamentUpdateParameters(FilamentParameters): name="Find filaments", description="Get a list of filaments that matches the search query.", response_model_exclude_none=True, + response_model=list[Filament], ) async def find( *, @@ -156,14 +158,14 @@ async def find( title="Offset", description="Offset in the full result set if a limit is set.", ), -) -> list[Filament]: +) -> JSONResponse: sort_by: dict[str, SortOrder] = {} if sort is not None: for sort_item in sort.split(","): field, direction = sort_item.split(":") sort_by[field] = SortOrder[direction.upper()] - db_items = await filament.find( + db_items, total_count = await filament.find( db=db, vendor_name=vendor_name, vendor_id=vendor_id, @@ -174,7 +176,15 @@ async def find( limit=limit, offset=offset, ) - return [Filament.from_db(db_item) for db_item in db_items] + + # Set x-total-count header for pagination + return JSONResponse( + content=jsonable_encoder( + (Filament.from_db(db_item) for db_item in db_items), + exclude_none=True, + ), + headers={"x-total-count": str(total_count)}, + ) @router.get( diff --git a/spoolman/api/v1/spool.py b/spoolman/api/v1/spool.py index 703c955..3eefc25 100644 --- a/spoolman/api/v1/spool.py +++ b/spoolman/api/v1/spool.py @@ -5,6 +5,7 @@ from datetime import datetime from typing import Annotated, Optional from fastapi import APIRouter, Depends, Query +from fastapi.encoders import jsonable_encoder from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse from pydantic import BaseModel, Field @@ -67,6 +68,7 @@ class SpoolUseParameters(BaseModel): name="Find spool", description="Get a list of spools that matches the search query.", response_model_exclude_none=True, + response_model=list[Spool], ) async def find( *, @@ -129,14 +131,14 @@ async def find( title="Offset", description="Offset in the full result set if a limit is set.", ), -) -> list[Spool]: +) -> JSONResponse: sort_by: dict[str, SortOrder] = {} if sort is not None: for sort_item in sort.split(","): field, direction = sort_item.split(":") sort_by[field] = SortOrder[direction.upper()] - db_items = await spool.find( + db_items, total_count = await spool.find( db=db, filament_name=filament_name, filament_id=filament_id, @@ -150,7 +152,15 @@ async def find( limit=limit, offset=offset, ) - return [Spool.from_db(db_item) for db_item in db_items] + + # Set x-total-count header for pagination + return JSONResponse( + content=jsonable_encoder( + (Spool.from_db(db_item) for db_item in db_items), + exclude_none=True, + ), + headers={"x-total-count": str(total_count)}, + ) @router.get( diff --git a/spoolman/api/v1/vendor.py b/spoolman/api/v1/vendor.py index 8f9f348..6722c00 100644 --- a/spoolman/api/v1/vendor.py +++ b/spoolman/api/v1/vendor.py @@ -3,7 +3,9 @@ from typing import Annotated, Optional from fastapi import APIRouter, Depends, Query +from fastapi.encoders import jsonable_encoder from fastapi.exceptions import RequestValidationError +from fastapi.responses import JSONResponse from pydantic import BaseModel, Field from pydantic.error_wrappers import ErrorWrapper from sqlalchemy.ext.asyncio import AsyncSession @@ -44,6 +46,7 @@ class VendorUpdateParameters(VendorParameters): name="Find vendor", description="Get a list of vendors that matches the search query.", response_model_exclude_none=True, + response_model=list[Vendor], ) async def find( db: Annotated[AsyncSession, Depends(get_db_session)], @@ -70,21 +73,28 @@ async def find( title="Offset", description="Offset in the full result set if a limit is set.", ), -) -> list[Vendor]: +) -> JSONResponse: sort_by: dict[str, SortOrder] = {} if sort is not None: for sort_item in sort.split(","): field, direction = sort_item.split(":") sort_by[field] = SortOrder[direction.upper()] - db_items = await vendor.find( + db_items, total_count = await vendor.find( db=db, name=name, sort_by=sort_by, limit=limit, offset=offset, ) - return [Vendor.from_db(db_item) for db_item in db_items] + # Set x-total-count header for pagination + return JSONResponse( + content=jsonable_encoder( + (Vendor.from_db(db_item) for db_item in db_items), + exclude_none=True, + ), + headers={"x-total-count": str(total_count)}, + ) @router.get( diff --git a/spoolman/database/filament.py b/spoolman/database/filament.py index 0f6b1e6..04c7cbb 100644 --- a/spoolman/database/filament.py +++ b/spoolman/database/filament.py @@ -3,7 +3,7 @@ import logging from typing import Optional -from sqlalchemy import select +from sqlalchemy import func, select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import contains_eager, joinedload @@ -78,11 +78,13 @@ async def find( # noqa: C901 sort_by: Optional[dict[str, SortOrder]] = None, limit: Optional[int] = None, offset: int = 0, -) -> list[models.Filament]: +) -> tuple[list[models.Filament], int]: """Find a list of filament objects by search criteria. Sort by a field by passing a dict with the field name as key and the sort order as value. The field name can contain nested fields, e.g. vendor.name. + + Returns a tuple containing the list of items and the total count of matching items. """ stmt = ( select(models.Filament) @@ -100,6 +102,14 @@ async def find( # noqa: C901 if article_number is not None: stmt = stmt.where(models.Filament.article_number.ilike(f"%{article_number}%")) + total_count = None + + if limit is not None: + total_count_stmt = stmt.with_only_columns(func.count()) + total_count = (await db.execute(total_count_stmt)).scalar() + + stmt = stmt.offset(offset).limit(limit) + if sort_by is not None: for fieldstr, order in sort_by.items(): field = parse_nested_field(models.Filament, fieldstr) @@ -108,11 +118,12 @@ async def find( # noqa: C901 elif order == SortOrder.DESC: stmt = stmt.order_by(field.desc()) - if limit is not None: - stmt = stmt.offset(offset).limit(limit) - rows = await db.execute(stmt) - return list(rows.scalars().all()) + result = list(rows.scalars().all()) + if total_count is None: + total_count = len(result) + + return result, total_count async def update( diff --git a/spoolman/database/spool.py b/spoolman/database/spool.py index 2d5fbd2..467b8fe 100644 --- a/spoolman/database/spool.py +++ b/spoolman/database/spool.py @@ -4,7 +4,7 @@ from datetime import datetime, timezone from typing import Optional import sqlalchemy -from sqlalchemy import case +from sqlalchemy import case, func from sqlalchemy.exc import NoResultFound from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import contains_eager, joinedload @@ -90,11 +90,13 @@ async def find( # noqa: C901, PLR0912 sort_by: Optional[dict[str, SortOrder]] = None, limit: Optional[int] = None, offset: int = 0, -) -> list[models.Spool]: +) -> tuple[list[models.Spool], int]: """Find a list of spool objects by search criteria. Sort by a field by passing a dict with the field name as key and the sort order as value. The field name can contain nested fields, e.g. filament.name. + + Returns a tuple containing the list of items and the total count of matching items. """ stmt = ( sqlalchemy.select(models.Spool) @@ -125,6 +127,14 @@ async def find( # noqa: C901, PLR0912 ), ) + total_count = None + + if limit is not None: + total_count_stmt = stmt.with_only_columns(func.count()) + total_count = (await db.execute(total_count_stmt)).scalar() + + stmt = stmt.offset(offset).limit(limit) + if sort_by is not None: for fieldstr, order in sort_by.items(): field = parse_nested_field(models.Spool, fieldstr) @@ -133,11 +143,12 @@ async def find( # noqa: C901, PLR0912 elif order == SortOrder.DESC: stmt = stmt.order_by(field.desc()) - if limit is not None: - stmt = stmt.offset(offset).limit(limit) - rows = await db.execute(stmt) - return list(rows.scalars().all()) + result = list(rows.scalars().all()) + if total_count is None: + total_count = len(result) + + return result, total_count async def update( diff --git a/spoolman/database/utils.py b/spoolman/database/utils.py index 44dae1f..5f95ad2 100644 --- a/spoolman/database/utils.py +++ b/spoolman/database/utils.py @@ -1,6 +1,7 @@ """Utility functions for the database module.""" from enum import Enum +from typing import Any import sqlalchemy @@ -12,7 +13,7 @@ class SortOrder(Enum): DESC = 2 -def parse_nested_field(base_obj: type[models.Base], field: str) -> sqlalchemy.Column: +def parse_nested_field(base_obj: type[models.Base], field: str) -> sqlalchemy.Column[Any]: """Parse a nested field string into a sqlalchemy field object.""" fields = field.split(".") if not hasattr(base_obj, fields[0]): diff --git a/spoolman/database/vendor.py b/spoolman/database/vendor.py index 8ecb2d3..a9e66a9 100644 --- a/spoolman/database/vendor.py +++ b/spoolman/database/vendor.py @@ -2,7 +2,7 @@ from typing import Optional -from sqlalchemy import select +from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession from spoolman.database import models @@ -41,12 +41,23 @@ async def find( sort_by: Optional[dict[str, SortOrder]] = None, limit: Optional[int] = None, offset: int = 0, -) -> list[models.Vendor]: - """Find a list of vendor objects by search criteria.""" +) -> tuple[list[models.Vendor], int]: + """Find a list of vendor objects by search criteria. + + Returns a tuple containing the list of items and the total count of matching items. + """ stmt = select(models.Vendor) if name is not None: stmt = stmt.where(models.Vendor.name.ilike(f"%{name}%")) + total_count = None + + if limit is not None: + total_count_stmt = stmt.with_only_columns(func.count()) + total_count = (await db.execute(total_count_stmt)).scalar() + + stmt = stmt.offset(offset).limit(limit) + if sort_by is not None: for fieldstr, order in sort_by.items(): field = getattr(models.Vendor, fieldstr) @@ -55,11 +66,12 @@ async def find( elif order == SortOrder.DESC: stmt = stmt.order_by(field.desc()) - if limit is not None: - stmt = stmt.offset(offset).limit(limit) - rows = await db.execute(stmt) - return list(rows.scalars().all()) + result = list(rows.scalars().all()) + if total_count is None: + total_count = len(result) + + return result, total_count async def update(