from itertools import islice from typing import List, Union from sqlalchemy import text, and_ from sqlalchemy.sql.expression import Select, FunctionElement from sqlalchemy.types import Numeric from sqlalchemy.ext.compiler import compiles try: from sqlalchemy.dialects.postgresql import insert as pg_insert # pylint: disable=unused-import except ImportError: pg_insert = None from .tables import AccountAddress class greatest(FunctionElement): # pylint: disable=invalid-name type = Numeric() name = 'greatest' @compiles(greatest) def default_greatest(element, compiler, **kw): return "greatest(%s)" % compiler.process(element.clauses, **kw) @compiles(greatest, 'sqlite') def sqlite_greatest(element, compiler, **kw): return "max(%s)" % compiler.process(element.clauses, **kw) def chunk(rows, step): it, total = iter(rows), len(rows) for _ in range(0, total, step): yield list(islice(it, step)) total -= step def constrain_single_or_list(constraints, column, value, convert=lambda x: x): if value is not None: if isinstance(value, list): value = [convert(v) for v in value] if len(value) == 1: constraints[column] = value[0] elif len(value) > 1: constraints[f"{column}__in"] = value else: constraints[column] = convert(value) return constraints def in_account_ids(account_ids: Union[List[str], str]): if isinstance(account_ids, list): if len(account_ids) > 1: return AccountAddress.c.account.in_(account_ids) account_ids = account_ids[0] return AccountAddress.c.account == account_ids def query(table, s: Select, **constraints) -> Select: limit = constraints.pop('limit', None) if limit is not None: s = s.limit(limit) offset = constraints.pop('offset', None) if offset is not None: s = s.offset(offset) order_by = constraints.pop('order_by', None) if order_by: if isinstance(order_by, str): s = s.order_by(text(order_by)) elif isinstance(order_by, list): s = s.order_by(text(', '.join(order_by))) else: raise ValueError("order_by must be string or list") group_by = constraints.pop('group_by', None) if group_by is not None: s = s.group_by(text(group_by)) account_ids = constraints.pop('account_ids', []) if account_ids: s = s.where(in_account_ids(account_ids)) if constraints: s = s.where( constraints_to_clause(table, constraints) ) return s def constraints_to_clause(tables, constraints): clause = [] for key, constraint in constraints.items(): if key.endswith('__not'): col, op = key[:-len('__not')], '__ne__' elif key.endswith('__is_null'): col = key[:-len('__is_null')] op = '__eq__' constraint = None elif key.endswith('__is_not_null'): col = key[:-len('__is_not_null')] op = '__ne__' constraint = None elif key.endswith('__lt'): col, op = key[:-len('__lt')], '__lt__' elif key.endswith('__lte'): col, op = key[:-len('__lte')], '__le__' elif key.endswith('__gt'): col, op = key[:-len('__gt')], '__gt__' elif key.endswith('__gte'): col, op = key[:-len('__gte')], '__ge__' elif key.endswith('__like'): col, op = key[:-len('__like')], 'like' elif key.endswith('__not_like'): col, op = key[:-len('__not_like')], 'notlike' elif key.endswith('__in') or key.endswith('__not_in'): if key.endswith('__in'): col, op, one_val_op = key[:-len('__in')], 'in_', '__eq__' else: col, op, one_val_op = key[:-len('__not_in')], 'notin_', '__ne__' if isinstance(constraint, Select): pass elif constraint: if isinstance(constraint, (list, set, tuple)): if len(constraint) == 1: op = one_val_op constraint = next(iter(constraint)) elif isinstance(constraint, str): constraint = text(constraint) else: raise ValueError(f"{col} requires a list, set or string as constraint value.") else: continue else: col, op = key, '__eq__' attr = None if '.' in col: table_name, col = col.split('.') _table = None for table in tables: if table.name == table_name.lower(): _table = table break if _table is not None: attr = getattr(_table.c, col) else: raise ValueError(f"Table '{table_name}' not available: {', '.join([t.name for t in tables])}.") else: for table in tables: attr = getattr(table.c, col, None) if attr is not None: break if attr is None: raise ValueError(f"Attribute '{col}' not found on tables: {', '.join([t.name for t in tables])}.") clause.append(getattr(attr, op)(constraint)) return and_(*clause)