From 5a634a7b6f248e5dead65f420ec5d30acd9724bf Mon Sep 17 00:00:00 2001
From: abhiabhi94 <13880786+abhiabhi94@users.noreply.github.com>
Date: Thu, 15 Jul 2021 19:09:40 +0530
Subject: [PATCH] Fixed #32906 -- Added docs and tests for using key and index
 lookups on JSONBAgg results.

Co-authored-by: Mariusz Felisiak <felisiak.mariusz@gmail.com>
---
 docs/ref/contrib/postgres/aggregates.txt | 26 +++++++++++++-
 tests/postgres_tests/test_aggregates.py  | 45 +++++++++++++++++++++++-
 2 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/docs/ref/contrib/postgres/aggregates.txt b/docs/ref/contrib/postgres/aggregates.txt
index 6d85d5c4783..9223f17ab86 100644
--- a/docs/ref/contrib/postgres/aggregates.txt
+++ b/docs/ref/contrib/postgres/aggregates.txt
@@ -125,7 +125,8 @@ General-purpose aggregation functions
 .. class:: JSONBAgg(expressions, distinct=False, filter=None, default=None, ordering=(), **extra)
 
     Returns the input values as a ``JSON`` array, or ``default`` if there are
-    no values.
+    no values. You can query the result using :lookup:`key and index lookups
+    <jsonfield.key>`.
 
     .. attribute:: distinct
 
@@ -145,6 +146,29 @@ General-purpose aggregation functions
 
         Examples are the same as for :attr:`ArrayAgg.ordering`.
 
+    Usage example::
+
+        class Room(models.Model):
+            number = models.IntegerField(unique=True)
+
+        class HotelReservation(model.Model):
+            room = models.ForeignKey('Room', on_delete=models.CASCADE)
+            start = models.DateTimeField()
+            end = models.DateTimeField()
+            requirements = models.JSONField(blank=True, null=True)
+
+        >>> from django.contrib.postgres.aggregates import JSONBAgg
+        >>> Room.objects.annotate(
+        ...     requirements=JSONBAgg(
+        ...         'hotelreservation__requirements',
+        ...         ordering='-hotelreservation__start',
+        ...     )
+        ... ).filter(requirements__0__sea_view=True).values('number', 'requirements')
+        <QuerySet [{'number': 102, 'requirements': [
+            {'parking': False, 'sea_view': True, 'double_bed': False},
+            {'parking': True, 'double_bed': True}
+        ]}]>
+
     .. deprecated:: 4.0
 
         If there are no rows and ``default`` is not provided, ``JSONBAgg``
diff --git a/tests/postgres_tests/test_aggregates.py b/tests/postgres_tests/test_aggregates.py
index 7ae7b16c9f2..a963aae3c4a 100644
--- a/tests/postgres_tests/test_aggregates.py
+++ b/tests/postgres_tests/test_aggregates.py
@@ -4,10 +4,11 @@ from django.db.models import (
 from django.db.models.fields.json import KeyTextTransform, KeyTransform
 from django.db.models.functions import Cast, Concat, Substr
 from django.test.utils import Approximate, ignore_warnings
+from django.utils import timezone
 from django.utils.deprecation import RemovedInDjango50Warning
 
 from . import PostgreSQLTestCase
-from .models import AggregateTestModel, StatTestModel
+from .models import AggregateTestModel, HotelReservation, Room, StatTestModel
 
 try:
     from django.contrib.postgres.aggregates import (
@@ -392,6 +393,48 @@ class TestGeneralAggregate(PostgreSQLTestCase):
         )
         self.assertEqual(values, {'jsonbagg': ['en', 'pl']})
 
+    def test_jsonb_agg_key_index_transforms(self):
+        room101 = Room.objects.create(number=101)
+        room102 = Room.objects.create(number=102)
+        datetimes = [
+            timezone.datetime(2018, 6, 20),
+            timezone.datetime(2018, 6, 24),
+            timezone.datetime(2018, 6, 28),
+        ]
+        HotelReservation.objects.create(
+            datespan=(datetimes[0].date(), datetimes[1].date()),
+            start=datetimes[0],
+            end=datetimes[1],
+            room=room102,
+            requirements={'double_bed': True, 'parking': True},
+        )
+        HotelReservation.objects.create(
+            datespan=(datetimes[1].date(), datetimes[2].date()),
+            start=datetimes[1],
+            end=datetimes[2],
+            room=room102,
+            requirements={'double_bed': False, 'sea_view': True, 'parking': False},
+        )
+        HotelReservation.objects.create(
+            datespan=(datetimes[0].date(), datetimes[2].date()),
+            start=datetimes[0],
+            end=datetimes[2],
+            room=room101,
+            requirements={'sea_view': False},
+        )
+        values = Room.objects.annotate(
+            requirements=JSONBAgg(
+                'hotelreservation__requirements',
+                ordering='-hotelreservation__start',
+            )
+        ).filter(requirements__0__sea_view=True).values('number', 'requirements')
+        self.assertSequenceEqual(values, [
+            {'number': 102, 'requirements': [
+                {'double_bed': False, 'sea_view': True, 'parking': False},
+                {'double_bed': True, 'parking': True},
+            ]},
+        ])
+
     def test_string_agg_array_agg_ordering_in_subquery(self):
         stats = []
         for i, agg in enumerate(AggregateTestModel.objects.order_by('char_field')):