How to properly wrap an Iterator to create an async generator?

Viewed 15

Context

I have a bunch of data Extractor all of which implement a job method. Most extractors rely on async clients such as aiohttp. For the others, I used the following pattern to wrap blocking code (such as pandas so usefull read_sql)

class SQLExtractor(Extractor):
    sessions: list[Connection]
    @staticmethod
    async def _read_sql_async(stmt, con):
        """
        Async wrapper for pandas.read_sql
        """
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, pd.read_sql, stmt, con)

    async def job(self, item: str, worker_id: int):
        """
        item: name of the table to ingest
        """
        params: TableMetadata = self.conf.get(f"tables.{item}")
        df = await self._read_sql_async(
            params.origin,
            con=self.sessions[worker_id],
        )
        return df

Current need

job are now generators, allowing them to return multiple dataframes. Applied to my SQLExtractor, I want to ingest big tables in multiple chunks. Provided with a chunksize arg, read_sql returns an Iterator[DataFrame]. It is thus possible to simply say

conn = engine.connect().execution_options(stream_results=True)
for chunk_dataframe in pd.read_sql(
        "SELECT * FROM users", conn, chunksize=1000):
    print(f"Got dataframe w/{len(chunk_dataframe)} row

However, the implementation of this feature is completely blocking and rely on blocking clients. Sure I could simply

class SQLExtractor(Extractor):
    sessions: list[Connection]
    async def job(self, item: str, worker_id: int):
    """
    item: name of the table to ingest
    """
    params: SQLTableMetadata = self.conf.get(f"tables.{item}")
    batch = 0
    for d in pd.read_sql(sql=params.origin, con=self.sessions[worker_id], chunksize=params.offset):
        batch += 1
        try:
            yield e
        except Exception as e:
            yield e

but 1) my sessions are unable to work in parallel and 2) send only send batches at the end of the iteration rather than 'as the chunks arrive' (so to speak). Is there a way I could wrap it or should I necessarily rewrite the loop to

df = await self._read_sql_async(f"{params.origin} LIMIT 0,10000", con=self.sessions[worker_id])
df = await self._read_sql_async(f"{params.origin} LIMIT 10000,20000", con=self.sessions[worker_id])
....

and yield the chunks?

0 Answers
Related