From: Mike Bayer Date: Fri, 11 Jul 2025 16:51:44 +0000 (-0400) Subject: support JSONB subscripting syntax X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=2ac0c459669d0200a4d15162ffc0a8765c60cb07;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git 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 --- diff --git a/doc/build/changelog/unreleased_20/10927.rst b/doc/build/changelog/unreleased_20/10927.rst new file mode 100644 index 0000000000..765cffd5f9 --- /dev/null +++ b/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. diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py index 75e7670c1f..43251ce277 100644 --- a/lib/sqlalchemy/dialects/postgresql/base.py +++ b/lib/sqlalchemy/dialects/postgresql/base.py @@ -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. diff --git a/test/dialect/postgresql/test_compiler.py b/test/dialect/postgresql/test_compiler.py index f98ea9645b..6c8289867e 100644 --- a/test/dialect/postgresql/test_compiler.py +++ b/test/dialect/postgresql/test_compiler.py @@ -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 diff --git a/test/dialect/postgresql/test_query.py b/test/dialect/postgresql/test_query.py index fc68e08ed4..fc9f7f7918 100644 --- a/test/dialect/postgresql/test_query.py +++ b/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"}, ) diff --git a/test/dialect/postgresql/test_types.py b/test/dialect/postgresql/test_types.py index 2e0d1ea5f6..b78c35d359 100644 --- a/test/dialect/postgresql/test_types.py +++ b/test/dialect/postgresql/test_types.py @@ -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",)