diff options
Diffstat (limited to 'django')
| -rw-r--r-- | django/contrib/postgres/aggregates/general.py | 10 | ||||
| -rw-r--r-- | django/contrib/postgres/aggregates/mixins.py | 47 |
2 files changed, 53 insertions, 4 deletions
diff --git a/django/contrib/postgres/aggregates/general.py b/django/contrib/postgres/aggregates/general.py index 806ecd1b78..4b2da0b101 100644 --- a/django/contrib/postgres/aggregates/general.py +++ b/django/contrib/postgres/aggregates/general.py @@ -1,14 +1,16 @@ from django.contrib.postgres.fields import ArrayField, JSONField from django.db.models.aggregates import Aggregate +from .mixins import OrderableAggMixin + __all__ = [ 'ArrayAgg', 'BitAnd', 'BitOr', 'BoolAnd', 'BoolOr', 'JSONBAgg', 'StringAgg', ] -class ArrayAgg(Aggregate): +class ArrayAgg(OrderableAggMixin, Aggregate): function = 'ARRAY_AGG' - template = '%(function)s(%(distinct)s%(expressions)s)' + template = '%(function)s(%(distinct)s%(expressions)s %(ordering)s)' @property def output_field(self): @@ -49,9 +51,9 @@ class JSONBAgg(Aggregate): return value -class StringAgg(Aggregate): +class StringAgg(OrderableAggMixin, Aggregate): function = 'STRING_AGG' - template = "%(function)s(%(distinct)s%(expressions)s, '%(delimiter)s')" + template = "%(function)s(%(distinct)s%(expressions)s, '%(delimiter)s'%(ordering)s)" def __init__(self, expression, delimiter, distinct=False, **extra): distinct = 'DISTINCT ' if distinct else '' diff --git a/django/contrib/postgres/aggregates/mixins.py b/django/contrib/postgres/aggregates/mixins.py new file mode 100644 index 0000000000..b270a5b653 --- /dev/null +++ b/django/contrib/postgres/aggregates/mixins.py @@ -0,0 +1,47 @@ +from django.db.models.expressions import F, OrderBy + + +class OrderableAggMixin: + + def __init__(self, expression, ordering=(), **extra): + if not isinstance(ordering, (list, tuple)): + ordering = [ordering] + ordering = ordering or [] + # Transform minus sign prefixed strings into an OrderBy() expression. + ordering = ( + (OrderBy(F(o[1:]), descending=True) if isinstance(o, str) and o[0] == '-' else o) + for o in ordering + ) + super().__init__(expression, **extra) + self.ordering = self._parse_expressions(*ordering) + + def resolve_expression(self, *args, **kwargs): + self.ordering = [expr.resolve_expression(*args, **kwargs) for expr in self.ordering] + return super().resolve_expression(*args, **kwargs) + + def as_sql(self, compiler, connection): + if self.ordering: + self.extra['ordering'] = 'ORDER BY ' + ', '.join(( + ordering_element.as_sql(compiler, connection)[0] + for ordering_element in self.ordering + )) + else: + self.extra['ordering'] = '' + return super().as_sql(compiler, connection) + + def get_source_expressions(self): + return self.source_expressions + self.ordering + + def get_source_fields(self): + # Filter out fields contributed by the ordering expressions as + # these should not be used to determine which the return type of the + # expression. + return [ + e._output_field_or_none + for e in self.get_source_expressions()[:self._get_ordering_expressions_index()] + ] + + def _get_ordering_expressions_index(self): + """Return the index at which the ordering expressions start.""" + source_expressions = self.get_source_expressions() + return len(source_expressions) - len(self.ordering) |
