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?