<@ULVA73B9P> what is the best practice of using ai...
# ask-marvin
h
@Marvin what is the best practice of using aiomysql in Prefect flow?
m
thought for 102.1 seconds
Best practice in Prefect 3.x: use async flows/tasks and manage
aiomysql
connections with a pool
. Recommended approach - Make your flow
async def
- Make DB tasks
async def
- Create/use an
aiomysql
pool, not one-off connections - Open the pool once per flow run when possible, and close it cleanly - Use Prefect task retries/timeouts for transient DB failures - Limit concurrency so you don’t overwhelm MySQL - Don’t call
asyncio.run()
inside a Prefect task/flow Example pattern:
Copy code
python
from prefect import flow, task
import aiomysql
import asyncio

@task(retries=3, retry_delay_seconds=5, timeout_seconds=30)
async def fetch_rows(pool, query: str):
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(query)
            return await cur.fetchall()

@flow
async def my_flow():
    pool = await aiomysql.create_pool(
        host="localhost",
        port=3306,
        user="user",
        password="password",
        db="mydb",
        minsize=1,
        maxsize=10,
        autocommit=False,
    )

    try:
        rows = await fetch_rows(pool, "SELECT * FROM my_table")
        return rows
    finally:
        pool.close()
        await pool.wait_closed()
Why this is best practice - Async tasks fit DB I/O well: Prefect 3 supports async flows/tasks cleanly. - Pooling is important: avoids reconnect overhead and helps control DB pressure. - Lifecycle matters: always close the pool in
finally
. - Retries/timeouts: useful for transient connection drops or lock waits. - Concurrency control: if you fan out many queries, keep it under your DB/pool limits. For concurrent queries:
Copy code
python
from prefect import flow, task
import aiomysql
import asyncio

@task
async def fetch_one(pool, user_id: int):
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute("SELECT * FROM users WHERE id=%s", (user_id,))
            return await cur.fetchone()

@flow
async def concurrent_flow():
    pool = await aiomysql.create_pool(
        host="localhost",
        user="user",
        password="password",
        db="mydb",
        minsize=1,
        maxsize=5,
    )

    try:
        results = await asyncio.gather(
            *(fetch_one(pool, i) for i in range(1, 6))
        )
        return results
    finally:
        pool.close()
        await pool.wait_closed()
Important caveats - Do not create a brand new connection in every tiny task unless you really need isolation. - Do not call
asyncio.run()
inside tasks; Prefect already manages the event loop. - Be careful mixing sync and async DB code in the same flow; prefer one model. - If transactions matter, explicitly
commit()
/
rollback()
. Transactional example:
Copy code
python
from prefect import task

@task
async def insert_row(pool, name: str):
    async with pool.acquire() as conn:
        try:
            async with conn.cursor() as cur:
                await cur.execute(
                    "INSERT INTO users(name) VALUES(%s)",
                    (name,)
                )
            await conn.commit()
        except Exception:
            await conn.rollback()
            raise
If you want an even more robust/idiomatic option, consider SQLAlchemy async with
mysql+aiomysql
, since it gives you better pooling and session management. Prefect also has a SQLAlchemy integration: https://docs.prefect.io/integrations/prefect-sqlalchemy General Prefect docs: https://docs.prefect.io If you want, I can give you: 1. a pure
aiomysql
production template
for Prefect, or 2. the SQLAlchemy async version, which is usually what I’d recommend for larger projects.