MagicStack / asyncpg

A fast PostgreSQL Database Client Library for Python/asyncio.
Apache License 2.0
6.99k stars 404 forks source link

Does asyncpg support jsonb_path_query? #963

Closed Lawouach closed 2 years ago

Lawouach commented 2 years ago

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 commented 2 years ago

635 didn't seem to help

Lawouach commented 2 years ago

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 https://github.com/sqlalchemy/sqlalchemy/issues/8216