gongzhiyongandClaude Sonnet 4.6
48d7dd7993
- Add return_dto=None to delete_attachment decorator to fix Litestar ImproperlyConfiguredException on 204 status code with None return type - Catch BaseException instead of Exception in create_tables() to handle anyio ExceptionGroup from concurrent gunicorn workers - Enable debug=True temporarily to capture detailed error traces Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
138 lines
4.6 KiB
Python
138 lines
4.6 KiB
Python
"""PostgreSQL models, engine, and session management."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
|
|
from sqlalchemy import BigInteger, DateTime, ForeignKey, String, Text, func
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
|
|
|
|
from app.config import settings
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Engine & session factory
|
|
# ---------------------------------------------------------------------------
|
|
|
|
engine = create_async_engine(
|
|
settings.database_url,
|
|
echo=False,
|
|
pool_size=5,
|
|
max_overflow=10,
|
|
pool_pre_ping=True,
|
|
)
|
|
|
|
async_session_factory = async_sessionmaker(engine, expire_on_commit=False)
|
|
|
|
|
|
async def get_session() -> AsyncSession:
|
|
"""Yield a new async session (for use in route handlers)."""
|
|
async with async_session_factory() as session:
|
|
yield session
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ORM base and models
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
def _utcnow() -> datetime:
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
class Conversation(Base):
|
|
__tablename__ = "conversations"
|
|
|
|
id: Mapped[str] = mapped_column(
|
|
String(64), primary_key=True, default=lambda: str(uuid.uuid4())
|
|
)
|
|
title: Mapped[str] = mapped_column(String(512), default="New conversation")
|
|
created_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True), default=_utcnow, server_default=func.now()
|
|
)
|
|
updated_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True), default=_utcnow, onupdate=_utcnow, server_default=func.now()
|
|
)
|
|
|
|
messages: Mapped[list[Message]] = relationship(
|
|
back_populates="conversation",
|
|
cascade="all, delete-orphan",
|
|
order_by="Message.created_at",
|
|
)
|
|
|
|
|
|
class Message(Base):
|
|
__tablename__ = "messages"
|
|
|
|
id: Mapped[str] = mapped_column(
|
|
String(64), primary_key=True, default=lambda: str(uuid.uuid4())
|
|
)
|
|
conversation_id: Mapped[str] = mapped_column(
|
|
String(64), ForeignKey("conversations.id", ondelete="CASCADE"), index=True
|
|
)
|
|
role: Mapped[str] = mapped_column(String(32)) # "human", "ai", "system"
|
|
content: Mapped[str] = mapped_column(Text, default="")
|
|
created_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True), default=_utcnow, server_default=func.now()
|
|
)
|
|
|
|
conversation: Mapped[Conversation] = relationship(back_populates="messages")
|
|
|
|
|
|
class Attachment(Base):
|
|
__tablename__ = "attachments"
|
|
|
|
id: Mapped[str] = mapped_column(
|
|
String(64), primary_key=True, default=lambda: str(uuid.uuid4())
|
|
)
|
|
conversation_id: Mapped[str | None] = mapped_column(
|
|
String(64),
|
|
ForeignKey("conversations.id", ondelete="SET NULL"),
|
|
nullable=True,
|
|
index=True,
|
|
)
|
|
message_id: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
|
filename: Mapped[str] = mapped_column(String(512))
|
|
content_type: Mapped[str] = mapped_column(String(128), default="application/octet-stream")
|
|
blob_url: Mapped[str] = mapped_column(Text)
|
|
size_bytes: Mapped[int] = mapped_column(BigInteger, default=0)
|
|
created_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True), default=_utcnow, server_default=func.now()
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Table creation helper
|
|
# ---------------------------------------------------------------------------
|
|
|
|
async def create_tables() -> None:
|
|
"""Create all tables if they don't exist.
|
|
|
|
Wrapped in try/except to handle the race condition when multiple
|
|
gunicorn workers call create_all simultaneously and PostgreSQL
|
|
raises a UniqueViolation on the pg_type_typname_nsp_index.
|
|
|
|
We catch BaseException (not just Exception) because anyio wraps
|
|
concurrent errors in an ExceptionGroup which is a BaseException
|
|
subclass and would otherwise crash the Litestar lifespan.
|
|
"""
|
|
try:
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
except BaseException:
|
|
# Table likely already created by another worker -- safe to ignore
|
|
import logging
|
|
logging.getLogger(__name__).debug(
|
|
"create_tables race condition (another worker likely created tables)",
|
|
exc_info=True,
|
|
)
|
|
|
|
|
|
async def dispose_engine() -> None:
|
|
"""Dispose the engine (for clean shutdown)."""
|
|
await engine.dispose()
|