grocery-manager/app/data/util.py

#
# Copyright (c) Daniel Sheffield 2023
#
# All rights reserved
#
# THIS SOFTWARE IS PROVIDED AS IS WITHOUT WARRANTY
from typing import List, Set, Tuple, Iterable
from psycopg.sql import (
    Identifier,
    Literal,
    SQL,
    Composable,
)


def get_where_include_exclude(
        table: str,
        col: str,
        include: Iterable[str],
        exclude: Iterable[str]
) -> SQL:
    return SQL("""
  ({identifier} = ANY({include}) OR ARRAY[]::text[] @> {include}::text[])
AND
  (NOT {identifier} = ANY({exclude}) OR {identifier} IS NULL)
""").format(
        identifier=Identifier(table, col),
        include=Literal(include),
        exclude=Literal(exclude)
    )


def get_select(alias_to_sql: dict[str,Composable]) -> Composable:
    select = SQL(""",
    """).join([
        SQL(' ').join([
            v, SQL('AS'), Identifier(k)
        ]) for k, v in alias_to_sql.items()
    ])
    return SQL("""
    """).join([SQL("SELECT"), *select])


def get_from(
    base: str,
    table_to_join_on: dict[Tuple[str, Tuple[str,str]]]
) -> Composable:
    joins = [
        SQL("{table} ON {table_column} = {other_column}").format(
            table=Identifier(table),
            table_column=Identifier(table, table_column),
            other_column=Identifier(*other_column) if isinstance(other_column,tuple) else Identifier(other_column)
        ) for table, (table_column, other_column) in table_to_join_on.items()
    ]
    return SQL('').join([SQL("""FROM {base}
LEFT JOIN """).format(base=Identifier(base)),
        SQL("""
LEFT JOIN """).join(joins)])


def get_groupby(alias_to_sql: dict[str, Composable], formatter=None) -> Composable:
    groupby = SQL(""",
""").join([
        formatter(v) if formatter is not None and isinstance(
            v, SQL
        ) else v for k, v in alias_to_sql.items()
    ])
    return SQL("""
""").join([SQL("GROUP BY"), *groupby])