Using the init plugin

In our example application, we’ve seen that we need to manage the database engine within the scope of the application’s lifespan, and the session within the scope of a request. This is a common pattern, and the SQLAlchemyInitPlugin plugin provides assistance for this.

In our latest update, we leverage two features of the plugin:

  1. The plugin will automatically create a database engine for us and manage it within the scope of the application’s lifespan.

  2. The plugin will automatically create a database session for us and manage it within the scope of a request.

We access the database session via dependency injection, using the db_session parameter.

Here’s the updated code:

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

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

The most notable difference is that we no longer need the db_connection() lifespan context manager - the plugin handles this for us. It also handles the creation of the tables in our database if we supply our metadata and set create_all=True when creating a SQLAlchemyAsyncConfig instance.

Additionally, we have a new db_session dependency available to us, which we use in our provide_transaction() dependency provider, instead of creating our own session.

Next steps

Next up, we’ll make one final change to our application, and then we’ll be recap!