support JSONB subscripting syntax

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
(cherry picked from commit 2ac0c45966)

doc/build/changelog/unreleased_20/10927.rst

@@ -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.

lib/sqlalchemy/dialects/postgresql/base.py

@@ -1880,9 +1880,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
@@ -3227,6 +3241,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,
@@ -3255,6 +3270,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.

test/dialect/postgresql/test_compiler.py

@@ -2671,22 +2671,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

test/dialect/postgresql/test_query.py

@@ -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"},
         )

test/dialect/postgresql/test_types.py

@@ -5882,15 +5882,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",
@@ -5908,33 +5902,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) = "
@@ -5966,6 +5937,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)
@@ -5985,6 +6001,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"],
@@ -5992,7 +6010,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(
@@ -6354,6 +6372,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",)