47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
from typing import AsyncIterator
|
|
import uuid
|
|
|
|
from fastapi import Request
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
|
from sqlalchemy import String, Uuid
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
# User Thread Mapping Table
|
|
class UserThreadMapping(Base):
|
|
"""Mapping of users to their chat threads.
|
|
|
|
Used to keep track of which user owns which chat thread.
|
|
Also stores meta data like thread name.
|
|
|
|
1:N relationship between user and thread. A thread belongs to exactly one user.
|
|
A user can have multiple threads.
|
|
Thread_id is unique in the table.
|
|
"""
|
|
|
|
__tablename__ = "userThreads"
|
|
id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4)
|
|
userId: Mapped[str] = mapped_column(String(255), nullable=False)
|
|
threadId: Mapped[str] = mapped_column(String(255), unique=True, nullable=False)
|
|
threadName: Mapped[str] = mapped_column(String(255), nullable=False)
|
|
|
|
|
|
# Dependency that pulls the sessionmaker off app.state
|
|
# This is set in app.py on startup in @asynccontextmanager
|
|
# TODO: If we use SQLAlchemy in other places, we can move this to a shared module
|
|
async def get_async_db_session(request: Request) -> AsyncIterator[AsyncSession]:
|
|
SessionLocal: async_sessionmaker[AsyncSession] = (
|
|
request.app.state.checkpoint_sessionmaker
|
|
)
|
|
async with SessionLocal() as session:
|
|
yield session
|
|
|
|
|
|
# Optional helper to init tables at startup (demo only)
|
|
async def init_models(engine) -> None:
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|