where_builder.py

from typing import Union, Any, List, Tuple, Dict
import copy

_ExpressionTuple = Tuple[str, str, Any]
_Expression = Union[str, _ExpressionTuple]
_ExpressionOrBuilder = Union[_Expression, "WhereClauseBuilder"]


class WhereClauseBuilder:
    """
    Simple SQL WHERE expression builder.
    Returns (sql, params) tuple which is consumed by asyncpg.
    Usage:
        >>> where_clause, params = (
        ...    WhereClauseBuilder()
        ...   .add_AND(("age", ">=", 18))
        ...   .add_OR(
        ...      WhereClauseBuilder()
        ...      .add_OR("is_admin IS TRUE")
        ...      .add_OR(("role", "=", "admin"))
        ...    )
        ...    .build()
        ... )
        >>> where_clause
        'age >= $1 OR (is_admin = $2 OR role = $3)'
        >>> params
        [18, True, 'admin']
    """

    def __init__(self):
        self.conditions: List[Tuple[str, _ExpressionOrBuilder]] = []
        self.params: Dict[str, Any] = {}

    def add_AND(self, expression: _ExpressionOrBuilder) -> "WhereClauseBuilder":
        builder = copy.deepcopy(self)
        builder.conditions.append(("AND", expression))
        return builder

    def add_OR(self, expression: _ExpressionOrBuilder) -> "WhereClauseBuilder":
        builder = copy.deepcopy(self)
        builder.conditions.append(("OR", expression))
        return builder

    def build(self, param_offset: int = 0) -> Tuple[str, list]:
        where_clause = ""
        params = []

        for concat_operator, expression in self.conditions:
            if isinstance(expression, str):
                if where_clause:
                    where_clause += f" {concat_operator} "
                where_clause += f"{expression}"

            elif isinstance(expression, tuple):
                if where_clause:
                    where_clause += f" {concat_operator} "

                field, expr_operator, value = expression
                params.append(value)
                number = len(params) + param_offset
                where_clause += f"{field} {expr_operator} ${number}"

            elif isinstance(expression, WhereClauseBuilder):
                inner_where, inner_params = expression.build(
                    param_offset=len(params) + param_offset
                )
                if where_clause:
                    where_clause += f" {concat_operator} "
                where_clause += f"({inner_where})"
                params += inner_params

        return where_clause, params