From: Guilherme Martins Crocetti <24530683+gmcrocetti@users.noreply.github.com>
Date: Wed, 4 Jan 2023 15:27:26 +0000 (-0500)
Subject: Implement missing `#-`, `@?` and `@@` Postgres' JSONB operators.
X-Git-Tag: rel_2_0_0rc2~7^2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=6b0c3dd5fd0c0225636ce759bc4433c91c5efb83;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

Implement missing `#-`, `@?` and `@@` Postgres' JSONB operators.

Fixes #7147.
Closes: #9038
Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/9038
Pull-request-sha: 8647aaf2d9f48c55c152673828deb8ed54966a11

Change-Id: Id2f611ed8080a2837c70d2ea4b41abc46d2bb026
---

diff --git a/doc/build/changelog/unreleased_20/7147.rst b/doc/build/changelog/unreleased_20/7147.rst
new file mode 100644
index 0000000000..629c232278
--- /dev/null
+++ b/doc/build/changelog/unreleased_20/7147.rst
@@ -0,0 +1,11 @@
+.. change::
+    :tags: json, postgresql
+    :tickets: 7147
+
+    Implemented missing ``JSONB`` operations:
+
+    * ``@@`` using :meth:`_postgresql.JSONB.Comparator.path_match`
+    * ``@?`` using :meth:`_postgresql.JSONB.Comparator.path_exists`
+    * ``#-`` using :meth:`_postgresql.JSONB.Comparator.delete_path`
+
+    Pull request curtesy of Guilherme Martins Crocetti.
diff --git a/lib/sqlalchemy/dialects/postgresql/json.py b/lib/sqlalchemy/dialects/postgresql/json.py
index b136ff4d5c..9c29360062 100644
--- a/lib/sqlalchemy/dialects/postgresql/json.py
+++ b/lib/sqlalchemy/dialects/postgresql/json.py
@@ -7,7 +7,10 @@
 # mypy: ignore-errors
 
 
+from .array import ARRAY
+from .array import array as _pg_array
 from ... import types as sqltypes
+from ...sql import cast
 from ...sql import operators
 
 
@@ -65,6 +68,27 @@ CONTAINED_BY = operators.custom_op(
     eager_grouping=True,
 )
 
+DELETE_PATH = operators.custom_op(
+    "#-",
+    precedence=idx_precedence,
+    natural_self_precedent=True,
+    eager_grouping=True,
+)
+
+PATH_EXISTS = operators.custom_op(
+    "@?",
+    precedence=idx_precedence,
+    natural_self_precedent=True,
+    eager_grouping=True,
+)
+
+PATH_MATCH = operators.custom_op(
+    "@@",
+    precedence=idx_precedence,
+    natural_self_precedent=True,
+    eager_grouping=True,
+)
+
 
 class JSONPathType(sqltypes.JSON.JSONPathType):
     def _processor(self, dialect, super_proc):
@@ -279,7 +303,10 @@ class JSONB(JSON):
     It also adds additional operators specific to JSONB, including
     :meth:`.JSONB.Comparator.has_key`, :meth:`.JSONB.Comparator.has_all`,
     :meth:`.JSONB.Comparator.has_any`, :meth:`.JSONB.Comparator.contains`,
-    and :meth:`.JSONB.Comparator.contained_by`.
+    :meth:`.JSONB.Comparator.contained_by`,
+    :meth:`.JSONB.Comparator.delete_path`,
+    :meth:`.JSONB.Comparator.path_exists` and
+    :meth:`.JSONB.Comparator.path_match`.
 
     Like the :class:`_types.JSON` type, the :class:`_postgresql.JSONB`
     type does not detect
@@ -340,4 +367,40 @@ class JSONB(JSON):
                 CONTAINED_BY, other, result_type=sqltypes.Boolean
             )
 
+        def delete_path(self, array):
+            """JSONB expression. Deletes field or array element specified in
+            the argument array.
+
+            The input may be a list of strings that will be coerced to an
+            ``ARRAY`` or an instance of :meth:`_postgres.array`.
+
+            .. versionadded:: 2.0
+            """
+            if not isinstance(array, _pg_array):
+                array = _pg_array(array)
+            right_side = cast(array, ARRAY(sqltypes.TEXT))
+            return self.operate(DELETE_PATH, right_side, result_type=JSONB)
+
+        def path_exists(self, other):
+            """Boolean expression. Test for presence of item given by the
+            argument JSONPath expression.
+
+            .. versionadded:: 2.0
+            """
+            return self.operate(
+                PATH_EXISTS, other, result_type=sqltypes.Boolean
+            )
+
+        def path_match(self, other):
+            """Boolean expression. Test if JSONPath predicate given by the
+            argument JSONPath expression matches.
+
+            Only the first item of the result is taken into account.
+
+            .. versionadded:: 2.0
+            """
+            return self.operate(
+                PATH_MATCH, other, result_type=sqltypes.Boolean
+            )
+
     comparator_factory = Comparator
diff --git a/test/dialect/postgresql/test_types.py b/test/dialect/postgresql/test_types.py
index e9d5e561f8..22993ae7bb 100644
--- a/test/dialect/postgresql/test_types.py
+++ b/test/dialect/postgresql/test_types.py
@@ -5683,6 +5683,24 @@ class JSONBTest(JSONTest):
             lambda self: self.jsoncol.contained_by({"foo": "1", "bar": None}),
             "test_table.test_column <@ %(test_column_1)s",
         ),
+        (
+            lambda self: self.jsoncol.delete_path(["a", "b"]),
+            "test_table.test_column #- CAST(ARRAY[%(param_1)s, "
+            "%(param_2)s] AS TEXT[])",
+        ),
+        (
+            lambda self: self.jsoncol.delete_path(array(["a", "b"])),
+            "test_table.test_column #- CAST(ARRAY[%(param_1)s, "
+            "%(param_2)s] AS TEXT[])",
+        ),
+        (
+            lambda self: self.jsoncol.path_exists("$.k1"),
+            "test_table.test_column @? %(test_column_1)s",
+        ),
+        (
+            lambda self: self.jsoncol.path_match("$.k1[0] > 2"),
+            "test_table.test_column @@ %(test_column_1)s",
+        ),
     )
     def test_where(self, whereclause_fn, expected):
         super().test_where(whereclause_fn, expected)
@@ -5712,6 +5730,19 @@ class JSONBRoundTripTest(JSONRoundTripTest):
         go("$.k1.k2", 0)
         go("$.k1.r6v1", 1)
 
+    @testing.combinations(
+        ["k1", "r6v1", "subr", 1],
+        array(["k1", "r6v1", "subr", 1]),
+        argnames="path",
+    )
+    def test_delete_path(self, connection, path):
+        self._fixture_data(connection)
+        q = select(self.data_table.c.data.delete_path(path)).where(
+            self.data_table.c.name == "r6"
+        )
+        res = connection.scalar(q)
+        eq_(res, {"k1": {"r6v1": {"subr": [1, 3]}}})
+
 
 class JSONBSuiteTest(suite.JSONTest):
     __requires__ = ("postgresql_jsonb",)