Support async ORM - piccolo

[virajkanwade]

No description provided.

[virajkanwade]

@apotterri

Piccolo currently supports
SQLiteEngine
PostgresEngine
CockroachEngine

CockroachEngine inherits from PostgresEngine.

https://piccolo-orm.readthedocs.io/en/latest/piccolo/engines/index.html#engine-types

from _appmap import wrapt
from _appmap.env import Env

from piccolo.engine.postgres import PostgresEngine
from piccolo.engine.sqlite import SQLiteEngine


async def install_extension_run_querystring(wrapped, self_obj, args, kwargs):
    print('> inside Postgres run_querystring')
    print('  query_id: ', self_obj.query_id + 1)  # this will be incremented inside the function
    print('  args:', args)
    print('  kwargs:', kwargs)

    query, query_args = args[0].compile_string(
            engine_type=self_obj.engine_type
    )

    print('    querystring:', args[0].__str__())
    print('      query: ', query)
    print('      query_args: ', query_args)

    response = await wrapped(*args, **kwargs)

    print('  query_id: ', self_obj.query_id)  # this will be incremented inside the function
    print('  response:', response)

    return response


async def install_extension_run_ddl(wrapped, self_obj, args, kwargs):
    print('> inside Postgres run_ddl')
    print('  query_id: ', self_obj.query_id + 1)  # this will be incremented inside the function
    print('  args:', args)
    print('  kwargs:', kwargs)

    print('    querystring:', args[0])

    response = await wrapped(*args, **kwargs)

    print('  query_id: ', self_obj.query_id)  # this will be incremented inside the function
    print('  response:', response)

    return response


if Env.current.enabled:
    PostgresEngine.run_querystring = wrapt.wrap_function_wrapper("piccolo.engine.postgres", "PostgresEngine.run_querystring", install_extension_run_querystring)
    PostgresEngine.run_ddl = wrapt.wrap_function_wrapper("piccolo.engine.postgres", "PostgresEngine.run_ddl", install_extension_run_ddl)

    SQLiteEngine.run_querystring = wrapt.wrap_function_wrapper("piccolo.engine.postgres", "SQLiteEngine.run_querystring", install_extension_run_querystring)
    SQLiteEngine.run_ddl = wrapt.wrap_function_wrapper("piccolo.engine.postgres", "SQLiteEngine.run_ddl", install_extension_run_ddl)