Function init_async_engine

trinity/buffer/schema/sql_schema.py:176–203

Create an async SQLAlchemy engine and table classes. Returns: For task schema: (async_engine, table_cls) For experience/sft/dpo schema: (async_engine, meta_cls, blob_cls)

(db_url: str, table_name: str, schema_type: Optional[str])

async def init_async_engine(db_url: str, table_name: str, schema_type: Optional[str]) -> Tuple:
    """Create an async SQLAlchemy engine and table classes.

    Returns:
        For task schema: (async_engine, table_cls)
        For experience/sft/dpo schema: (async_engine, meta_cls, blob_cls)
    """
    from trinity.buffer.utils import to_async_url

    logger = get_logger(__name__)
    async_url = to_async_url(db_url)
    engine = create_async_engine(async_url, pool_pre_ping=True)

    if schema_type is None:
        schema_type = "task"

    classes = _create_table_classes(table_name, schema_type)

    try:
        async with engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
        logger.info(f"Created async tables for {table_name} (schema={schema_type}).")
    except OperationalError:
        logger.warning(f"Failed to create async tables for {table_name}, assuming they exist.")

    if schema_type == "task":
        return engine, classes[0]
    return engine, classes[0], classes[1]