From e03083917db03757e48f8edac4c8491b72c8a3c4 Mon Sep 17 00:00:00 2001 From: Devin Cox Date: Fri, 9 Aug 2024 13:56:56 -0700 Subject: [PATCH] Fixed #35586 -- Added support for set-returning database functions. Aggregation optimization didn't account for not referenced set-returning annotations on Postgres. Co-authored-by: Simon Charette --- django/db/models/expressions.py | 2 ++ django/db/models/sql/query.py | 9 +++++++++ docs/ref/models/expressions.txt | 10 ++++++++++ docs/releases/5.2.txt | 4 ++++ tests/annotations/models.py | 8 ++++++++ tests/annotations/tests.py | 21 +++++++++++++++++++++ 6 files changed, 54 insertions(+) diff --git a/django/db/models/expressions.py b/django/db/models/expressions.py index dc09e43fda1..eee0eafc837 100644 --- a/django/db/models/expressions.py +++ b/django/db/models/expressions.py @@ -182,6 +182,8 @@ class BaseExpression: allowed_default = False # Can the expression be used during a constraint validation? constraint_validation_compatible = True + # Does the expression possibly return more than one row? + set_returning = False def __init__(self, output_field=None): if output_field is not None: diff --git a/django/db/models/sql/query.py b/django/db/models/sql/query.py index c1e2fc1d4f1..aef3f48f10a 100644 --- a/django/db/models/sql/query.py +++ b/django/db/models/sql/query.py @@ -491,6 +491,11 @@ class Query(BaseExpression): ) or having ) + set_returning_annotations = { + alias + for alias, annotation in self.annotation_select.items() + if getattr(annotation, "set_returning", False) + } # Decide if we need to use a subquery. # # Existing aggregations would cause incorrect results as @@ -510,6 +515,7 @@ class Query(BaseExpression): or qualify or self.distinct or self.combinator + or set_returning_annotations ): from django.db.models.sql.subqueries import AggregateQuery @@ -551,6 +557,9 @@ class Query(BaseExpression): if annotation.get_group_by_cols(): annotation_mask.add(annotation_alias) inner_query.set_annotation_mask(annotation_mask) + # Annotations that possibly return multiple rows cannot + # be masked as they might have an incidence on the query. + annotation_mask |= set_returning_annotations # Add aggregates to the outer AggregateQuery. This requires making # sure all columns referenced by the aggregates are selected in the diff --git a/docs/ref/models/expressions.txt b/docs/ref/models/expressions.txt index 1b6a208d016..7833580ac13 100644 --- a/docs/ref/models/expressions.txt +++ b/docs/ref/models/expressions.txt @@ -1095,6 +1095,16 @@ calling the appropriate methods on the wrapped expression. :py:data:`NotImplemented` which forces the expression to be computed on the database. + .. attribute:: set_returning + + .. versionadded:: 5.2 + + Tells Django that this expression contains a set-returning function, + enforcing subquery evaluation. It's used, for example, to allow some + Postgres set-returning functions (e.g. ``JSONB_PATH_QUERY``, + ``UNNEST``, etc.) to skip optimization and be properly evaluated when + annotations spawn rows themselves. Defaults to ``False``. + .. method:: resolve_expression(query=None, allow_joins=True, reuse=None, summarize=False, for_save=False) Provides the chance to do any preprocessing or validation of diff --git a/docs/releases/5.2.txt b/docs/releases/5.2.txt index 02a068e5af6..a15e6692053 100644 --- a/docs/releases/5.2.txt +++ b/docs/releases/5.2.txt @@ -218,6 +218,10 @@ Models * Added support for validation of model constraints which use a :class:`~django.db.models.GeneratedField`. +* The new :attr:`.Expression.set_returning` attribute specifies that the + expression contains a set-returning function, enforcing subquery evaluation. + This is necessary for many Postgres set-returning functions. + Requests and Responses ~~~~~~~~~~~~~~~~~~~~~~ diff --git a/tests/annotations/models.py b/tests/annotations/models.py index fbb9ca69884..914770d2fe7 100644 --- a/tests/annotations/models.py +++ b/tests/annotations/models.py @@ -58,3 +58,11 @@ class Company(models.Model): class Ticket(models.Model): active_at = models.DateTimeField() duration = models.DurationField() + + +class JsonModel(models.Model): + data = models.JSONField(default=dict, blank=True) + id = models.IntegerField(primary_key=True) + + class Meta: + required_db_features = {"supports_json_field"} diff --git a/tests/annotations/tests.py b/tests/annotations/tests.py index 703847e1dd5..29660a827e2 100644 --- a/tests/annotations/tests.py +++ b/tests/annotations/tests.py @@ -1,7 +1,9 @@ import datetime from decimal import Decimal +from unittest import skipUnless from django.core.exceptions import FieldDoesNotExist, FieldError +from django.db import connection from django.db.models import ( BooleanField, Case, @@ -15,6 +17,7 @@ from django.db.models import ( FloatField, Func, IntegerField, + JSONField, Max, OuterRef, Q, @@ -43,6 +46,7 @@ from .models import ( Company, DepartmentStore, Employee, + JsonModel, Publisher, Store, Ticket, @@ -1167,6 +1171,23 @@ class NonAggregateAnnotationTestCase(TestCase): with self.assertRaisesMessage(ValueError, msg): Book.objects.annotate(**{crafted_alias: Value(1)}) + @skipUnless(connection.vendor == "postgresql", "PostgreSQL tests") + @skipUnlessDBFeature("supports_json_field") + def test_set_returning_functions(self): + class JSONBPathQuery(Func): + function = "jsonb_path_query" + output_field = JSONField() + set_returning = True + + test_model = JsonModel.objects.create( + data={"key": [{"id": 1, "name": "test1"}, {"id": 2, "name": "test2"}]}, id=1 + ) + qs = JsonModel.objects.annotate( + table_element=JSONBPathQuery("data", Value("$.key[*]")) + ).filter(pk=test_model.pk) + + self.assertEqual(qs.count(), len(qs)) + class AliasTests(TestCase): @classmethod