]> git.ipfire.org Git - thirdparty/sqlalchemy/sqlalchemy.git/commitdiff
support JSONB subscripting syntax
authorMike Bayer <mike_mp@zzzcomputing.com>
Fri, 11 Jul 2025 16:51:44 +0000 (12:51 -0400)
committerMike Bayer <mike_mp@zzzcomputing.com>
Sun, 13 Jul 2025 15:58:48 +0000 (11:58 -0400)
Added support for PostgreSQL 14+ JSONB subscripting syntax. When connected
to PostgreSQL 14 or later, JSONB columns now automatically use the native
subscript notation ``jsonb_col['key']`` instead of the arrow operator
``jsonb_col -> 'key'`` for both read and write operations. This provides
better compatibility with PostgreSQL's native JSONB subscripting feature
while maintaining backward compatibility with older PostgreSQL versions.
JSON columns continue to use the traditional arrow syntax regardless of
PostgreSQL version.

Fixes: #10927
Change-Id: I4b3a8a55a71f2ca3d95416a7b350b785574631eb

doc/build/changelog/unreleased_20/10927.rst [new file with mode: 0644]
lib/sqlalchemy/dialects/postgresql/base.py
test/dialect/postgresql/test_compiler.py
test/dialect/postgresql/test_query.py
test/dialect/postgresql/test_types.py

diff --git a/doc/build/changelog/unreleased_20/10927.rst b/doc/build/changelog/unreleased_20/10927.rst
new file mode 100644 (file)
index 0000000..765cffd
--- /dev/null
@@ -0,0 +1,12 @@
+.. change::
+    :tags: usecase, postgresql
+    :tickets: 10927
+
+    Added support for PostgreSQL 14+ JSONB subscripting syntax. When connected
+    to PostgreSQL 14 or later, JSONB columns now automatically use the native
+    subscript notation ``jsonb_col['key']`` instead of the arrow operator
+    ``jsonb_col -> 'key'`` for both read and write operations. This provides
+    better compatibility with PostgreSQL's native JSONB subscripting feature
+    while maintaining backward compatibility with older PostgreSQL versions.
+    JSON columns continue to use the traditional arrow syntax regardless of
+    PostgreSQL version.
index 75e7670c1fee48945d934c1de9d9feaa40314006..43251ce2773aeaf49d34582fa4150d55dbf9d3ab 100644 (file)
@@ -1874,9 +1874,23 @@ class PGCompiler(compiler.SQLCompiler):
 
         kw["eager_grouping"] = True
 
-        return self._generate_generic_binary(
-            binary, " -> " if not _cast_applied else " ->> ", **kw
-        )
+        if (
+            not _cast_applied
+            and isinstance(binary.left.type, _json.JSONB)
+            and self.dialect._supports_jsonb_subscripting
+        ):
+            # for pg14+JSONB use subscript notation: col['key'] instead
+            # of col -> 'key'
+            return "%s[%s]" % (
+                self.process(binary.left, **kw),
+                self.process(binary.right, **kw),
+            )
+        else:
+            # Fall back to arrow notation for older versions or when cast
+            # is applied
+            return self._generate_generic_binary(
+                binary, " -> " if not _cast_applied else " ->> ", **kw
+            )
 
     def visit_json_path_getitem_op_binary(
         self, binary, operator, _cast_applied=False, **kw
@@ -3233,6 +3247,7 @@ class PGDialect(default.DefaultDialect):
     _backslash_escapes = True
     _supports_create_index_concurrently = True
     _supports_drop_index_concurrently = True
+    _supports_jsonb_subscripting = True
 
     def __init__(
         self,
@@ -3261,6 +3276,8 @@ class PGDialect(default.DefaultDialect):
         )
         self.supports_identity_columns = self.server_version_info >= (10,)
 
+        self._supports_jsonb_subscripting = self.server_version_info >= (14,)
+
     def get_isolation_level_values(self, dbapi_conn):
         # note the generic dialect doesn't have AUTOCOMMIT, however
         # all postgresql dialects should include AUTOCOMMIT.
index f98ea9645b043e40d90bd7289c67ca0654752094..6c8289867e2cf171fddcce3516acb7d0f43906c2 100644 (file)
@@ -2675,22 +2675,87 @@ class CompileTest(fixtures.TestBase, AssertsCompiledSQL):
         )
 
     @testing.combinations(
-        (lambda col: col["foo"] + " ", "(x -> %(x_1)s) || %(param_1)s"),
+        (
+            lambda col: col["foo"] + " ",
+            "(x -> %(x_1)s) || %(param_1)s",
+            "x[%(x_1)s] || %(param_1)s",
+        ),
         (
             lambda col: col["foo"] + " " + col["bar"],
             "(x -> %(x_1)s) || %(param_1)s || (x -> %(x_2)s)",
+            "x[%(x_1)s] || %(param_1)s || x[%(x_2)s]",
         ),
-        argnames="expr, expected",
+        argnames="expr, json_expected, jsonb_expected",
     )
     @testing.combinations((JSON(),), (JSONB(),), argnames="type_")
-    def test_eager_grouping_flag(self, expr, expected, type_):
+    def test_eager_grouping_flag(
+        self, expr, json_expected, jsonb_expected, type_
+    ):
         """test #10479"""
         col = Column("x", type_)
 
         expr = testing.resolve_lambda(expr, col=col)
 
+        # Choose expected result based on type
+        expected = (
+            jsonb_expected if isinstance(type_, JSONB) else json_expected
+        )
         self.assert_compile(expr, expected)
 
+    @testing.variation("pgversion", ["pg14", "pg13"])
+    def test_jsonb_subscripting(self, pgversion):
+        """test #10927 - PostgreSQL 14+ JSONB subscripting syntax"""
+        data = table("data", column("id", Integer), column("x", JSONB))
+
+        dialect = postgresql.dialect()
+
+        if pgversion.pg13:
+            dialect._supports_jsonb_subscripting = False
+
+        # Test SELECT with JSONB indexing
+        stmt = select(data.c.x["key"])
+        self.assert_compile(
+            stmt,
+            (
+                "SELECT data.x[%(x_1)s] AS anon_1 FROM data"
+                if pgversion.pg14
+                else "SELECT data.x -> %(x_1)s AS anon_1 FROM data"
+            ),
+            dialect=dialect,
+        )
+
+        # Test UPDATE with JSONB indexing (the original issue case)
+        stmt = update(data).values({data.c.x["new_key"]: data.c.x["old_key"]})
+        self.assert_compile(
+            stmt,
+            (
+                "UPDATE data SET x[%(x_1)s]=(data.x[%(x_2)s])"
+                if pgversion.pg14
+                else "UPDATE data SET x -> %(x_1)s=(data.x -> %(x_2)s)"
+            ),
+            dialect=dialect,
+        )
+
+    def test_json_still_uses_arrow_syntax(self):
+        """test #10927 - JSON type still uses arrow syntax even on PG 14+"""
+        data = table("data", column("id", Integer), column("x", JSON))
+
+        # Test PostgreSQL 14+ still uses arrow syntax for JSON (not JSONB)
+
+        # Test SELECT with JSON indexing
+        stmt = select(data.c.x["key"])
+        self.assert_compile(
+            stmt,
+            "SELECT data.x -> %(x_1)s AS anon_1 FROM data",
+        )
+
+        # Test UPDATE with JSON indexing
+        stmt = update(data).values({data.c.x["new_key"]: data.c.x["old_key"]})
+        self.assert_compile(
+            stmt,
+            "UPDATE data SET x -> %(x_1)s=(data.x -> %(x_2)s)",
+        )
+
     def test_range_custom_object_hook(self):
         # See issue #8884
         from datetime import date
index fc68e08ed4df7bad503d8b664e9a20e64060bde8..fc9f7f79188f8111864106c0ecface23b66d76f2 100644 (file)
@@ -1799,7 +1799,12 @@ class TableValuedRoundTripTest(fixtures.TestBase):
         eq_(connection.execute(stmt).all(), [(1, "foo"), (2, "bar")])
 
 
-class RequiresCastTest(fixtures.TablesTest):
+class JSONUpdateTest(fixtures.TablesTest):
+    """round trip tests related to using JSON and JSONB in UPDATE statements
+    with PG-specific features
+
+    """
+
     __only_on__ = "postgresql"
     __backend__ = True
 
@@ -1825,6 +1830,8 @@ class RequiresCastTest(fixtures.TablesTest):
         )
 
     def test_update_values(self, connection):
+        t = self.tables["t"]
+
         value = values(
             Column("id", Integer),
             Column("uuid", Uuid),
@@ -1848,8 +1855,149 @@ class RequiresCastTest(fixtures.TablesTest):
             ]
         )
         connection.execute(
-            self.tables["t"]
-            .update()
+            t.update()
             .values(uuid=value.c.uuid, j=value.c.j, jb=value.c.jb)
-            .where(self.tables["t"].c.id == value.c.id)
+            .where(t.c.id == value.c.id)
+        )
+
+        updated_data = connection.execute(t.select().order_by(t.c.id))
+        eq_(
+            [(str(row.uuid), row.j, row.jb) for row in updated_data],
+            [
+                (
+                    "8b6ec1ec-b979-4d0b-b2ce-9acc6e4c2943",
+                    {"foo": 1},
+                    {"foo_jb": 1},
+                ),
+                (
+                    "a2123bcb-7ea3-420a-8284-1db4b2759d79",
+                    {"bar": 2},
+                    {"bar_jb": 2},
+                ),
+            ],
+        )
+
+    @testing.only_on("postgresql>=14")
+    def test_jsonb_element_update_basic(self, connection):
+        """Test updating individual JSONB elements with subscript syntax
+
+        test #10927
+
+        """
+        t = self.tables["t"]
+
+        # Insert test data with complex JSONB
+        connection.execute(
+            t.insert(),
+            [
+                {
+                    "id": 10,
+                    "jb": {
+                        "user": {"name": "Alice", "age": 30},
+                        "active": True,
+                    },
+                },
+                {
+                    "id": 11,
+                    "jb": {
+                        "user": {"name": "Bob", "age": 25},
+                        "active": False,
+                    },
+                },
+            ],
+        )
+
+        # Update specific elements using JSONB subscript syntax
+        # This tests the new JSONB subscripting feature from issue #10927
+        connection.execute(
+            t.update()
+            .values({t.c.jb["user"]["name"]: "Alice Updated"})
+            .where(t.c.id == 10)
+        )
+
+        connection.execute(
+            t.update().values({t.c.jb["active"]: True}).where(t.c.id == 11)
+        )
+
+        results = connection.execute(
+            t.select().where(t.c.id.in_([10, 11])).order_by(t.c.id)
+        )
+
+        eq_(
+            [row.jb for row in results],
+            [
+                {"user": {"name": "Alice Updated", "age": 30}, "active": True},
+                {"user": {"name": "Bob", "age": 25}, "active": True},
+            ],
+        )
+
+    @testing.only_on("postgresql>=14")
+    def test_jsonb_element_update_multiple_keys(self, connection):
+        """Test updating multiple JSONB elements in a single statement
+
+        test #10927
+
+        """
+        t = self.tables["t"]
+
+        connection.execute(
+            t.insert(),
+            {
+                "id": 20,
+                "jb": {
+                    "config": {"theme": "dark", "lang": "en"},
+                    "version": 1,
+                },
+            },
+        )
+
+        # Update multiple elements at once
+        connection.execute(
+            t.update()
+            .values({t.c.jb["config"]["theme"]: "light", t.c.jb["version"]: 2})
+            .where(t.c.id == 20)
+        )
+
+        # Verify the updates
+        row = connection.execute(t.select().where(t.c.id == 20)).one()
+
+        eq_(
+            row.jb,
+            {"config": {"theme": "light", "lang": "en"}, "version": 2},
+        )
+
+    @testing.only_on("postgresql>=14")
+    def test_jsonb_element_update_array_element(self, connection):
+        """Test updating JSONB array elements
+
+        test #10927
+
+        """
+        t = self.tables["t"]
+
+        # Insert test data with arrays
+        connection.execute(
+            t.insert(),
+            {
+                "id": 30,
+                "jb": {
+                    "tags": ["python", "sql", "postgres"],
+                    "priority": "high",
+                },
+            },
+        )
+
+        # Update array element
+        connection.execute(
+            t.update()
+            .values({t.c.jb["tags"][1]: "postgresql"})
+            .where(t.c.id == 30)
+        )
+
+        # Verify the update
+        row = connection.execute(t.select().where(t.c.id == 30)).fetchone()
+
+        eq_(
+            row.jb,
+            {"tags": ["python", "postgresql", "postgres"], "priority": "high"},
         )
index 2e0d1ea5f6ba30231b1fd98d808be27caefef51b..b78c35d35903109b0764e3caf98e206874246d33 100644 (file)
@@ -6028,15 +6028,9 @@ class JSONTest(AssertsCompiledSQL, fixtures.TestBase):
     def any_(self):
         return any_(array([7]))
 
+    # Test combinations that use path (#>) and astext (->> and #>>) operators
+    # These don't change between JSON and JSONB
     @testing.combinations(
-        (
-            lambda self: self.jsoncol["bar"] == None,  # noqa
-            "(test_table.test_column -> %(test_column_1)s) IS NULL",
-        ),
-        (
-            lambda self: self.jsoncol["bar"] != None,  # noqa
-            "(test_table.test_column -> %(test_column_1)s) IS NOT NULL",
-        ),
         (
             lambda self: self.jsoncol[("foo", 1)] == None,  # noqa
             "(test_table.test_column #> %(test_column_1)s) IS NULL",
@@ -6054,33 +6048,10 @@ class JSONTest(AssertsCompiledSQL, fixtures.TestBase):
             "CAST((test_table.test_column ->> %(test_column_1)s) AS INTEGER) "
             "= %(param_1)s",
         ),
-        (
-            lambda self: self.jsoncol["bar"].cast(Integer) == 5,
-            "CAST((test_table.test_column -> %(test_column_1)s) AS INTEGER) "
-            "= %(param_1)s",
-        ),
         (
             lambda self: self.jsoncol[("foo", 1)].astext == None,  # noqa
             "(test_table.test_column #>> %(test_column_1)s) IS NULL",
         ),
-        (
-            lambda self: self.jsoncol["bar"] == 42,
-            "(test_table.test_column -> %(test_column_1)s) = %(param_1)s",
-        ),
-        (
-            lambda self: self.jsoncol["bar"] != 42,
-            "(test_table.test_column -> %(test_column_1)s) != %(param_1)s",
-        ),
-        (
-            lambda self: self.jsoncol["bar"] == self.any_,
-            "(test_table.test_column -> %(test_column_1)s) = "
-            "ANY (ARRAY[%(param_1)s])",
-        ),
-        (
-            lambda self: self.jsoncol["bar"] != self.any_,
-            "(test_table.test_column -> %(test_column_1)s) != "
-            "ANY (ARRAY[%(param_1)s])",
-        ),
         (
             lambda self: self.jsoncol["bar"].astext == self.any_,
             "(test_table.test_column ->> %(test_column_1)s) = "
@@ -6112,6 +6083,51 @@ class JSONTest(AssertsCompiledSQL, fixtures.TestBase):
             "WHERE %s" % expected,
         )
 
+    # Test combinations that use subscript (->) operator
+    # These differ between JSON (always ->) and JSONB ([] on PG 14+)
+    @testing.combinations(
+        (
+            lambda self: self.jsoncol["bar"] == None,  # noqa
+            "(test_table.test_column -> %(test_column_1)s) IS NULL",
+        ),
+        (
+            lambda self: self.jsoncol["bar"] != None,  # noqa
+            "(test_table.test_column -> %(test_column_1)s) IS NOT NULL",
+        ),
+        (
+            lambda self: self.jsoncol["bar"].cast(Integer) == 5,
+            "CAST((test_table.test_column -> %(test_column_1)s) AS INTEGER) "
+            "= %(param_1)s",
+        ),
+        (
+            lambda self: self.jsoncol["bar"] == 42,
+            "(test_table.test_column -> %(test_column_1)s) = %(param_1)s",
+        ),
+        (
+            lambda self: self.jsoncol["bar"] != 42,
+            "(test_table.test_column -> %(test_column_1)s) != %(param_1)s",
+        ),
+        (
+            lambda self: self.jsoncol["bar"] == self.any_,
+            "(test_table.test_column -> %(test_column_1)s) = "
+            "ANY (ARRAY[%(param_1)s])",
+        ),
+        (
+            lambda self: self.jsoncol["bar"] != self.any_,
+            "(test_table.test_column -> %(test_column_1)s) != "
+            "ANY (ARRAY[%(param_1)s])",
+        ),
+        id_="as",
+    )
+    def test_where_subscript(self, whereclause_fn, expected):
+        whereclause = whereclause_fn(self)
+        stmt = select(self.test_table).where(whereclause)
+        self.assert_compile(
+            stmt,
+            "SELECT test_table.id, test_table.test_column FROM test_table "
+            "WHERE %s" % expected,
+        )
+
     def test_path_typing(self):
         col = column("x", JSON())
         is_(col["q"].type._type_affinity, types.JSON)
@@ -6131,6 +6147,8 @@ class JSONTest(AssertsCompiledSQL, fixtures.TestBase):
 
         is_(col["q"]["p"].astext.type.__class__, MyType)
 
+    # Test column selection that uses subscript (->) operator
+    # This differs between JSON (always ->) and JSONB ([] on PG 14+)
     @testing.combinations(
         (
             lambda self: self.jsoncol["foo"],
@@ -6138,7 +6156,7 @@ class JSONTest(AssertsCompiledSQL, fixtures.TestBase):
             True,
         )
     )
-    def test_cols(self, colclause_fn, expected, from_):
+    def test_cols_subscript(self, colclause_fn, expected, from_):
         colclause = colclause_fn(self)
         stmt = select(colclause)
         self.assert_compile(
@@ -6500,6 +6518,68 @@ class JSONBTest(JSONTest):
     def test_where_jsonb(self, whereclause_fn, expected):
         super().test_where(whereclause_fn, expected)
 
+    # Override test_where_subscript to provide JSONB-specific expectations
+    # JSONB uses subscript syntax (e.g., col['key']) on PostgreSQL 14+
+    @testing.combinations(
+        (
+            lambda self: self.jsoncol["bar"] == None,  # noqa
+            "test_table.test_column[%(test_column_1)s] IS NULL",
+        ),
+        (
+            lambda self: self.jsoncol["bar"] != None,  # noqa
+            "test_table.test_column[%(test_column_1)s] IS NOT NULL",
+        ),
+        (
+            lambda self: self.jsoncol["bar"].cast(Integer) == 5,
+            "CAST(test_table.test_column[%(test_column_1)s] AS INTEGER) "
+            "= %(param_1)s",
+        ),
+        (
+            lambda self: self.jsoncol["bar"] == 42,
+            "test_table.test_column[%(test_column_1)s] = %(param_1)s",
+        ),
+        (
+            lambda self: self.jsoncol["bar"] != 42,
+            "test_table.test_column[%(test_column_1)s] != %(param_1)s",
+        ),
+        (
+            lambda self: self.jsoncol["bar"] == self.any_,
+            "test_table.test_column[%(test_column_1)s] = "
+            "ANY (ARRAY[%(param_1)s])",
+        ),
+        (
+            lambda self: self.jsoncol["bar"] != self.any_,
+            "test_table.test_column[%(test_column_1)s] != "
+            "ANY (ARRAY[%(param_1)s])",
+        ),
+        id_="as",
+    )
+    def test_where_subscript(self, whereclause_fn, expected):
+        whereclause = whereclause_fn(self)
+        stmt = select(self.test_table).where(whereclause)
+        self.assert_compile(
+            stmt,
+            "SELECT test_table.id, test_table.test_column FROM test_table "
+            "WHERE %s" % expected,
+        )
+
+    # Override test_cols_subscript to provide JSONB-specific expectations
+    # JSONB uses subscript syntax (e.g., col['key']) on PostgreSQL 14+
+    @testing.combinations(
+        (
+            lambda self: self.jsoncol["foo"],
+            "test_table.test_column[%(test_column_1)s] AS anon_1",
+            True,
+        )
+    )
+    def test_cols_subscript(self, colclause_fn, expected, from_):
+        colclause = colclause_fn(self)
+        stmt = select(colclause)
+        self.assert_compile(
+            stmt,
+            ("SELECT %s" + (" FROM test_table" if from_ else "")) % expected,
+        )
+
 
 class JSONBRoundTripTest(JSONRoundTripTest):
     __requires__ = ("postgresql_jsonb",)