procrastinate-org / procrastinate

PostgreSQL-based Task Queue for Python
https://procrastinate.readthedocs.io/
MIT License
834 stars 52 forks source link

Easy change of default connection pool class #1022

Closed ewjoachim closed 5 months ago

ewjoachim commented 5 months ago

Discussed in https://github.com/procrastinate-org/procrastinate/discussions/1011

Originally posted by **stinovlas** March 22, 2024 `PsycopgConnector` currently uses `psycopg_pool.AsyncConnectionPool` by default. It's possible to pass `pool` argument to `App.open_async` which overrides this, but it's not as easy to pass it to procrastinate `worker` / `defer` CLI commands. I ended up with overriding `open_async` method of `PsycopgConnector` like this: ```python import psycopg_pool from procrastinate import PsycopgConnector class NullPoolPsycopgConnector(PsycopgConnector): async def open_async(self, pool: psycopg_pool.AsyncConnectionPool | None = None) -> None: """Override open_async to use null connection pool by default.""" if pool is None: pool = psycopg_pool.AsyncNullConnectionPool(**self._pool_args) return await super().open_async(pool) ``` This works, but it doesn't feel very clean. I'd rather have `connection_pool_cls` kwarg on `PsycopgConnector` where I could pass `AsyncNullConnectionPool` (or any other class inheriting from `AsyncConnectionPool`). If specified, this class would be used instead of `AsyncConnectionPool` in `PsycopgConnector._create_pool`. It would be used like this: ```python app = TasksApp( connector=PsycopgConnector( connection_pool_cls=psycopg_pool.AsyncNullConnectionPool, json_dumps=orjson.dumps, json_loads=orjson.loads, **connector_options ), import_paths=[...], ) ``` What do you think?

-> we'll probably do a pool_factory rather than a class, but otherwise, let's roll with it!