wag-managment-api-service-v.../Services/PostgresDb/Models/response.py

91 lines
2.4 KiB
Python

"""
Response handler for PostgreSQL query results.
This module provides a wrapper class for SQLAlchemy query results,
adding convenience methods for accessing data and managing query state.
"""
from typing import Any, Dict, List, Optional, TypeVar, Generic, Union
from sqlalchemy.orm import Query
T = TypeVar("T")
class PostgresResponse(Generic[T]):
"""
Wrapper for PostgreSQL/SQLAlchemy query results.
Attributes:
query: SQLAlchemy query object
first: Whether to return first result only
data: Query results (lazy loaded)
count: Total count of results
Properties:
all: All results as list
first_item: First result only
"""
def __init__(
self,
query: Query,
first: bool = False,
status: bool = True,
message: str = "",
error: Optional[str] = None,
):
self._query = query
self._first = first
self.status = status
self.message = message
self.error = error
self._data: Optional[Union[List[T], T]] = None
self._count: Optional[int] = None
@property
def query(self) -> Query:
"""Get query object."""
return self._query
@property
def data(self) -> Union[List[T], T, None]:
"""
Lazy load and return query results.
Returns first item if first=True, otherwise returns all results.
"""
if self._data is None:
results = self._query.all()
self._data = results[0] if self._first and results else results
return self._data
@property
def count(self) -> int:
"""Lazy load and return total count of results."""
if self._count is None:
self._count = self._query.count()
return self._count
@property
def all(self) -> List[T]:
"""Get all results as list."""
return (
self.data
if isinstance(self.data, list)
else [self.data] if self.data else []
)
@property
def first(self) -> Optional[T]:
"""Get first result only."""
return self.data if self._first else (self.data[0] if self.data else None)
def as_dict(self) -> Dict[str, Any]:
"""Convert response to dictionary format."""
return {
"status": self.status,
"message": self.message,
"data": self.data,
"count": self.count,
"error": self.error,
}