Does asyncpg support jsonb_path_query?

[Lawouach]

• asyncpg version: 0.26.0
• PostgreSQL version: 14.5
• Do you use a PostgreSQL SaaS? If so, which? Can you reproduce
  the issue with a local PostgreSQL install?: Yes but that happens with a local PostgreSQL too
• Python version: 3.10.4
• Platform: x86_64 GNU/Linux
• Do you use pgbouncer?: no
• Did you install asyncpg with pip?: yes
• Can the issue be reproduced under both asyncio and
  uvloop?: yes

Does asyncpg supports the jsonb_path_query function? I get an error when trying to use it (whereas it works with psycopg2):

import asyncio
import json

from sqlalchemy import Table, Column, MetaData, func, select, create_engine
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.ext.asyncio import create_async_engine

async def run_async():
    engine = create_async_engine("postgresql+asyncpg://test:secret@localhost:5432/test", echo=True)
    metadata = MetaData(bind=engine)
    tt = Table('tt', metadata, Column('col', JSONB))

    async with engine.begin() as conn:
        await conn.run_sync(metadata.drop_all)
        await conn.run_sync(metadata.create_all)

        await conn.execute(
            tt.insert().values(
                col={
                    "provider": "A", "services": [
                        {"service": "B", "data": {"argument_name": "C"}}    
                    ]}
                )
            )

    async with engine.connect() as conn:
        s = select([func.jsonb_path_query(
            tt.c.col,
            "$ ? (@.provider == $provider).services[*] ? (@.service == $svc).data ? (@.argument_name == $argname)",
            json.dumps({"provider": "A", "svc": "B", "argname": "C"})
        )])
        r = await conn.execute(s)
        print(r.scalars().all())
        
    await engine.dispose()


def run_sync():
    engine = create_engine("postgresql://test:secret@localhost:5432/test", echo=True)
    metadata = MetaData(bind=engine)
    tt = Table('tt', metadata, Column('col', JSONB))
        
    metadata.drop_all(engine)
    metadata.create_all(engine)

    with engine.connect() as conn:
        conn.execute(
            tt.insert().values(
                col={
                    "provider": "A", "services": [
                        {"service": "B", "data": {"argument_name": "C"}}    
                    ]}
                )
            )

    with engine.connect() as conn:
        s = select([func.jsonb_path_query(
            tt.c.col,
            "$ ? (@.provider == $provider).services[*] ? (@.service == $svc).data ? (@.argument_name == $argname)",
            json.dumps({"provider": "A", "svc": "B", "argname": "C"})
        )])
        r = conn.execute(s)
        print(r.scalars().all())
        
    engine.dispose()

run_sync()
print("#####################################")
asyncio.run(run_async())

Results in:

2022-10-07 14:35:02,112 INFO sqlalchemy.engine.Engine SELECT jsonb_path_query(tt.col, %(jsonb_path_query_2)s, %(jsonb_path_query_3)s) AS jsonb_path_query_1 
FROM tt
2022-10-07 14:35:02,112 INFO sqlalchemy.engine.Engine [generated in 0.00047s] {'jsonb_path_query_2': '$ ? (@.provider == $provider).services[*] ? (@.service == $svc).data ? (@.argument_name == $argname)', 'jsonb_path_query_3': '{"provider": "A", "svc": "B", "argname": "C"}'}
[{'argument_name': 'C'}]

and

sqlalchemy.exc.ProgrammingError: (sqlalchemy.dialects.postgresql.asyncpg.ProgrammingError) <class 'asyncpg.exceptions.UndefinedFunctionError'>: function jsonb_path_query(jsonb, character varying, character varying) does not exist
HINT:  No function matches the given name and argument types. You might need to add explicit type casts.
[SQL: SELECT jsonb_path_query(tt.col, %s, %s) AS jsonb_path_query_1 
FROM tt]
[parameters: ('$ ? (@.provider == $provider).services[*] ? (@.service == $svc).data ? (@.argument_name == $argname)', '{"provider": "A", "svc": "B", "argname": "C"}')]

Is there something I'm missing to enable perhaps?

[Lawouach]

#635 didn't seem to help

[Lawouach]

Okay. See, sometimes opening an issue and your eyes open as well.

It was right there, the error was clear. I wasn't using the appropriate types for the arguments of the function. Here is what works:

from sqlalchemy import cast, literal_column
from sqlalchemy.dialects.postgresql import JSONB

s = select([func.jsonb_path_query(
            cast(tt.c.col, JSONB),
            literal_column("'$ ? (@.provider == $provider).services[*] ? (@.service == $svc).data ? (@.argument_name == $argname)'::jsonpath"),
            cast({"provider": "A", "svc": "B", "argname": "C"}, JSONB)
        )])

Note, that when you'll use SQLAlchemy 2, you'll probably prefer something like this:

from sqlalchemy import cast
from sqlalchemy.dialects.postgresql import JSONB, JSONPATH

s = select([func.jsonb_path_query(
            cast(tt.c.col, JSONB),
            cast("$ ? (@.provider == $provider).services[*] ? (@.service == $svc).data ? (@.argument_name == $argname)", JSONPATH),
            cast({"provider": "A", "svc": "B", "argname": "C"}, JSONB)
        )])

see sqlalchemy/sqlalchemy#8216