Как работать с пулом соединений asyncpg в агенте Faust?

Question

Я хочу, чтобы агент Faust записывал в таблицу PostgreSQL. Я хотел бы использовать пул соединений asyncpg, но не могу найти чистый способ вставить его в код инициализации приложения.

Answer 1

Просто добавьте приведенную ниже функцию в свое приложение faust.

class KafkaWorker(faust.App):
    def __init__(self, *args: List, **kwargs: Dict) -> None:
        self.broker : str = kwargs.pop('broker')
        self._db_pool = None

        super().__init__(*args, broker=KafkaWorker._broker_faust_string(self.broker), **kwargs)

    async def db_pool(self) -> Pool:
        ''' '''
        if not self._db_pool:
            logging.warning('kafka.db_pool initialization...')
            self._db_pool = await db.db_pool()
            logging.warning('kafka.db_pool initialization...done ✓')

        return self._db_pool

где db.db_pool

from os import environ

import asyncpg
from asyncpg.pool import Pool
from sqlalchemy.orm import scoped_session, sessionmaker
from sqlalchemy.ext.declarative import declarative_base


session = scoped_session(sessionmaker())
Base = declarative_base()


async def db_pool() -> Pool:
    return await asyncpg.create_pool(
        dsn=environ.get('DB_CNX_STRING', 'postgresql://postgres:postgres@postgres:5432/actions')
    )

а затем вы получаете доступ к нему как

@kafka.agent(actions_topic)
async def store_actions(actions: StreamT):
    async for action in actions:
        db_pool = await current_agent().app.db_pool()
        async with db_pool.acquire() as conn:
            try:
                yield await save_action(conn, action.to_representation())
            except StoreException:
                logger.exception(f'Error while inserting action in DB, continuing....')
            finally:
                yield action.id