Fix JSON extraction with -> in function arguments

django_ormql/expressions.py

@@ -1140,6 +1140,41 @@ class JSONExtractScalar(JSONExtract):
     sql_type = "TEXT"
 
 
+@register
+class Lambda(Expression):
+    # We do not support lambdas, but the parser emits them for structurs like fun(field->path)
+    # which we actually want to interpret as JSON access. Unfortunately, we need to recursively
+    # insert our JSONExtract node at the innermost level if we have nested JSON extraction, which
+    # makes this more complex.
+    node_class = expressions.Lambda
+
+    def _replace_recursively(self, root_field, expr):
+        if not isinstance(
+            expr, (expressions.JSONExtract, expressions.JSONExtractScalar)
+        ):
+            raise QueryNotSupported("Invalid usage of JSON lookup")
+        if isinstance(expr.this, expressions.Column):
+            return expr.__class__(
+                this=expressions.JSONExtract(this=root_field, expression=expr.this),
+                expression=expr.expression,
+            )
+        else:
+            return expr.__class__(
+                this=self._replace_recursively(root_field, expr.this),
+                expression=expr.expression,
+            )
+
+    def to_django(self, expression, recurse, **kwargs):
+        root_field = expressions.Column(this=expression.expressions[0])
+        new_expr = self._replace_recursively(root_field, expression.this)
+        return expression_to_django(new_expr, recurse, **kwargs)
+
+    def to_sql_type(self, expression, recurse, **kwargs):
+        root_field = expressions.Column(this=expression.expressions[0])
+        new_expr = self._replace_recursively(root_field, expression.this)
+        return expression_to_sql_type(new_expr, **kwargs)
+
+
 @register
 class TypeInfo(FuncExpression):
     func_name = "type_info"

tests/test_jsonfield.py

@@ -6,7 +6,7 @@
 def test_select_json(engine_t1):
     res = engine_t1.query(
         """
-        SELECT address
+        SELECT address, TYPE_INFO(address) as type_info
         FROM customers
         WHERE name = "CA"
         """
@@ -20,7 +20,8 @@ def test_select_json(engine_t1):
                     "name": "Heidelberg",
                     "state": {"code": "BW", "country": {"code": "DE"}},
                 },
-            }
+            },
+            "type_info": "JSONB",
         }
     ]
 
@@ -29,12 +30,14 @@ def test_select_json(engine_t1):
 def test_select_json_key(engine_t1):
     res = engine_t1.query(
         """
-        SELECT address->city->state AS state
+        SELECT address->city->state AS state, TYPE_INFO(address->city->state) AS type_info
         FROM customers
         WHERE name = "CA"
         """
     )
-    assert list(res) == [{"state": {"code": "BW", "country": {"code": "DE"}}}]
+    assert list(res) == [
+        {"state": {"code": "BW", "country": {"code": "DE"}}, "type_info": "JSONB"}
+    ]
 
     res = engine_t1.query(
         """
@@ -50,12 +53,14 @@ def test_select_json_key(engine_t1):
 
     res = engine_t1.query(
         """
-        SELECT address->city->state->>code AS state
+        SELECT address->city->state->>code AS state,
+            LOWER(address->city->state->>code) AS lower_code,
+            TYPE_INFO(address->city->state->>code) AS type_info
         FROM customers
         WHERE name = "CA"
         """
     )
-    assert list(res) == [{"state": "BW"}]
+    assert list(res) == [{"state": "BW", "lower_code": "bw", "type_info": "TEXT"}]
 
 
 @pytest.mark.django_db