How to work with asyncpg connection pool in Faust agent?

Question

I want Faust agent to write to PostgreSQL table. I'd like to use asyncpg connection pool but cannot find a clean way to inject it into the app initialization code.

Answer 1

Simply add the below function to your faust App

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

where db.db_pool is

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')
    )

and then you access it as

@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