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 52f4721da9..156dccb32f 100644 --- a/lib/sqlalchemy/dialects/postgresql/base.py +++ b/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. diff --git a/test/dialect/postgresql/test_compiler.py b/test/dialect/postgresql/test_compiler.py index 2a763593b2..72adf953e4 100644 --- a/test/dialect/postgresql/test_compiler.py +++ b/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 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 c96cb90378..4613ebf32c 100644 --- a/test/dialect/postgresql/test_types.py +++ b/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",)