1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
|
from __future__ import annotations
from contextlib import asynccontextmanager
from datetime import date, datetime
from typing import AsyncIterator
from uuid import UUID
import anyio
from rich import get_console
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from sqlalchemy.orm import Mapped
from litestar.plugins.sqlalchemy import base, repository
console = get_console()
# the SQLAlchemy base includes a declarative model for you to use in your models.
# The `Base` class includes a `UUID` based primary key (`id`)
class Author(base.UUIDBase):
name: Mapped[str]
dob: Mapped[date]
dod: Mapped[date | None]
class AuthorRepository(repository.SQLAlchemyAsyncRepository[Author]):
"""Author repository."""
model_type = Author
engine = create_async_engine(
"sqlite+aiosqlite:///test.sqlite",
future=True,
)
session_factory = async_sessionmaker(engine, expire_on_commit=False)
# let's make a simple context manager as an example here.
@asynccontextmanager
async def repository_factory() -> AsyncIterator[AuthorRepository]:
async with session_factory() as db_session:
try:
yield AuthorRepository(session=db_session)
except Exception: # noqa: BLE001
await db_session.rollback()
else:
await db_session.commit()
async def create_author() -> Author:
async with repository_factory() as repo:
obj = await repo.add(
Author(
name="F. Scott Fitzgerald",
dob=datetime.strptime("1896-09-24", "%Y-%m-%d").date(),
)
)
console.print(f"Created Author record for {obj.name} with primary key {obj.id}.")
return obj
async def update_author(obj: Author) -> Author:
async with repository_factory() as repo:
obj = await repo.update(obj)
console.print(f"Updated Author record for {obj.name} with primary key {obj.id}.")
return obj
async def remove_author(id: UUID) -> Author:
async with repository_factory() as repo:
obj = await repo.delete(id)
console.print(f"Deleted Author record for {obj.name} with primary key {obj.id}.")
return obj
async def get_author_if_exists(id: UUID) -> Author | None:
async with repository_factory() as repo:
obj = await repo.get_one_or_none(id=id)
if obj is not None:
console.print(f"Found Author record for {obj.name} with primary key {obj.id}.")
else:
console.print(f"Could not find Author with primary key {id}.")
return obj
async def run_script() -> None:
"""Load data from a fixture."""
async with engine.begin() as conn:
await conn.run_sync(base.UUIDBase.metadata.create_all)
# 1) create a new Author record.
console.print("1) Adding a new record")
author = await create_author()
author_id = author.id
# 2) Let's update the Author record.
console.print("2) Updating a record.")
author.dod = datetime.strptime("1940-12-21", "%Y-%m-%d").date()
await update_author(author)
# 3) Let's delete the record we just created.
console.print("3) Removing a record.")
await remove_author(author_id)
# 4) Let's verify the record no longer exists.
console.print("4) Select one or none.")
_should_be_none = await get_author_if_exists(author_id)
if __name__ == "__main__":
anyio.run(run_script)
|
