44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
from typing import Any, List, Optional, TypeVar, Union
|
|
|
|
from sqlalchemy.orm import Query
|
|
from sqlalchemy.orm.session import Session
|
|
|
|
T = TypeVar('T')
|
|
|
|
class QueryResponse:
|
|
"""Handler for SQLAlchemy query results with error handling."""
|
|
|
|
def __init__(self, db: Session, query: Query, first: bool = False):
|
|
self.db = db
|
|
self.first = first
|
|
self.__query = query
|
|
|
|
def get(self, index: int) -> Optional[T]:
|
|
"""Get item at specific index if it exists."""
|
|
count = self.count
|
|
if count and not index > count:
|
|
return self.data[index - 1]
|
|
return None
|
|
|
|
@property
|
|
def data(self) -> Union[Optional[T], List[T]]:
|
|
"""Get query results with error handling."""
|
|
try:
|
|
if self.first:
|
|
return self.__query.first()
|
|
return self.__query.all()
|
|
except Exception as e:
|
|
# Handle any database errors by rolling back
|
|
self.db.rollback()
|
|
return None if self.first else []
|
|
|
|
@property
|
|
def count(self) -> int:
|
|
"""Get total count of query results."""
|
|
return self.__query.count()
|
|
|
|
@property
|
|
def query(self) -> Query:
|
|
"""Get the underlying SQLAlchemy query."""
|
|
return self.__query
|