Fixed #32776 -- Added support for Array subqueries on PostgreSQL.

This commit is contained in:
Hannes Ljungberg 2020-11-20 21:35:04 +01:00 committed by Mariusz Felisiak
parent 49ca6bbc44
commit a06b977a91
5 changed files with 119 additions and 1 deletions

View File

@ -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)

View File

@ -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}]

View File

@ -30,6 +30,7 @@ a number of PostgreSQL specific data types.
aggregates aggregates
constraints constraints
expressions
fields fields
forms forms
functions functions

View File

@ -131,6 +131,11 @@ Minor features
:class:`~django.contrib.postgres.operations.AddConstraintNotValid` on :class:`~django.contrib.postgres.operations.AddConstraintNotValid` on
PostgreSQL. PostgreSQL.
* The new
:class:`ArraySubquery() <django.contrib.postgres.expressions.ArraySubquery>`
expression allows using subqueries to construct lists of values on
PostgreSQL.
:mod:`django.contrib.redirects` :mod:`django.contrib.redirects`
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

View File

@ -10,7 +10,7 @@ from django.core.exceptions import FieldError
from django.core.management import call_command from django.core.management import call_command
from django.db import IntegrityError, connection, models from django.db import IntegrityError, connection, models
from django.db.models.expressions import Exists, OuterRef, RawSQL, Value 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 import TransactionTestCase, modify_settings, override_settings
from django.test.utils import isolate_apps from django.test.utils import isolate_apps
from django.utils import timezone from django.utils import timezone
@ -28,6 +28,7 @@ try:
from psycopg2.extras import NumericRange from psycopg2.extras import NumericRange
from django.contrib.postgres.aggregates import ArrayAgg 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 import ArrayField
from django.contrib.postgres.fields.array import ( from django.contrib.postgres.fields.array import (
IndexTransform, SliceTransform, IndexTransform, SliceTransform,
@ -551,6 +552,66 @@ class TestQuerying(PostgreSQLTestCase):
1, 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): class TestDateTimeExactQuerying(PostgreSQLTestCase):