Chapter 3: SQLModel

In models/pokemon.py, we’ll define our SQLModel. By setting the ID as an optional primary key, we let PostgreSQL generate a new ID when creating an entity. Using default_factory, the model automatically inserts the creation timestamp along with three required fields: name, number, and an enum region.

from __future__ import annotations

from datetime import datetime, timezone

import enum
from sqlmodel import Field, SQLModel

class Region(enum.Enum):
    kanto = "Kanto"
    johto = "Johto"


class Pokemon(SQLModel, table=True):
    id: int | None = Field(default=None, primary_key=True)
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    name: str
    number: int
    region: Region

We’ll create a new db.py module to handle database connectivity. Here, we’ll read PostgreSQL credentials from environment variables, initialise the engine, and create a function that yields a Session object.

from os import environ
from sqlmodel import SQLModel, create_engine, Session

username = environ.get("POSTGRES_USER")
password = environ.get("POSTGRES_PASSWORD")
host = environ.get("POSTGRES_HOST")
db = environ.get("POSTGRES_DB")
postgres_url = f"postgresql://{username}:{password}@{host}:5432/{db}"

engine = create_engine(postgres_url)

get_session():
    with Session(engine) as session:
        yield session