|
| 1 | +# The Database file (db.py) |
| 2 | + |
| 3 | +The database setup is fairly straightforward, we will go through it line by |
| 4 | +line. |
| 5 | + |
| 6 | +## Imports |
| 7 | + |
| 8 | +```python linenums="1" |
| 9 | +"""Set up the database connection and session.""" "" |
| 10 | +from collections.abc import AsyncGenerator |
| 11 | +from typing import Any |
| 12 | + |
| 13 | +from sqlalchemy import MetaData |
| 14 | +from sqlalchemy.ext.asyncio import ( |
| 15 | + AsyncSession, |
| 16 | + async_sessionmaker, |
| 17 | + create_async_engine, |
| 18 | +) |
| 19 | +from sqlalchemy.orm import DeclarativeBase |
| 20 | +``` |
| 21 | + |
| 22 | +Lines 1 to 11 are the imports. The only thing to note here is that we are using |
| 23 | +the `AsyncGenerator` type hint for the `get_db` function. This is because we are |
| 24 | +using the `yield` keyword in the function, which makes it a generator. The |
| 25 | +`AsyncGenerator` type hint is a special type hint that is used for asynchronous |
| 26 | +generators. |
| 27 | + |
| 28 | +## Database Connection String |
| 29 | + |
| 30 | +```python linenums="13" |
| 31 | +DATABASE_URL = "postgresql+asyncpg://postgres:postgres@localhost/postgres" |
| 32 | +# DATABASE_URL = "sqlite+aiosqlite:///./test.db" |
| 33 | +``` |
| 34 | + |
| 35 | +We set a variable to be used later which contains the database URL. We are using |
| 36 | +PostgreSQL, but you can use any database that SQLAlchemy supports. The commented |
| 37 | +out line is for SQLite, which is a good choice for testing. You can comment out |
| 38 | +the PostgreSQL line (**13**)and uncomment the SQLite line (**14**)to use SQLite |
| 39 | +instead. |
| 40 | + |
| 41 | +This is a basic connection string, in reality you would want to use environment |
| 42 | +variables to store the user/password and database name. |
| 43 | + |
| 44 | +## The Base Class |
| 45 | + |
| 46 | +```python linenums="20" |
| 47 | +class Base(DeclarativeBase): |
| 48 | + """Base class for SQLAlchemy models. |
| 49 | +
|
| 50 | + All other models should inherit from this class. |
| 51 | + """ |
| 52 | + |
| 53 | + metadata = MetaData( |
| 54 | + naming_convention={ |
| 55 | + "ix": "ix_%(column_0_label)s", |
| 56 | + "uq": "uq_%(table_name)s_%(column_0_name)s", |
| 57 | + "ck": "ck_%(table_name)s_%(constraint_name)s", |
| 58 | + "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s", |
| 59 | + "pk": "pk_%(table_name)s", |
| 60 | + } |
| 61 | + ) |
| 62 | +``` |
| 63 | + |
| 64 | +This takes the `DeclarativeBase` class from SQLAlchemy and adds a `metadata` |
| 65 | +attribute to it. This is used to define the naming convention for the database |
| 66 | +tables. **This is not required**, but it is a good idea to set this up for |
| 67 | +consistency. |
| 68 | + |
| 69 | +We will use this class as the base class for all of our future models. |
| 70 | + |
| 71 | +## The database engine and session |
| 72 | + |
| 73 | +```python linenums="37" |
| 74 | +async_engine = create_async_engine(DATABASE_URL, echo=False) |
| 75 | +``` |
| 76 | + |
| 77 | +Here on line 37 we create the database engine. The `create_async_engine` |
| 78 | +function takes the database URL and returns an engine, the connection to the |
| 79 | +database. The `echo` parameter is set to `False` to prevent SQLAlchemy from |
| 80 | +outputting all of the SQL commands it is running. Note that it uses the |
| 81 | +`DATABASE_URL` variable we set earlier. |
| 82 | + |
| 83 | +```python linenums="38" |
| 84 | +async_session = async_sessionmaker(async_engine, expire_on_commit=False) |
| 85 | +``` |
| 86 | + |
| 87 | +Next, we create the session. The `async_sessionmaker` function takes the engine |
| 88 | +and returns a session. The `expire_on_commit` parameter is set to `False` to |
| 89 | +prevent SQLAlchemy from expiring objects on commit. This is required for |
| 90 | +`asyncpg` to work properly. |
| 91 | + |
| 92 | +We will NOT use this session directly, instead we will use the `get_db` function |
| 93 | +below to get and release a session. |
| 94 | + |
| 95 | +## The `get_db()` function |
| 96 | + |
| 97 | +```python linenums="41" |
| 98 | +async def get_db() -> AsyncGenerator[AsyncSession, Any]: |
| 99 | + """Get a database session. |
| 100 | +
|
| 101 | + To be used for dependency injection. |
| 102 | + """ |
| 103 | + async with async_session() as session, session.begin(): |
| 104 | + yield session |
| 105 | +``` |
| 106 | + |
| 107 | +This function is used to get a database session as a generator function. This |
| 108 | +function is used for dependency injection, which is a fancy way of saying that |
| 109 | +we will use it to pass the database session to other functions. Since we have |
| 110 | +used the `with` statement, the session will be automatically closed (and data |
| 111 | +comitted) when the function returns, usually after the related route is |
| 112 | +complete. |
| 113 | + |
| 114 | +!!! note |
| 115 | + Note that in line **46** we are using a combined `with` statement. This |
| 116 | + is a shortcut for using two nested `with` statements, one for the |
| 117 | + `async_session` and one for the `session.begin()`. |
| 118 | + |
| 119 | +## The `init_models()` function |
| 120 | + |
| 121 | +This function is used to create the database tables. It is called by the |
| 122 | +`lifespan()` function at startup. |
| 123 | + |
| 124 | +!!! note |
| 125 | + This function is only used in our demo, in a real application you would |
| 126 | + use a migration tool like |
| 127 | + [Alembic](https://alembic.sqlalchemy.org/en/latest/){:target="_blank"} |
| 128 | + instead. |
| 129 | + |
| 130 | +```python linenums="50" |
| 131 | +async def init_models() -> None: |
| 132 | + """Create tables if they don't already exist. |
| 133 | +
|
| 134 | + In a real-life example we would use Alembic to manage migrations. |
| 135 | + """ |
| 136 | + async with async_engine.begin() as conn: |
| 137 | + # await conn.run_sync(Base.metadata.drop_all) # noqa: ERA001 |
| 138 | + await conn.run_sync(Base.metadata.create_all) |
| 139 | +``` |
| 140 | + |
| 141 | +This function shows how to run a `syncronous` function in an `async` context |
| 142 | +using the `async_engine` object directly instead of the `async_session` object. |
| 143 | +On line **57** we use the `run_sync` method to run the `create_all` method of |
| 144 | +the `Base.metadata` object (a syncronous function). This will create all of the |
| 145 | +tables defined in the models. |
| 146 | + |
| 147 | +Next, we will look at the models themselves and the Schemas used to validate |
| 148 | +them within FastAPI. |
0 commit comments