117 lines
3.8 KiB
Python
117 lines
3.8 KiB
Python
"""Database models and async engine (pgvector / pg16).
|
|
|
|
SQLAlchemy 2.0 async with asyncpg. The users table stores profile + tool
|
|
selection; the preferences table leaves room for per-user preferences and
|
|
future semantic-memory vectors (pgvector column commented out deliberately so
|
|
nothing depends on the extension at first launch).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
|
|
from sqlalchemy import (
|
|
Boolean,
|
|
DateTime,
|
|
ForeignKey,
|
|
MetaData,
|
|
String,
|
|
func,
|
|
select,
|
|
)
|
|
from sqlalchemy.ext.asyncio import (
|
|
AsyncSession,
|
|
async_sessionmaker,
|
|
create_async_engine,
|
|
)
|
|
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
|
|
|
|
from .config import settings
|
|
|
|
# Recommend enabling pgvector: CREATE EXTENSION IF NOT EXISTS vector;
|
|
# then uncomment the `embedding` column on ToolPreference.
|
|
naming_convention = {
|
|
"ix": "ix_%(column_0_label)s",
|
|
"uq": "uq_%(table_name)s_%(column_0_name)s",
|
|
"ck": "ck_%(table_name)s_%(constraint_name)s",
|
|
"fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
|
|
"pk": "pk_%(table_name)s",
|
|
}
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
metadata = MetaData(naming_convention=naming_convention)
|
|
|
|
|
|
class User(Base):
|
|
__tablename__ = "users"
|
|
|
|
id: Mapped[int] = mapped_column(primary_key=True)
|
|
username: Mapped[str] = mapped_column(String(64), unique=True, index=True)
|
|
full_name: Mapped[str] = mapped_column(String(128))
|
|
password_hash: Mapped[str] = mapped_column(String(256))
|
|
is_admin: Mapped[bool] = mapped_column(Boolean, default=False)
|
|
can_chat: Mapped[bool] = mapped_column(Boolean, default=True)
|
|
can_image: Mapped[bool] = mapped_column(Boolean, default=True)
|
|
can_video: Mapped[bool] = mapped_column(Boolean, default=True)
|
|
can_audio: Mapped[bool] = mapped_column(Boolean, default=True)
|
|
can_docs: Mapped[bool] = mapped_column(Boolean, default=False)
|
|
is_active: Mapped[bool] = mapped_column(Boolean, default=True)
|
|
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
|
|
|
|
preferences: Mapped[list["ToolPreference"]] = relationship(
|
|
back_populates="user", cascade="all, delete-orphan"
|
|
)
|
|
|
|
def __repr__(self) -> str:
|
|
return f"<User {self.username!r} admin={self.is_admin}>"
|
|
|
|
|
|
class ToolPreference(Base):
|
|
"""Per-user saved preference for a tool slot (e.g. favourite editor)."""
|
|
|
|
__tablename__ = "tool_preferences"
|
|
|
|
id: Mapped[int] = mapped_column(primary_key=True)
|
|
user_id: Mapped[int] = mapped_column(
|
|
ForeignKey("users.id", ondelete="CASCADE"), index=True
|
|
)
|
|
tool_id: Mapped[str] = mapped_column(String(64), index=True)
|
|
value: Mapped[str] = mapped_column(String(512), default="")
|
|
|
|
# Semantic memory vector (optional, pgvector).
|
|
# embedding: Mapped[Vector | None] = mapped_column(Vector(1536), nullable=True)
|
|
|
|
user: Mapped[User] = relationship(back_populates="preferences")
|
|
|
|
|
|
engine = create_async_engine(
|
|
settings.DATABASE_URL,
|
|
echo=settings.DEBUG,
|
|
pool_pre_ping=True,
|
|
)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
|
|
async def init_db() -> None:
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
|
|
async def get_session() -> AsyncSession: # FastAPI dependency
|
|
async with SessionLocal() as session:
|
|
yield session
|
|
|
|
|
|
from contextlib import asynccontextmanager
|
|
|
|
|
|
@asynccontextmanager
|
|
async def session_scope():
|
|
"""Explicit async context manager for service/route use (not a FastAPI dep)."""
|
|
async with SessionLocal() as session:
|
|
yield session
|
|
|
|
|
|
async def get_user_by_username(session: AsyncSession, username: str) -> User | None:
|
|
result = await session.execute(select(User).where(User.username == username))
|
|
return result.scalar_one_or_none() |