Add PostgreSQL JSON operators to the AST compiler

pydal/ast_translate.py

@@ -27,6 +27,7 @@
 
 # Op names that translate as straight BinOp(name, left, right) with no
 # opts and no structural transformation.
+# Specialized literal type hints are handled before the generic dispatch.
 _PLAIN_BINOPS = frozenset(
     {
         "lt",
@@ -124,6 +125,13 @@ def _select_to_ast(sel: Select) -> ast.Node:
 
 def _field_type(node: Any) -> Optional[str]:
     """Best-effort field-type hint for a left operand."""
+    op_name = getattr(getattr(node, "op", None), "__name__", None)
+    if op_name in ("json_key_value", "json_path_value"):
+        return "string"
+    if op_name in ("lower", "upper", "replace", "substring", "coalesce"):
+        nested_type = _field_type(getattr(node, "first", None))
+        if nested_type:
+            return nested_type
     t = getattr(node, "type", None)
     return t if isinstance(t, str) else None
 
@@ -160,6 +168,20 @@ def _expr_to_ast(expr) -> ast.Node:
             table=getattr(expr, "tablename", None),
         )
 
+    # ---------- JSON operators ----------
+    if name in ("json_key", "json_key_value"):
+        if not isinstance(s, (str, int)):
+            raise TypeError("Key must be a string or int")
+        key_type = "integer" if isinstance(s, int) else "string"
+        return ast.BinOp(name, to_ast(f), to_ast(s, type_hint=key_type))
+    if name in ("json_path", "json_path_value"):
+        path_type = "json_path" if isinstance(s, list) else "string"
+        return ast.BinOp(name, to_ast(f), to_ast(s, type_hint=path_type))
+    if name == "json_contains":
+        # json_contains accepts an already-serialized JSON document. Keep it
+        # as text so parameterized compilation does not serialize it again.
+        return ast.BinOp(name, to_ast(f), to_ast(s, type_hint="string"))
+
     # ---------- comparisons (fold None into is_null/is_not_null) ----------
     if name in ("eq", "ne"):
         if s is None:

pydal/compilers/postgres.py

@@ -18,6 +18,49 @@
 
 @compilers.register_for(Postgres)
 class PostgresCompiler(SQLCompiler):
+    def _json_operand(self, node, sql_type):
+        rendered = self.visit(node)
+        if self._ctx is not None and isinstance(node, ast.Literal):
+            rendered = "%s::%s" % (rendered, sql_type)
+        return rendered
+
+    def _json_path_operand(self, node):
+        if isinstance(node, ast.Literal) and node.type == "json_path":
+            if self._ctx is not None:
+                return "%s::text[]" % self._ctx.bind(node.value)
+            if self.adapter is not None and isinstance(self.adapter, PostgresPsyco):
+                return str(self.adapter.adapt(node.value))
+            return "ARRAY[%s]::text[]" % ",".join(
+                str(self._represent(value, "string")) for value in node.value
+            )
+        return self._json_operand(node, "text[]")
+
+    def op_json_key(self, l, r, _):
+        key_type = "integer" if getattr(r, "type", None) == "integer" else "text"
+        return "%s->%s" % (
+            self.visit(l),
+            self._json_operand(r, key_type),
+        )
+
+    def op_json_key_value(self, l, r, _):
+        key_type = "integer" if getattr(r, "type", None) == "integer" else "text"
+        return "%s->>%s" % (
+            self.visit(l),
+            self._json_operand(r, key_type),
+        )
+
+    def op_json_path(self, l, r, _):
+        return "%s#>%s" % (self.visit(l), self._json_path_operand(r))
+
+    def op_json_path_value(self, l, r, _):
+        return "%s#>>%s" % (self.visit(l), self._json_path_operand(r))
+
+    def op_json_contains(self, l, r, _):
+        left = self.visit(l)
+        if not isinstance(l, ast.FieldRef):
+            left = "(%s)" % left
+        return "%s::jsonb@>%s::jsonb" % (left, self.visit(r))
+
     def _render_insert(self, n: ast.Insert, table: str, cols: str, values: str) -> str:
         sql = super()._render_insert(n, table, cols, values)
         if self.adapter is None:

tests/__init__.py

@@ -18,6 +18,7 @@
 from .ast_translate import *
 from .connection import *
 from .cross_dialect import *
+from .postgres_json import *
 from .driver_io import *
 from .tier2_units import *
 from .tier4_units import *