Add documentation about how to use the async tools (session, etc)

[tiangolo]

Privileged issue

• I'm @tiangolo or he asked me directly to create an issue here.

Issue Content

Add documentation about how to use the async tools (session, etc).

Funding

• You can sponsor this specific effort via a Polar.sh pledge below
• We receive the pledge once the issue is completed & verified

[alvynabranches]

First of all we have to import from sqlalchemy which I dont like. It is a bit confusing. Hence I would like to import it from sqlmodel itself.

Like how we have

with Session(engine) as session:
    session.exec()

we should have

async with AsyncSession(engine) as session:
    await session.exec()

so it becomes easier for us to make sessions.

[MatsiukMykola]

from collections.abc import AsyncGenerator
from typing import Annotated, Callable

from fastapi import Depends
from sqlalchemy import create_engine
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker, Session

from core.config import settings

##############

# region 'Async Session'
engine = create_async_engine(
    settings.SQLALCHEMY_DATABASE_URI_ASYNC,
    future=True,
    echo=settings.LOCAL_DEV,
    hide_parameters=not settings.LOCAL_DEV,
    connect_args={
        # https://www.postgresql.org/docs/current/runtime-config.html
        "server_settings": {
            "application_name": f"{settings.PROJECT_NAME} {settings.VERSION} async",
            "jit": "off",
        },
    },
)

AsyncSessionFactory = sessionmaker(
    bind=engine,
    autoflush=False,
    expire_on_commit=False,
    class_=AsyncSession,
)


async def get_db() -> AsyncGenerator:
    yield AsyncSessionFactory


Session = Annotated[AsyncSession, Depends(get_db)]
# endregion


##############
@router.post(...)
async def some_function(
    session: Session)
    async with session() as db:  # noqa
         ...
        stmt = select(Model1)
        items = await db.execute(stmt)
        data = items.all()
        return data

asyncpg, works fine, except:

1. asyncpg cant execute sql-files with multistatements
2. asyncpg not shows params statements

[deshetti]

A full working example of the following is here: https://github.com/deshetti/sqlmodel-async-example
Opened a PR to add documentation to the docs: #633

from contextlib import asynccontextmanager
from typing import Optional

from fastapi import FastAPI
from pydantic import BaseModel
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker
from sqlmodel import Field, SQLModel

# Initialize FastAPI application
app = FastAPI()


# Define User model for SQLModel
class User(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    age: int


# Define UserCreate model for Pydantic validation
# For id field to not show up on the OpenAPI spec
class UserCreate(BaseModel):
    name: str
    age: int


# Database connection string
DATABASE_URL = "postgresql+asyncpg://postgres:postgres@localhost/sampledb"

# Create an asynchronous engine for the database
engine = create_async_engine(
    DATABASE_URL,
    echo=True,
    future=True,
    pool_size=20,
    max_overflow=20,
    pool_recycle=3600,
)


# Ayschronous Context manager for handling database sessions
@asynccontextmanager
async def get_session() -> AsyncSession:
    async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
    async with async_session() as session:
        yield session


# Function to create a new user in the database
async def create_user(user: User) -> User:
    async with get_session() as session:
        session.add(user)
        await session.commit()
        await session.refresh(user)
    return user


# Event handler for startup event of FastAPI application
@app.on_event("startup")
async def on_startup():
    async with engine.begin() as conn:
        # For SQLModel, this will create the tables (but won't drop existing ones)
        await conn.run_sync(SQLModel.metadata.create_all)


# Endpoint to create a new user
@app.post("/users/", response_model=User)
async def create_user_endpoint(user: UserCreate):
    db_user = User(**user.dict())
    result = await create_user(db_user)
    return result


# Main entry point of the application
if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)

[PookieBuns]

First of all we have to import from sqlalchemy which I dont like. It is a bit confusing. Hence I would like to import it from sqlmodel itself.

Like how we have

with Session(engine) as session:
    session.exec()

we should have

async with AsyncSession(engine) as session:
    await session.exec()

so it becomes easier for us to make sessions.

@tiangolo Can I work on this? Add a sqlmodel version of asyncsession as well as create_async_engine

[PookieBuns]

First of all we have to import from sqlalchemy which I dont like. It is a bit confusing. Hence I would like to import it from sqlmodel itself.

Like how we have

with Session(engine) as session:
    session.exec()

we should have

async with AsyncSession(engine) as session:
    await session.exec()

so it becomes easier for us to make sessions.

It looks like sqlmodel has already implemented its own asyncsession but is just not importable from the root directory. Currently I believe you need to import it through
from sqlmodel.ext.asyncio.session import AsyncSession

If we want to adhere to sqlalchemy model import structure it should be
from sqlmodel.ext.asyncio import AsyncSession

However, according to @alvynabranches the proposed solution is
from sqlmodel import AsyncSession

@tiangolo what do you think?

[1st]

@tiangolo when do you plan to add a fully-functional async support to this library?

We are using the the SQL Alchemy at the moment and recently rewrote our code to use async. But problem is in the async m2m and m2o relations, that we need to use .awaitable_attrs that our code a bit ugly. Look on this:

docs = await folder.awaitable_attrs.documents
docs_to_authors = {doc.id: await doc.awaitable_attrs.authors for doc in docs}

Do you plan to get rid of such awaitable_attrs syntax in your version and have something like this?

docs = await folder.documents
docs_to_authors = {doc.id: await doc.authors for doc in docs}

so that under the hood it will know that "documents" and "authors" are actually awaitable props and we need to use await keyword for them. Otherwise mypy needs to catch such a problems and notify us that we forgot to add await keyword for such m2m and m2o relations when trying to load them from DB.

[ryanrain2016]

@tiangolo when do you plan to add a fully-functional async support to this library?

We are using the the SQL Alchemy at the moment and recently rewrote our code to use async. But problem is in the async m2m and m2o relations, that we need to use .awaitable_attrs that our code a bit ugly. Look on this:

docs = await folder.awaitable_attrs.documents
docs_to_authors = {doc.id: await doc.awaitable_attrs.authors for doc in docs}

Do you plan to get rid of such awaitable_attrs syntax in your version and have something like this?

docs = await folder.documents
docs_to_authors = {doc.id: await doc.authors for doc in docs}

so that under the hood it will know that "documents" and "authors" are actually awaitable props and we need to use await keyword for them. Otherwise mypy needs to catch such a problems and notify us that we forgot to add await keyword for such m2m and m2o relations when trying to load them from DB.

If you want to fetch all related data, I suggest you to try selectinload

from sqlalchemy.orm import selectinload
...
(await session.exec(select(User).options(selectinload(User.roles))).all()

[tonyf]

When we were using SQLAlchemy with Pydantic, we introduced a mix-in to handle the fetching of awaitable_attrs.

class ModelInDBMixin(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    created_at: Optional[dt.datetime] = None
    updated_at: Optional[dt.datetime] = None

    @classmethod
    async def from_data(cls, obj: Base, **kwargs):
        fields = [
            field
            for field in dir(obj)
            if not field.startswith("_")
            and field != "awaitable_attrs"
            and field in cls.model_fields.keys()
        ]
        data = await asyncio.gather(
            *[get_attr_or_none(obj, field) for field in fields]
        )
        data_dict = {k: v for (k, v) in zip(fields, data) if v is not None}
        data_dict.update(kwargs)
        try:
            return cls.model_validate(data_dict)
        except Exception as e:
            raise e

Usage looking like User.from_data(db_obj) just in the same way you'd call from_orm in the original api.

Now, with SQLModel, we're monkeypatching sqlmodel_validate:

import asyncio
from collections.abc import Mapping
from typing import Any

import sqlmodel._compact  # type: ignore
from pydantic_core import PydanticUndefined as Undefined
from sqlmodel._compat import (ObjectWithUpdateWrapper, _TSQLModel,
                              is_table_model_class, partial_init)


async def get_attr_or_none(obj: Any, attr: str) -> Any:
    try:
        # make sure attr is not a hybrid_method
        if hasattr(getattr(obj, attr), "__call__"):
            return None
        return await getattr(obj.awaitable_attrs, attr)
    except Exception:
        return None

def sqlmodel_validate(
    cls: type[_TSQLModel],
    obj: Any,
    *,
    strict: bool | None = None,
    from_attributes: bool | None = None,
    context: Mapping[str, Any] | None = None,
    update: Mapping[str, Any] | None = None,
) -> _TSQLModel:
    asyncio_loop = asyncio.get_event_loop()
    if not is_table_model_class(cls):
        new_obj: _TSQLModel = cls.__new__(cls)
    else:
        # If table, create the new instance normally to make SQLAlchemy create
        # the _sa_instance_state attribute
        with partial_init():
            new_obj = cls()
    
    # SQLModel Override to get class SQLAlchemy __dict__ attributes and
    # set them back in after creating the object
    old_dict = new_obj.__dict__.copy()
    use_obj = obj
    
    # Handle awaitable attributes if they exist
    if hasattr(obj, 'awaitable_attrs'):
        # Get all model fields that might be awaitable
        fields = [
            field for field in cls.model_fields.keys()
            if not field.startswith("_")
        ]
        # Fetch all awaitable attributes
        awaited_values = asyncio_loop.run_until_complete(
            asyncio.gather(*[
                get_attr_or_none(obj, field) 
                for field in fields
            ])
        )
        # Create a dictionary of non-None awaited values
        awaited_dict = {
            k: v for k, v in zip(fields, awaited_values)
            if v is not None
        }
        # If obj is a dict, update it with awaited values
        if isinstance(obj, dict):
            use_obj = {**obj, **awaited_dict}
        else:
            # Create a new object that combines original and awaited attributes
            use_obj = ObjectWithUpdateWrapper(obj=obj, update=awaited_dict)
    
    # Apply any explicit updates
    if isinstance(use_obj, dict) and update:
        use_obj = {**use_obj, **update}
    elif update:
        use_obj = ObjectWithUpdateWrapper(obj=use_obj, update=update)

    cls.__pydantic_validator__.validate_python(
        use_obj,
        strict=strict,
        from_attributes=from_attributes,
        context=context,
        self_instance=new_obj,
    )
    
    # Capture fields set to restore it later
    fields_set = new_obj.__pydantic_fields_set__.copy()
    if not is_table_model_class(cls):
        # If not table, normal Pydantic code, set __dict__
        new_obj.__dict__ = {**old_dict, **new_obj.__dict__}
    else:
        # Do not set __dict__, instead use setattr to trigger SQLAlchemy
        # instrumentation
        for key, value in {**old_dict, **new_obj.__dict__}.items():
            setattr(new_obj, key, value)
    
    # Restore fields set
    object.__setattr__(new_obj, "__pydantic_fields_set__", fields_set)
    
    # Get and set any relationship objects
    if is_table_model_class(cls):
        for key in new_obj.__sqlmodel_relationships__:
            value = getattr(use_obj, key, Undefined)
            if value is not Undefined:
                setattr(new_obj, key, value)
    
    return new_obj


sqlmodel._compat.sqlmodel_validate = sqlmodel_validate

[prerna-pramod-gh]

Here's a working example showing how to use async tools like AsyncSession with SQLModel:

from sqlmodel import SQLModel, Field
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker
import asyncio

# Define your model
class Hero(SQLModel, table=True):
    id: int | None = Field(default=None, primary_key=True)
    name: str
    secret_name: str

# Setup async engine and sessionmaker
DATABASE_URL = "sqlite+aiosqlite:///./test.db"
engine = create_async_engine(DATABASE_URL, echo=True)
async_session = sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)

# Create database tables
async def init_db():
    async with engine.begin() as conn:
        await conn.run_sync(SQLModel.metadata.create_all)

# Insert a record using async session
async def create_hero():
    async with async_session() as session:
        async with session.begin():
            hero = Hero(name="Deadpond", secret_name="Dive Wilson")
            session.add(hero)

# Run setup and example
async def main():
    await init_db()
    await create_hero()

asyncio.run(main())