Using the serialization plugin

Our next improvement is to leverage the SQLAlchemySerializationPlugin so that we can receive and return our SQLAlchemy models directly to and from our handlers.

Here’s the code:

  1from collections.abc import AsyncGenerator
  2from contextlib import asynccontextmanager
  3
  4from advanced_alchemy.extensions.litestar import SQLAlchemySerializationPlugin
  5from sqlalchemy import select
  6from sqlalchemy.exc import IntegrityError, NoResultFound
  7from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
  8from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
  9
 10from litestar import Litestar, get, post, put
 11from litestar.datastructures import State
 12from litestar.di import NamedDependency
 13from litestar.exceptions import ClientException, NotFoundException
 14from litestar.params import FromPath, FromQuery
 15from litestar.status_codes import HTTP_409_CONFLICT
 16
 17
 18class Base(DeclarativeBase): ...
 19
 20
 21class TodoItem(Base):
 22    __tablename__ = "todo_items"
 23
 24    title: Mapped[str] = mapped_column(primary_key=True)
 25    done: Mapped[bool]
 26
 27
 28@asynccontextmanager
 29async def db_connection(app: Litestar) -> AsyncGenerator[None, None]:
 30    engine = getattr(app.state, "engine", None)
 31    if engine is None:
 32        engine = create_async_engine("sqlite+aiosqlite:///todo.sqlite", echo=True)
 33        app.state.engine = engine
 34
 35    async with engine.begin() as conn:
 36        await conn.run_sync(Base.metadata.create_all)
 37
 38    try:
 39        yield
 40    finally:
 41        await engine.dispose()
 42
 43
 44sessionmaker = async_sessionmaker(expire_on_commit=False)
 45
 46
 47async def provide_transaction(state: State) -> AsyncGenerator[AsyncSession, None]:
 48    async with sessionmaker(bind=state.engine) as session:
 49        try:
 50            async with session.begin():
 51                yield session
 52        except IntegrityError as exc:
 53            raise ClientException(
 54                status_code=HTTP_409_CONFLICT,
 55                detail=str(exc),
 56            ) from exc
 57
 58
 59async def get_todo_by_title(todo_name: str, session: AsyncSession) -> TodoItem:
 60    query = select(TodoItem).where(TodoItem.title == todo_name)
 61    result = await session.execute(query)
 62    try:
 63        return result.scalar_one()
 64    except NoResultFound as e:
 65        raise NotFoundException(detail=f"TODO {todo_name!r} not found") from e
 66
 67
 68async def get_todo_list(done: bool | None, session: AsyncSession) -> list[TodoItem]:
 69    query = select(TodoItem)
 70    if done is not None:
 71        query = query.where(TodoItem.done.is_(done))
 72
 73    result = await session.execute(query)
 74    return list(result.scalars().all())
 75
 76
 77@get("/")
 78async def get_list(transaction: NamedDependency[AsyncSession], done: FromQuery[bool | None] = None) -> list[TodoItem]:
 79    return await get_todo_list(done, transaction)
 80
 81
 82@post("/")
 83async def add_item(data: TodoItem, transaction: NamedDependency[AsyncSession]) -> TodoItem:
 84    transaction.add(data)
 85    return data
 86
 87
 88@put("/{item_title:str}")
 89async def update_item(
 90    item_title: FromPath[str], data: TodoItem, transaction: NamedDependency[AsyncSession]
 91) -> TodoItem:
 92    todo_item = await get_todo_by_title(item_title, transaction)
 93    todo_item.title = data.title
 94    todo_item.done = data.done
 95    return todo_item
 96
 97
 98app = Litestar(
 99    [get_list, add_item, update_item],
100    dependencies={"transaction": provide_transaction},
101    lifespan=[db_connection],
102    plugins=[SQLAlchemySerializationPlugin()],
103)

  1from collections.abc import AsyncGenerator
  2from contextlib import asynccontextmanager
  3
  4from advanced_alchemy.extensions.litestar import SQLAlchemySerializationPlugin
  5from sqlalchemy import select
  6from sqlalchemy.exc import IntegrityError, NoResultFound
  7from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
  8from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
  9
 10from litestar import Litestar, get, post, put
 11from litestar.datastructures import State
 12from litestar.di import NamedDependency
 13from litestar.exceptions import ClientException, NotFoundException
 14from litestar.params import FromPath, FromQuery
 15from litestar.status_codes import HTTP_409_CONFLICT
 16
 17
 18class Base(DeclarativeBase): ...
 19
 20
 21class TodoItem(Base):
 22    __tablename__ = "todo_items"
 23
 24    title: Mapped[str] = mapped_column(primary_key=True)
 25    done: Mapped[bool]
 26
 27
 28@asynccontextmanager
 29async def db_connection(app: Litestar) -> AsyncGenerator[None]:
 30    engine = getattr(app.state, "engine", None)
 31    if engine is None:
 32        engine = create_async_engine("sqlite+aiosqlite:///todo.sqlite", echo=True)
 33        app.state.engine = engine
 34
 35    async with engine.begin() as conn:
 36        await conn.run_sync(Base.metadata.create_all)
 37
 38    try:
 39        yield
 40    finally:
 41        await engine.dispose()
 42
 43
 44sessionmaker = async_sessionmaker(expire_on_commit=False)
 45
 46
 47async def provide_transaction(state: State) -> AsyncGenerator[AsyncSession]:
 48    async with sessionmaker(bind=state.engine) as session:
 49        try:
 50            async with session.begin():
 51                yield session
 52        except IntegrityError as exc:
 53            raise ClientException(
 54                status_code=HTTP_409_CONFLICT,
 55                detail=str(exc),
 56            ) from exc
 57
 58
 59async def get_todo_by_title(todo_name: str, session: AsyncSession) -> TodoItem:
 60    query = select(TodoItem).where(TodoItem.title == todo_name)
 61    result = await session.execute(query)
 62    try:
 63        return result.scalar_one()
 64    except NoResultFound as e:
 65        raise NotFoundException(detail=f"TODO {todo_name!r} not found") from e
 66
 67
 68async def get_todo_list(done: bool | None, session: AsyncSession) -> list[TodoItem]:
 69    query = select(TodoItem)
 70    if done is not None:
 71        query = query.where(TodoItem.done.is_(done))
 72
 73    result = await session.execute(query)
 74    return list(result.scalars().all())
 75
 76
 77@get("/")
 78async def get_list(transaction: NamedDependency[AsyncSession], done: FromQuery[bool | None] = None) -> list[TodoItem]:
 79    return await get_todo_list(done, transaction)
 80
 81
 82@post("/")
 83async def add_item(data: TodoItem, transaction: NamedDependency[AsyncSession]) -> TodoItem:
 84    transaction.add(data)
 85    return data
 86
 87
 88@put("/{item_title:str}")
 89async def update_item(
 90    item_title: FromPath[str], data: TodoItem, transaction: NamedDependency[AsyncSession]
 91) -> TodoItem:
 92    todo_item = await get_todo_by_title(item_title, transaction)
 93    todo_item.title = data.title
 94    todo_item.done = data.done
 95    return todo_item
 96
 97
 98app = Litestar(
 99    [get_list, add_item, update_item],
100    dependencies={"transaction": provide_transaction},
101    lifespan=[db_connection],
102    plugins=[SQLAlchemySerializationPlugin()],
103)

We’ve simply imported the plugin and added it to our app’s plugins list, and now we can receive and return our SQLAlchemy data models directly to and from our handler.

We’ve also been able to remove the TodoType and TodoCollectionType aliases, and the serialize_todo() function, making the implementation even more concise.

Compare handlers before and after Serialization Plugin

Once more, let’s compare the sets of application handlers before and after our refactoring:

 1from advanced_alchemy.extensions.litestar import SQLAlchemySerializationPlugin
 2from sqlalchemy.ext.asyncio import AsyncSession
 3
 4from litestar import Litestar, get, post, put
 5from litestar.di import NamedDependency
 6from litestar.params import FromPath, FromQuery
 7
 8
 9@get("/")
10async def get_list(transaction: NamedDependency[AsyncSession], done: FromQuery[bool | None] = None) -> list[TodoItem]:
11    return await get_todo_list(done, transaction)
12
13
14@post("/")
15async def add_item(data: TodoItem, transaction: NamedDependency[AsyncSession]) -> TodoItem:
16    transaction.add(data)
17    return data
18
19
20@put("/{item_title:str}")
21async def update_item(
22    item_title: FromPath[str], data: TodoItem, transaction: NamedDependency[AsyncSession]
23) -> TodoItem:
24    todo_item = await get_todo_by_title(item_title, transaction)
25    todo_item.title = data.title
26    todo_item.done = data.done
27    return todo_item
28
29
30app = Litestar(
31    [get_list, add_item, update_item],
32    dependencies={"transaction": provide_transaction},
33    lifespan=[db_connection],
34    plugins=[SQLAlchemySerializationPlugin()],
35)

 1from collections.abc import AsyncGenerator, Sequence
 2from contextlib import asynccontextmanager
 3from typing import Any
 4
 5from sqlalchemy import select
 6from sqlalchemy.exc import IntegrityError, NoResultFound
 7from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
 8from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
 9
10from litestar import Litestar, get, post, put
11from litestar.datastructures import State
12from litestar.exceptions import ClientException, NotFoundException
13from litestar.params import FromPath, FromQuery
14    return result.scalars().all()
15
16
17@get("/")
18async def get_list(state: State, done: FromQuery[bool | None] = None) -> TodoCollectionType:
19    async with sessionmaker(bind=state.engine) as session:
20        return [serialize_todo(todo) for todo in await get_todo_list(done, session)]
21
22
23@post("/")
24async def add_item(data: TodoType, state: State) -> TodoType:
25    new_todo = TodoItem(title=data["title"], done=data["done"])
26    async with sessionmaker(bind=state.engine) as session:
27        try:
28            async with session.begin():
29                session.add(new_todo)
30        except IntegrityError as e:
31            raise ClientException(
32                status_code=HTTP_409_CONFLICT,
33                detail=f"TODO {new_todo.title!r} already exists",
34            ) from e
35
36    return serialize_todo(new_todo)
37
38
39@put("/{item_title:str}")
40async def update_item(item_title: FromPath[str], data: TodoType, state: State) -> TodoType:
41    async with sessionmaker(bind=state.engine) as session, session.begin():
42        todo_item = await get_todo_by_title(item_title, session)
43        todo_item.title = data["title"]
44        todo_item.done = data["done"]
45    return serialize_todo(todo_item)

Very nice! But, we can do better.

Next steps

In our application, we’ve had to build a bit of scaffolding to integrate SQLAlchemy with our application. We’ve had to define the db_connection() lifespan context manager, and the provide_transaction() dependency provider.

Next we’ll look at how the SQLAlchemyInitPlugin can help us.