prod-wag-backend-automate-services-2/api_services/api_validations/response/base.py

"""
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, Optional, TypeVar, Generic, Union

from pydantic import BaseModel
from sqlalchemy.orm import Query


T = TypeVar("T")


class PostgresResponse(Generic[T]):
    """
    Wrapper for PostgreSQL/SQLAlchemy query results.

    Properties:
        count: Total count of results
        query: Get query object
        as_dict: Convert response to dictionary format
    """

    def __init__(self, query: Query, base_model: Optional[BaseModel] = None):
        self._query = query
        self._count: Optional[int] = None
        self._base_model: Optional[BaseModel] = base_model
        self.single = False

    @property
    def query(self) -> Query:
        """Get query object."""
        return self._query

    @property
    def data(self) -> Union[list[T], T]:
        """Get query object."""
        return self._query.all()

    @property
    def count(self) -> int:
        """Get query object."""
        return self._query.count()

    @property
    def to_dict(self, **kwargs) -> list[dict]:
        """Get query object."""
        if self._base_model:
            return [self._base_model(**item.to_dict()).model_dump(**kwargs) for item in self.data]
        return [item.to_dict() for item in self.data]

    @property
    def as_dict(self) -> Dict[str, Any]:
        """Convert response to dictionary format."""
        return {
            "query": str(self.query),
            "count": self.count,
            "data": self.to_dict,
        }


class PostgresResponseSingle(Generic[T]):
    """
    Wrapper for PostgreSQL/SQLAlchemy query results.

    Properties:
        count: Total count of results
        query: Get query object
        as_dict: Convert response to dictionary format
        data: Get query object
    """

    def __init__(self, query: Query, base_model: Optional[BaseModel] = None):
        self._query = query
        self._count: Optional[int] = None
        self._base_model: Optional[BaseModel] = base_model
        self.single = True

    @property
    def query(self) -> Query:
        """Get query object."""
        return self._query

    @property
    def to_dict(self, **kwargs) -> dict:
        """Get query object."""
        if self._base_model:
            return self._base_model(**self._query.first().to_dict()).model_dump(**kwargs)
        return self._query.first().to_dict()

    @property
    def data(self) -> T:
        """Get query object."""
        return self._query.first()

    @property
    def count(self) -> int:
        """Get query object."""
        return self._query.count()

    @property
    def as_dict(self) -> Dict[str, Any]:
        """Convert response to dictionary format."""
        return {"query": str(self.query),"data": self.to_dict, "count": self.count}


class ResultQueryJoin:
    """
    ResultQueryJoin
    params:
        list_of_instrumented_attributes: list of instrumented attributes
        query: query object
    """

    def __init__(self, list_of_instrumented_attributes, query):
        """Initialize ResultQueryJoin"""
        self.list_of_instrumented_attributes = list_of_instrumented_attributes
        self._query = query

    @property
    def query(self):
        """Get query object."""
        return self._query

    @property
    def to_dict(self):
        """Convert response to dictionary format."""
        list_of_dictionaries, result = [], dict()
        for user_orders_shipping_iter in self.query.all():
            for index, instrumented_attribute_iter in enumerate(self.list_of_instrumented_attributes):
                result[str(instrumented_attribute_iter)] = user_orders_shipping_iter[index]
            list_of_dictionaries.append(result)
        return list_of_dictionaries

    @property
    def count(self):
        """Get count of query."""
        return self.query.count()

    @property
    def data(self):
        """Get query object."""
        return self.query.all()

    @property
    def as_dict(self):
        """Convert response to dictionary format."""
        return {"query": str(self.query), "data": self.data, "count": self.count}


class ResultQueryJoinSingle:
    """
    ResultQueryJoinSingle
    params:
        list_of_instrumented_attributes: list of instrumented attributes
        query: query object
    """

    def __init__(self, list_of_instrumented_attributes, query):
        """Initialize ResultQueryJoinSingle"""
        self.list_of_instrumented_attributes = list_of_instrumented_attributes
        self._query = query

    @property
    def query(self):
        """Get query object."""
        return self._query

    @property
    def to_dict(self):
        """Convert response to dictionary format."""
        data, result = self.query.first(), dict()
        for index, instrumented_attribute_iter in enumerate(self.list_of_instrumented_attributes):
            result[str(instrumented_attribute_iter)] = data[index]
        return result

    @property
    def count(self):
        """Get count of query."""
        return self.query.count()

    @property
    def data(self):
        """Get query object."""
        return self._query.first()

    @property
    def as_dict(self):
        """Convert response to dictionary format."""
        return {"query": str(self.query), "data": self.data, "count": self.count}