Make _wrap_json public for every connector?

[denisSurkov]

Sometimes I need to extend connector. I cannot inherit from existing connector, like PsycopgConnector. But I want to reuse existing simple _wrap_json methods.

Is it okey to ask to make them as @classmethod and public (without _ before function name)?

---

Why I need to extend existing connector
I want to reuse same sqlalchemy open session with transaction. So I make something like this:

from typing import LiteralString, Any

from procrastinate import PsycopgConnector
from procrastinate.connector import BaseAsyncConnector, Pool
from sqlalchemy.ext.asyncio import AsyncSession
import psycopg.types.json


class ReuseSqlalchemyConnector(BaseAsyncConnector):
    """
    With PsycopgConnector:

    transaction_id 1:
        begin
        ...
        commit
    transaction_id 2:
        begin
        select unset(procrastinate_defer)
        commit

    With ReuseSqlalchemyConnector:

    transaction_id 1:
        begin
        ...
        select unset(procrastinate_defer)
        commit
    """

    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    def _wrap_value(self, value: Any) -> Any:
        if isinstance(value, dict):
            return psycopg.types.json.Jsonb(value)
        elif isinstance(value, list):
            return [self._wrap_value(item) for item in value]
        elif isinstance(value, tuple):
            return tuple([self._wrap_value(item) for item in value])
        else:
            return value

    def _wrap_json(self, arguments: dict[str, Any]):
        return {key: self._wrap_value(value) for key, value in arguments.items()} # would be awesome to just PsycopgConnector.wrap_json(arguments)

    async def open_async(self, pool: Pool | None = None) -> None:
        ...

    async def close_async(self, pool: Pool | None = None) -> None:
        ...

    async def execute_query_one_async(
        self, query: LiteralString, **arguments: Any
    ) -> dict[str, Any]:
        connection = await self.session.connection()
        result = await connection.exec_driver_sql(query, arguments)
        mapping = result.mappings()
        return mapping.fetchone()  # type: ignore


    async def execute_query_all_async(
        self, query: LiteralString, **arguments: Any
    ) -> list[dict[str, Any]]:
        connection = await self.session.connection()
        result = await connection.exec_driver_sql(query, parameters=self._wrap_json(arguments))
        mapping = result.mappings()
        return mapping.fetchall()

And then call it like this:

Lock = asyncio.Lock()

async with Lock:
    with app.replace_connector(ReuseSqlalchemyConnector(session)):
           ...

Any tips how to make ReuseConnector easier to implement are welcome. Thank you.

[medihack]

Sounds ok to me. We could even move them out of the class into normal functions (in all connectors) and add some unit tests. Are you open to creating a PR?

[denisSurkov]

Are you open to creating a PR?

Yes, I will do it 👍