From a06b977a91f043c509df781670fb4cf35cb437b7 Mon Sep 17 00:00:00 2001 From: Hannes Ljungberg Date: Fri, 20 Nov 2020 21:35:04 +0100 Subject: [PATCH] Fixed #32776 -- Added support for Array subqueries on PostgreSQL. --- django/contrib/postgres/expressions.py | 14 +++++ docs/ref/contrib/postgres/expressions.txt | 37 +++++++++++++ docs/ref/contrib/postgres/index.txt | 1 + docs/releases/4.0.txt | 5 ++ tests/postgres_tests/test_array.py | 63 ++++++++++++++++++++++- 5 files changed, 119 insertions(+), 1 deletion(-) create mode 100644 django/contrib/postgres/expressions.py create mode 100644 docs/ref/contrib/postgres/expressions.txt diff --git a/django/contrib/postgres/expressions.py b/django/contrib/postgres/expressions.py new file mode 100644 index 0000000000..ea7cbe038d --- /dev/null +++ b/django/contrib/postgres/expressions.py @@ -0,0 +1,14 @@ +from django.contrib.postgres.fields import ArrayField +from django.db.models import Subquery +from django.utils.functional import cached_property + + +class ArraySubquery(Subquery): + template = 'ARRAY(%(subquery)s)' + + def __init__(self, queryset, **kwargs): + super().__init__(queryset, **kwargs) + + @cached_property + def output_field(self): + return ArrayField(self.query.output_field) diff --git a/docs/ref/contrib/postgres/expressions.txt b/docs/ref/contrib/postgres/expressions.txt new file mode 100644 index 0000000000..79590c2393 --- /dev/null +++ b/docs/ref/contrib/postgres/expressions.txt @@ -0,0 +1,37 @@ +===================================== +PostgreSQL specific query expressions +===================================== + +.. module:: django.contrib.postgres.expressions + :synopsis: PostgreSQL specific query expressions + +These expressions are available from the +``django.contrib.postgres.expressions`` module. + +``ArraySubquery()`` expressions +=============================== + +.. class:: ArraySubquery(queryset) + +.. versionadded:: 4.0 + +``ArraySubquery`` is a :class:`~django.db.models.Subquery` that uses the +PostgreSQL ``ARRAY`` constructor to build a list of values from the queryset, +which must use :meth:`.QuerySet.values` to return only a single column. + +This class differs from :class:`~django.contrib.postgres.aggregates.ArrayAgg` +in the way that it does not act as an aggregate function and does not require +an SQL ``GROUP BY`` clause to build the list of values. + +For example, if you want to annotate all related books to an author as JSON +objects:: + + >>> from django.db.models import OuterRef + >>> from django.db.models.functions import JSONObject + >>> from django.contrib.postgres.expressions import ArraySubquery + >>> books = Book.objects.filter(author=OuterRef('pk')).values( + ... json=JSONObject(title='title', pages='pages') + ... ) + >>> author = Author.objects.annotate(books=ArraySubquery(books)).first() + >>> author.books + [{'title': 'Solaris', 'pages': 204}, {'title': 'The Cyberiad', 'pages': 295}] diff --git a/docs/ref/contrib/postgres/index.txt b/docs/ref/contrib/postgres/index.txt index 03ff6da1e0..7627043037 100644 --- a/docs/ref/contrib/postgres/index.txt +++ b/docs/ref/contrib/postgres/index.txt @@ -30,6 +30,7 @@ a number of PostgreSQL specific data types. aggregates constraints + expressions fields forms functions diff --git a/docs/releases/4.0.txt b/docs/releases/4.0.txt index 54d0fd4c1d..bc8e1639c8 100644 --- a/docs/releases/4.0.txt +++ b/docs/releases/4.0.txt @@ -131,6 +131,11 @@ Minor features :class:`~django.contrib.postgres.operations.AddConstraintNotValid` on PostgreSQL. +* The new + :class:`ArraySubquery() ` + expression allows using subqueries to construct lists of values on + PostgreSQL. + :mod:`django.contrib.redirects` ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/tests/postgres_tests/test_array.py b/tests/postgres_tests/test_array.py index 6fb34b3d92..62a495d5a7 100644 --- a/tests/postgres_tests/test_array.py +++ b/tests/postgres_tests/test_array.py @@ -10,7 +10,7 @@ from django.core.exceptions import FieldError from django.core.management import call_command from django.db import IntegrityError, connection, models from django.db.models.expressions import Exists, OuterRef, RawSQL, Value -from django.db.models.functions import Cast, Upper +from django.db.models.functions import Cast, JSONObject, Upper from django.test import TransactionTestCase, modify_settings, override_settings from django.test.utils import isolate_apps from django.utils import timezone @@ -28,6 +28,7 @@ try: from psycopg2.extras import NumericRange from django.contrib.postgres.aggregates import ArrayAgg + from django.contrib.postgres.expressions import ArraySubquery from django.contrib.postgres.fields import ArrayField from django.contrib.postgres.fields.array import ( IndexTransform, SliceTransform, @@ -551,6 +552,66 @@ class TestQuerying(PostgreSQLTestCase): 1, ) + def test_filter_by_array_subquery(self): + inner_qs = NullableIntegerArrayModel.objects.filter( + field__len=models.OuterRef('field__len'), + ).values('field') + self.assertSequenceEqual( + NullableIntegerArrayModel.objects.alias( + same_sized_fields=ArraySubquery(inner_qs), + ).filter(same_sized_fields__len__gt=1), + self.objs[0:2], + ) + + def test_annotated_array_subquery(self): + inner_qs = NullableIntegerArrayModel.objects.exclude( + pk=models.OuterRef('pk') + ).values('order') + self.assertSequenceEqual( + NullableIntegerArrayModel.objects.annotate( + sibling_ids=ArraySubquery(inner_qs), + ).get(order=1).sibling_ids, + [2, 3, 4, 5], + ) + + def test_group_by_with_annotated_array_subquery(self): + inner_qs = NullableIntegerArrayModel.objects.exclude( + pk=models.OuterRef('pk') + ).values('order') + self.assertSequenceEqual( + NullableIntegerArrayModel.objects.annotate( + sibling_ids=ArraySubquery(inner_qs), + sibling_count=models.Max('sibling_ids__len'), + ).values_list('sibling_count', flat=True), + [len(self.objs) - 1] * len(self.objs), + ) + + def test_annotated_ordered_array_subquery(self): + inner_qs = NullableIntegerArrayModel.objects.order_by('-order').values('order') + self.assertSequenceEqual( + NullableIntegerArrayModel.objects.annotate( + ids=ArraySubquery(inner_qs), + ).first().ids, + [5, 4, 3, 2, 1], + ) + + def test_annotated_array_subquery_with_json_objects(self): + inner_qs = NullableIntegerArrayModel.objects.exclude( + pk=models.OuterRef('pk') + ).values(json=JSONObject(order='order', field='field')) + siblings_json = NullableIntegerArrayModel.objects.annotate( + siblings_json=ArraySubquery(inner_qs), + ).values_list('siblings_json', flat=True).get(order=1) + self.assertSequenceEqual( + siblings_json, + [ + {'field': [2], 'order': 2}, + {'field': [2, 3], 'order': 3}, + {'field': [20, 30, 40], 'order': 4}, + {'field': None, 'order': 5}, + ], + ) + class TestDateTimeExactQuerying(PostgreSQLTestCase):