286 lines
8.0 KiB
Python
286 lines
8.0 KiB
Python
"""
|
|
Advanced filtering functionality for SQLAlchemy models.
|
|
|
|
This module provides a comprehensive set of filtering capabilities for SQLAlchemy models,
|
|
including pagination, ordering, and complex query building.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
import arrow
|
|
|
|
from typing import Any, TypeVar, Type, Union, Optional
|
|
|
|
from sqlalchemy import ColumnExpressionArgument
|
|
from sqlalchemy.orm import Query, Session
|
|
from sqlalchemy.sql.elements import BinaryExpression
|
|
|
|
from response import PostgresResponse
|
|
|
|
|
|
T = TypeVar("T", bound="QueryModel")
|
|
|
|
|
|
class QueryModel:
|
|
|
|
__abstract__ = True
|
|
pre_query = None
|
|
|
|
@classmethod
|
|
def _query(cls: Type[T], db: Session) -> Query:
|
|
"""Returns the query to use in the model."""
|
|
return cls.pre_query if cls.pre_query else db.query(cls)
|
|
|
|
@classmethod
|
|
def add_new_arg_to_args(
|
|
cls: Type[T],
|
|
args_list: tuple[BinaryExpression, ...],
|
|
argument: str,
|
|
value: BinaryExpression
|
|
) -> tuple[BinaryExpression, ...]:
|
|
"""
|
|
Add a new argument to the query arguments if it doesn't exist.
|
|
|
|
Args:
|
|
args_list: Existing query arguments
|
|
argument: Key of the argument to check for
|
|
value: New argument value to add
|
|
|
|
Returns:
|
|
Updated tuple of query arguments
|
|
"""
|
|
# Convert to set to remove duplicates while preserving order
|
|
new_args = list(dict.fromkeys(
|
|
arg for arg in args_list
|
|
if isinstance(arg, BinaryExpression)
|
|
))
|
|
|
|
# Check if argument already exists
|
|
if not any(
|
|
getattr(getattr(arg, "left", None), "key", None) == argument
|
|
for arg in new_args
|
|
):
|
|
new_args.append(value)
|
|
|
|
return tuple(new_args)
|
|
|
|
@classmethod
|
|
def get_not_expired_query_arg(
|
|
cls: Type[T],
|
|
args: tuple[BinaryExpression, ...]
|
|
) -> tuple[BinaryExpression, ...]:
|
|
"""
|
|
Add expiry date filtering to the query arguments.
|
|
|
|
Args:
|
|
args: Existing query arguments
|
|
|
|
Returns:
|
|
Updated tuple of query arguments with expiry filters
|
|
|
|
Raises:
|
|
AttributeError: If model does not have expiry_starts or expiry_ends columns
|
|
"""
|
|
try:
|
|
current_time = str(arrow.now())
|
|
starts = cls.expiry_starts <= current_time
|
|
ends = cls.expiry_ends > current_time
|
|
|
|
args = cls.add_new_arg_to_args(args, "expiry_ends", ends)
|
|
args = cls.add_new_arg_to_args(args, "expiry_starts", starts)
|
|
return args
|
|
|
|
except AttributeError as e:
|
|
raise AttributeError(
|
|
f"Model {cls.__name__} must have expiry_starts and expiry_ends columns"
|
|
) from e
|
|
|
|
@classmethod
|
|
def produce_query_to_add(cls: Type[T], filter_list, args):
|
|
"""
|
|
Adds query to main filter options
|
|
Args:
|
|
filter_list: Dictionary containing query parameters
|
|
args: Existing query arguments to add to
|
|
|
|
Returns:
|
|
Updated query arguments tuple
|
|
"""
|
|
if filter_list.get("query"):
|
|
for smart_iter in cls.filter_expr(**filter_list["query"]):
|
|
if key := getattr(getattr(smart_iter, "left", None), "key", None):
|
|
args = cls.add_new_arg_to_args(args, key, smart_iter)
|
|
return args
|
|
|
|
@classmethod
|
|
def convert(
|
|
cls: Type[T],
|
|
smart_options: dict[str, Any],
|
|
validate_model: Any = None
|
|
) -> Optional[tuple[BinaryExpression, ...]]:
|
|
"""
|
|
Convert smart options to SQLAlchemy filter expressions.
|
|
|
|
Args:
|
|
smart_options: Dictionary of filter options
|
|
validate_model: Optional model to validate against
|
|
|
|
Returns:
|
|
Tuple of SQLAlchemy filter expressions or None if validation fails
|
|
"""
|
|
if validate_model is not None:
|
|
# Add validation logic here if needed
|
|
pass
|
|
|
|
return tuple(cls.filter_expr(**smart_options))
|
|
|
|
@classmethod
|
|
def filter_by_one(
|
|
cls: Type[T],
|
|
db: Session,
|
|
system: bool = False,
|
|
**kwargs: Any
|
|
) -> PostgresResponse[T]:
|
|
"""
|
|
Filter single record by keyword arguments.
|
|
|
|
Args:
|
|
db: Database session
|
|
system: If True, skip status filtering
|
|
**kwargs: Filter criteria
|
|
|
|
Returns:
|
|
Query response with single record
|
|
"""
|
|
if "is_confirmed" not in kwargs and not system:
|
|
kwargs["is_confirmed"] = True
|
|
kwargs.pop("system", None)
|
|
query = cls._query(db).filter_by(**kwargs)
|
|
return PostgresResponse(
|
|
model=cls,
|
|
pre_query=cls._query(db),
|
|
query=query,
|
|
is_array=False
|
|
)
|
|
|
|
@classmethod
|
|
def filter_one(
|
|
cls: Type[T],
|
|
*args: Union[BinaryExpression, ColumnExpressionArgument],
|
|
db: Session,
|
|
) -> PostgresResponse[T]:
|
|
"""
|
|
Filter single record by expressions.
|
|
|
|
Args:
|
|
db: Database session
|
|
*args: Filter expressions
|
|
|
|
Returns:
|
|
Query response with single record
|
|
"""
|
|
args = cls.get_not_expired_query_arg(args)
|
|
query = cls._query(db).filter(*args)
|
|
return PostgresResponse(
|
|
model=cls,
|
|
pre_query=cls._query(db),
|
|
query=query,
|
|
is_array=False
|
|
)
|
|
|
|
@classmethod
|
|
def filter_one_system(
|
|
cls: Type[T],
|
|
*args: Union[BinaryExpression, ColumnExpressionArgument],
|
|
db: Session,
|
|
) -> PostgresResponse[T]:
|
|
"""
|
|
Filter single record by expressions without status filtering.
|
|
|
|
Args:
|
|
db: Database session
|
|
*args: Filter expressions
|
|
|
|
Returns:
|
|
Query response with single record
|
|
"""
|
|
query = cls._query(db).filter(*args)
|
|
return PostgresResponse(
|
|
model=cls,
|
|
pre_query=cls._query(db),
|
|
query=query,
|
|
is_array=False
|
|
)
|
|
|
|
@classmethod
|
|
def filter_all_system(
|
|
cls: Type[T],
|
|
*args: Union[BinaryExpression, ColumnExpressionArgument],
|
|
db: Session,
|
|
) -> PostgresResponse[T]:
|
|
"""
|
|
Filter multiple records by expressions without status filtering.
|
|
|
|
Args:
|
|
db: Database session
|
|
*args: Filter expressions
|
|
|
|
Returns:
|
|
Query response with matching records
|
|
"""
|
|
query = cls._query(db).filter(*args)
|
|
return PostgresResponse(
|
|
model=cls,
|
|
pre_query=cls._query(db),
|
|
query=query,
|
|
is_array=True
|
|
)
|
|
|
|
@classmethod
|
|
def filter_all(
|
|
cls: Type[T],
|
|
*args: Union[BinaryExpression, ColumnExpressionArgument],
|
|
db: Session,
|
|
) -> PostgresResponse[T]:
|
|
"""
|
|
Filter multiple records by expressions.
|
|
|
|
Args:
|
|
db: Database session
|
|
*args: Filter expressions
|
|
|
|
Returns:
|
|
Query response with matching records
|
|
"""
|
|
args = cls.get_not_expired_query_arg(args)
|
|
query = cls._query(db).filter(*args)
|
|
return PostgresResponse(
|
|
model=cls,
|
|
pre_query=cls._query(db),
|
|
query=query,
|
|
is_array=True
|
|
)
|
|
|
|
@classmethod
|
|
def filter_by_all_system(
|
|
cls: Type[T],
|
|
db: Session,
|
|
**kwargs: Any
|
|
) -> PostgresResponse[T]:
|
|
"""
|
|
Filter multiple records by keyword arguments.
|
|
|
|
Args:
|
|
db: Database session
|
|
**kwargs: Filter criteria
|
|
|
|
Returns:
|
|
Query response with matching records
|
|
"""
|
|
query = cls._query(db).filter_by(**kwargs)
|
|
return PostgresResponse(
|
|
model=cls,
|
|
pre_query=cls._query(db),
|
|
query=query,
|
|
is_array=True
|
|
)
|