From 10178197d57476f69688d4535e550a1ea3a5eac5 Mon Sep 17 00:00:00 2001
From: Allen Jonathan David <allenajdjonathan@gmail.com>
Date: Tue, 30 Aug 2022 22:56:18 +0530
Subject: [PATCH] Fixed #33966 -- Added support for using KeyTextTransform from
 lookup.

---
 django/db/models/fields/json.py      | 13 +++++++++++++
 docs/releases/4.2.txt                |  4 ++++
 docs/topics/db/queries.txt           | 27 +++++++++++++++++++++++++++
 tests/model_fields/test_jsonfield.py | 24 +++++++++++++++++++-----
 4 files changed, 63 insertions(+), 5 deletions(-)

diff --git a/django/db/models/fields/json.py b/django/db/models/fields/json.py
index 16231f24ac..7296fe42bc 100644
--- a/django/db/models/fields/json.py
+++ b/django/db/models/fields/json.py
@@ -4,6 +4,7 @@ from django import forms
 from django.core import checks, exceptions
 from django.db import NotSupportedError, connections, router
 from django.db.models import lookups
+from django.db.models.constants import LOOKUP_SEP
 from django.db.models.fields import TextField
 from django.db.models.lookups import PostgresOperatorLookup, Transform
 from django.utils.translation import gettext_lazy as _
@@ -379,6 +380,18 @@ class KeyTextTransform(KeyTransform):
             json_path = compile_json_path(key_transforms)
             return "(%s ->> %%s)" % lhs, tuple(params) + (json_path,)
 
+    @classmethod
+    def from_lookup(cls, lookup):
+        transform, *keys = lookup.split(LOOKUP_SEP)
+        if not keys:
+            raise ValueError("Lookup must contain key or index transforms.")
+        for key in keys:
+            transform = cls(key, transform)
+        return transform
+
+
+KT = KeyTextTransform.from_lookup
+
 
 class KeyTransformTextLookupMixin:
     """
diff --git a/docs/releases/4.2.txt b/docs/releases/4.2.txt
index e4394c5887..7ca4f0271e 100644
--- a/docs/releases/4.2.txt
+++ b/docs/releases/4.2.txt
@@ -216,6 +216,10 @@ Models
   allows performing actions that can fail after a database transaction is
   successfully committed.
 
+* The new :class:`KT() <django.db.models.fields.json.KT>` expression represents
+  the text value of a key, index, or path transform of
+  :class:`~django.db.models.JSONField`.
+
 Requests and Responses
 ~~~~~~~~~~~~~~~~~~~~~~
 
diff --git a/docs/topics/db/queries.txt b/docs/topics/db/queries.txt
index 6deeec5037..5114efb57d 100644
--- a/docs/topics/db/queries.txt
+++ b/docs/topics/db/queries.txt
@@ -1059,6 +1059,33 @@ To query for missing keys, use the ``isnull`` lookup::
     :lookup:`istartswith`, :lookup:`lt`, :lookup:`lte`, :lookup:`gt`, and
     :lookup:`gte`, as well as with :ref:`containment-and-key-lookups`.
 
+``KT()`` expressions
+~~~~~~~~~~~~~~~~~~~~
+
+.. versionadded:: 4.2
+
+.. module:: django.db.models.fields.json
+
+.. class:: KT(lookup)
+
+    Represents the text value of a key, index, or path transform of
+    :class:`~django.db.models.JSONField`. You can use the double underscore
+    notation in ``lookup`` to chain dictionary key and index transforms.
+
+    For example::
+
+        >>> from django.db.models.fields.json import KT
+        >>> Dog.objects.create(name="Shep", data={
+        ...     "owner": {"name": "Bob"},
+        ...     "breed": ["collie", "lhasa apso"],
+        ... })
+        <Dog: Shep>
+        >>> Dogs.objects.annotate(
+        ...     first_breed=KT("data__breed__1"),
+        ...     owner_name=KT("data__owner__name")
+        ... ).filter(first_breed__startswith="lhasa", owner_name="Bob")
+        <QuerySet [<Dog: Shep>]>
+
 .. note::
 
     Due to the way in which key-path queries work,
diff --git a/tests/model_fields/test_jsonfield.py b/tests/model_fields/test_jsonfield.py
index d6b9f031b1..2c32d8a4ea 100644
--- a/tests/model_fields/test_jsonfield.py
+++ b/tests/model_fields/test_jsonfield.py
@@ -27,6 +27,7 @@ from django.db.models import (
 )
 from django.db.models.expressions import RawSQL
 from django.db.models.fields.json import (
+    KT,
     KeyTextTransform,
     KeyTransform,
     KeyTransformFactory,
@@ -374,11 +375,7 @@ class TestQuerying(TestCase):
         qs = NullableJSONModel.objects.filter(value__isnull=False)
         self.assertQuerysetEqual(
             qs.filter(value__isnull=False)
-            .annotate(
-                key=KeyTextTransform(
-                    "f", KeyTransform("1", KeyTransform("d", "value"))
-                ),
-            )
+            .annotate(key=KT("value__d__1__f"))
             .values("key")
             .annotate(count=Count("key"))
             .order_by("count"),
@@ -1078,3 +1075,20 @@ class TestQuerying(TestCase):
             ).filter(chain=F("related_key__0")),
             [related_obj],
         )
+
+    def test_key_text_transform_from_lookup(self):
+        qs = NullableJSONModel.objects.annotate(b=KT("value__bax__foo")).filter(
+            b__contains="ar",
+        )
+        self.assertSequenceEqual(qs, [self.objs[7]])
+        qs = NullableJSONModel.objects.annotate(c=KT("value__o")).filter(
+            c__contains="uot",
+        )
+        self.assertSequenceEqual(qs, [self.objs[4]])
+
+    def test_key_text_transform_from_lookup_invalid(self):
+        msg = "Lookup must contain key or index transforms."
+        with self.assertRaisesMessage(ValueError, msg):
+            KT("value")
+        with self.assertRaisesMessage(ValueError, msg):
+            KT("")