Add support for JSONB subscripting introduced in Postgresql 14

[iurisilvio]

Add support for the subscript jsonb notation introduced in pg 14. Ref at https://www.postgresql.org/docs/current/datatype-json.html#JSONB-SUBSCRIPTING.

This notation seems to be allowed in more placed than the classical "arrow" notation.
For example it can be used in a set clause of an updated:

UPDATE foo SET jsonb_col['some_value']=jsonb_col -> 'some_other_val'
UPDATE foo SET jsonb_col['some_value']=jsonb_col['some_other_val']

I'm not sure if this should become the preferred syntax in postgresql 14+.
This new syntax is used initially only by in the update set clause, since it would be not strange to have different syntax depending on json/jsonb. May be revisited in the future

Original text below

Describe the bug

I don't know if it is an intentional unsupported feature or a bug.

I'm trying to copy a value from a key to another key and it fails on Postgres.

For Postgres 16, the SQL generated is:

UPDATE table SET column -> %(column_1)s=(column -> %(column_2)s)
[parameters: {'column_1': 'foo', 'column_2': 'bar'}]

The update fails with:

(psycopg2.errors.SyntaxError) syntax error at or near "->"
LINE 1: UPDATE table SET column -> 'foo'...

The expected syntax is UPDATE table SET column['foo']=column -> 'bar'.

Optional link from https://docs.sqlalchemy.org which documents the behavior that is expected

No response

SQLAlchemy Version in Use

2.0.21

DBAPI (i.e. the database driver)

psycopg2

Database Vendor and Major Version

PostgreSQL 16

Python Version

3.11

Operating system

Linux

To Reproduce

I changed a test and ran it for sqlite and it fails with another error.


    def test_extract_subobject(self, connection, metadata):
        sqlite_json = Table("json_test", metadata, Column("foo", sqlite.JSON))

        metadata.create_all(connection)

        value = {"json": {"foo": "bar"}}

        connection.execute(sqlite_json.insert(), dict(foo=value))

        eq_(
            connection.scalar(select(sqlite_json.c.foo["json"])), value["json"]
        )
+       connection.execute(
+           sqlite_json.update().values({sqlite_json.c.foo["otherjson"]: sqlite_json.c.foo["json"]})
+       )


I don't know the expected syntax for sqlite.

Error

Traceback (most recent call last):
  File "sqlalchemy/engine/base.py", line 1969, in _exec_single_context
    self.dialect.do_execute(
  File "sqlalchemy/engine/default.py", line 922, in do_execute
    cursor.execute(statement, parameters)
sqlite3.OperationalError: near "(": syntax error

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "_pytest/runner.py", line 341, in from_call
    result: Optional[TResult] = func()
  File "_pytest/runner.py", line 262, in <lambda>
    lambda: ihook(item=item, **kwds), when=when, reraise=reraise
  File "pluggy/_hooks.py", line 501, in __call__
    return self._hookexec(self.name, self._hookimpls.copy(), kwargs, firstresult)
  File "pluggy/_manager.py", line 119, in _hookexec
    return self._inner_hookexec(hook_name, methods, kwargs, firstresult)
  File "pluggy/_callers.py", line 181, in _multicall
    return outcome.get_result()
  File "pluggy/_result.py", line 99, in get_result
    raise exc.with_traceback(exc.__traceback__)
  File "pluggy/_callers.py", line 102, in _multicall
    res = hook_impl.function(*args)
  File "_pytest/runner.py", line 177, in pytest_runtest_call
    raise e
  File "_pytest/runner.py", line 169, in pytest_runtest_call
    item.runtest()
  File "_pytest/python.py", line 1792, in runtest
    self.ihook.pytest_pyfunc_call(pyfuncitem=self)
  File "pluggy/_hooks.py", line 501, in __call__
    return self._hookexec(self.name, self._hookimpls.copy(), kwargs, firstresult)
  File "pluggy/_manager.py", line 119, in _hookexec
    return self._inner_hookexec(hook_name, methods, kwargs, firstresult)
  File "pluggy/_callers.py", line 138, in _multicall
    raise exception.with_traceback(exception.__traceback__)
  File "pluggy/_callers.py", line 102, in _multicall
    res = hook_impl.function(*args)
  File "_pytest/python.py", line 194, in pytest_pyfunc_call
    result = testfunction(**testargs)
  File "sqlalchemy/test/dialect/test_sqlite.py", line 317, in test_extract_subobject
    connection.execute(sqlite_json.update().values({sqlite_json.c.foo["newkey"]: sqlite_json.c.foo["json"]}))
  File "sqlalchemy/engine/base.py", line 1416, in execute
    return meth(
  File "sqlalchemy/sql/elements.py", line 516, in _execute_on_connection
    return connection._execute_clauseelement(
  File "sqlalchemy/engine/base.py", line 1639, in _execute_clauseelement
    ret = self._execute_context(
  File "sqlalchemy/engine/base.py", line 1848, in _execute_context
    return self._exec_single_context(
  File "sqlalchemy/engine/base.py", line 1988, in _exec_single_context
    self._handle_dbapi_exception(
  File "sqlalchemy/engine/base.py", line 2343, in _handle_dbapi_exception
    raise sqlalchemy_exception.with_traceback(exc_info[2]) from e
  File "sqlalchemy/engine/base.py", line 1969, in _exec_single_context
    self.dialect.do_execute(
  File "sqlalchemy/engine/default.py", line 922, in do_execute
    cursor.execute(statement, parameters)
sqlalchemy.exc.OperationalError: (sqlite3.OperationalError) near "(": syntax error
[SQL: UPDATE json_test SET JSON_QUOTE(JSON_EXTRACT(foo, ?))=(JSON_QUOTE(JSON_EXTRACT(json_test.foo, ?)))]
[parameters: ('$."newkey"', '$."json"')]
(Background on this error at: https://sqlalche.me/e/20/e3q8)

Additional context

No response

[CaselIT]

Hi,

You are mentioning postgresql but the example is using sqlite.
Also the code provided is not runnable as is

[iurisilvio]

Sorry, you can reproduce it with:

from sqlalchemy import Column, MetaData, Table, create_engine
from sqlalchemy.dialects.sqlite import base as sqlite

engine = create_engine("sqlite://")
metadata = MetaData()


with engine.connect() as connection:
    sqlite_json = Table("json_test", metadata, Column("foo", sqlite.JSON))

    metadata.create_all(connection)

    value = {"json": {"foo": "bar"}}

    connection.execute(sqlite_json.insert(), dict(foo=value))

    connection.execute(
        sqlite_json.update().values({sqlite_json.c.foo["otherjson"]: sqlite_json.c.foo["json"]})
    )

It fails for Postgres too.

[CaselIT]

But they are you using sqlite.JSON in postgresql? That's a custom sqlite type

[iurisilvio]

For Postgres, I has able to use func.jsonb_set to make it work.

sqlite_json.update().values({sqlite_json.c.foo: sa.func.jsonb_set(sqlite_json.c.foo, '{"otherjson"}', sqlite_json.c.foo["json"])})

Works for me now, but I think the performance is not the same of writing only an attribute to the column.

[iurisilvio]

But they are you using sqlite.JSON in postgresql? That's a custom sqlite type

I'm using proper jsonb for postgres, I just wrote an example for sqlite now, reusing some test code.

[CaselIT]

Using postgresql json seems to work as expected.

import sqlalchemy as sa
from sqlalchemy.dialects import postgresql as pg

t = sa.table('foo', sa.column('bar', pg.JSON))
u = t.update().values({t.c.bar['foo']: t.c.bar['src']})

print(u.compile(dialect=pg.dialect()).string)

this prints a string that looks fine

'UPDATE foo SET bar -> %(bar_1)s=(foo.bar -> %(bar_2)s)'

[iurisilvio]

This SQL fails on Postgres 16 (at least for JSONB). I think it should be

UPDATE foo SET bar[%(bar_1)s]=(foo.bar -> %(bar_2)s)

[CaselIT]

Ok I though you wanted to generate something like 'UPDATE foo SET bar -> %(bar_1)s=(foo.bar -> %(bar_2)s)'.

Your proposed syntax will error too using json:

ERROR: cannot subscript type json because it does not support subscripting

[iurisilvio]

Here the SQL statements, ran in Postgres 16:

=> create table foo (
   id       integer primary key, 
   bar  jsonb not null default '{}'::jsonb
);
CREATE TABLE

=> update foo set bar -> 'a' = bar -> 'b';
ERROR:  syntax error at or near "->"
LINE 1: update foo set bar -> 'a' = bar -> 'b';
                           ^
=> update foo set bar['a'] = bar -> 'b';
UPDATE 0

[CaselIT]

Transformed this in to an use case, since it's not something that's currently implemented.

Do you know if the [] syntax is equivalent to the -> one and can be used instead of it in all places -> was used (plus the new ones?)

[iurisilvio]

I think for JSONB, they're equivalent. https://www.postgresql.org/docs/current/datatype-json.html#JSONB-SUBSCRIPTING

-- it works
update foo set bar['a'] = bar['b'];

For normal JSON, I did some tests here and subscription doesn't work.

create table foo
(
   id       integer primary key, 
   bar  json not null default '{}'::json 
);
insert into foo values (1, '{"foo": "bar"}');
=> select bar['foo'] from foo;
ERROR:  cannot subscript type json because it does not support subscripting

[CaselIT]

If so then 2.1 could maybe use this syntax in place of the current one for jsonb types.
I'm not sure if adding something like t.c.a_jsonb_col['some-key'].subscript_notation() or t.c.a_jsonb_col['some-key'].as_subscript() is a good option here.

@zzzeek do you have other ideas?

[zzzeek]

assuming I am seeing what's needed here, we only need the new syntax for the left hand side in an UPDATE SET clause, and on prior versions of PostgreSQL, nothing is supported anyway.

if that's all accurate then I'd seek to get the postgresql compiler to use this new syntax only for left hand side SET variables.

[CaselIT]

Any particular reason for not changing it in all cases, other than the current way works fine?

[zzzeek]

I like to keep changes minimal and backend agnostic if possible

[iurisilvio]

If that's all accurate

JSONB exists from Postgres 12, but JSONB subscripting is documented from Postgres 14, probably doesn't work before that. https://www.postgresql.org/docs/14/datatype-json.html#JSONB-SUBSCRIPTING

Needs some research to understand the details. Changing all cases may be problematic.

[zzzeek]

there is likely already infrastructure in the compiler to support doing special compilation for the "keys" of update.values(). pretty sure MySQL has something happening on that end.

[CaselIT]

Ok. I'll update the initial issue and target 2.0