diff --git a/src/infi/clickhouse_orm/query.py b/src/infi/clickhouse_orm/query.py index 897d45a..78da81e 100644 --- a/src/infi/clickhouse_orm/query.py +++ b/src/infi/clickhouse_orm/query.py @@ -297,6 +297,7 @@ def __init__(self, model_cls, database): self._order_by = [] self._where_q = Q() self._prewhere_q = Q() + self._having_q = Q() self._grouping_fields = [] self._grouping_with_totals = False self._fields = model_cls.fields().keys() @@ -392,6 +393,9 @@ def as_sql(self): if self._grouping_with_totals: sql += ' WITH TOTALS' + if self._having_q and not self._having_q.is_empty: + sql += '\nHAVING ' + self._having_q.to_sql(self._model_cls) + if self._order_by: sql += '\nORDER BY ' + self.order_by_as_sql() @@ -581,6 +585,15 @@ def _verify_mutation_allowed(self): assert not self._distinct, 'Mutations are not allowed after calling distinct()' assert not self._final, 'Mutations are not allowed after calling final()' + def _having_or_exclude(self, *q, **kwargs): + raise NotImplementedError('Cannot use `HAVING` with QuerySet, use AggregateQuerySet instead') + + def having(self, *args, **kwargs): + return self._having_or_exclude(*args, **kwargs) + + def having_exclude(self, *args, **kwargs): + return self._having_or_exclude(*args, _inverse=True, **kwargs) + def aggregate(self, *args, **kwargs): """ Returns an `AggregateQuerySet` over this query, with `args` serving as @@ -684,6 +697,34 @@ def with_totals(self): def _verify_mutation_allowed(self): raise AssertionError('Cannot mutate an AggregateQuerySet') + def _having_or_exclude(self, *q, **kwargs): + from .funcs import F + + inverse = kwargs.pop('_inverse', False) + + qs = copy(self) + + condition = Q() + for arg in q: + if isinstance(arg, Q): + condition &= arg + elif isinstance(arg, F): + condition &= Q(arg) + else: + raise TypeError('Invalid argument "%r" to queryset filter' % arg) + + if kwargs: + condition &= Q(**kwargs) + + if inverse: + condition = ~condition + if self._having_q: + condition = copy(self._having_q) & condition + + qs._having_q = condition + + return qs + # Expose only relevant classes in import * __all__ = [c.__name__ for c in [Q, QuerySet, AggregateQuerySet]]