src/db/postgres/models.py

"""SQLAlchemy database models."""

from uuid import uuid4
from sqlalchemy import Column, String, DateTime, Text, Integer, ForeignKey
from sqlalchemy.orm import relationship
from sqlalchemy.sql import func
from sqlalchemy.dialects.postgresql import JSONB
from src.db.postgres.connection import Base


class User(Base):
    """User model."""
    __tablename__ = "users"

    id = Column(String, primary_key=True, default=lambda: str(uuid4()))
    fullname = Column(String, nullable=False)
    email = Column(String, nullable=False, unique=True, index=True)
    password = Column(String, nullable=False)  # bcrypt-hashed
    company = Column(String)
    company_size = Column(String)
    function = Column(String)
    site = Column(String)
    role = Column(String)
    status = Column(String, nullable=False, default="active")  # active | inactive
    created_at = Column(DateTime(timezone=True), server_default=func.now())


class Document(Base):
    """Document model."""
    __tablename__ = "documents"

    id = Column(String, primary_key=True, default=lambda: str(uuid4()))
    user_id = Column(String, nullable=False, index=True)
    filename = Column(String, nullable=False)
    blob_name = Column(String, nullable=False, unique=True)
    file_size = Column(Integer)
    file_type = Column(String)  # pdf, docx, txt, etc.
    status = Column(String, default="uploaded")  # uploaded, processing, completed, failed
    processed_at = Column(DateTime(timezone=True))
    error_message = Column(Text)
    created_at = Column(DateTime(timezone=True), server_default=func.now())


class Room(Base):
    """Room model for chat sessions."""
    __tablename__ = "rooms"

    id = Column(String, primary_key=True, default=lambda: str(uuid4()))
    user_id = Column(String, nullable=False, index=True)
    title = Column(String, default="New Chat")
    created_at = Column(DateTime(timezone=True), server_default=func.now())
    updated_at = Column(DateTime(timezone=True), onupdate=func.now())

    status = Column(String, nullable=False, default="active")  # active | inactive

    messages = relationship("ChatMessage", back_populates="room", cascade="all, delete-orphan")


class ChatMessage(Base):
    """Chat message model."""
    __tablename__ = "chat_messages"

    id = Column(String, primary_key=True, default=lambda: str(uuid4()))
    room_id = Column(String, ForeignKey("rooms.id"), nullable=False, index=True)
    role = Column(String, nullable=False)  # user, assistant
    content = Column(Text, nullable=False)
    created_at = Column(DateTime(timezone=True), server_default=func.now())

    room = relationship("Room", back_populates="messages")
    sources = relationship("MessageSource", back_populates="message", cascade="all, delete-orphan")


class MessageSource(Base):
    """Sources (RAG references) attached to an assistant message."""
    __tablename__ = "message_sources"

    id = Column(String, primary_key=True, default=lambda: str(uuid4()))
    message_id = Column(String, ForeignKey("chat_messages.id", ondelete="CASCADE"), nullable=False, index=True)
    document_id = Column(String)
    filename = Column(Text)
    page_label = Column(Text)
    created_at = Column(DateTime(timezone=True), server_default=func.now())

    message = relationship("ChatMessage", back_populates="sources")


class DatabaseClient(Base):
    """User-registered external database connections."""
    __tablename__ = "databases"

    id = Column(String, primary_key=True, default=lambda: str(uuid4()))
    user_id = Column(String, nullable=False, index=True)
    name = Column(String, nullable=False)       # display name, e.g. "Prod DB"
    db_type = Column(String, nullable=False)    # postgres|mysql|sqlserver|supabase|bigquery|snowflake
    credentials = Column(JSONB, nullable=False) # per-type JSON; sensitive fields Fernet-encrypted
    status = Column(String, nullable=False, default="active")  # active | inactive
    created_at = Column(DateTime(timezone=True), server_default=func.now())
    updated_at = Column(DateTime(timezone=True), onupdate=func.now())


class Catalog(Base):
    """Per-user data catalog stored as a single jsonb row.

    `data` holds the full Pydantic Catalog (src/catalog/models.py:Catalog)
    serialized via `model_dump(mode="json")`. Read path uses
    `Catalog.model_validate(...)` to rehydrate.

    Dedicated table — kept separate from `langchain_pg_embedding` so unstructured
    embeddings and structured-catalog metadata never share storage.
    """
    __tablename__ = "data_catalog"

    user_id = Column(String, primary_key=True)
    data = Column(JSONB, nullable=False)
    schema_version = Column(String, nullable=False, default="1.0")
    generated_at = Column(DateTime(timezone=True), server_default=func.now())
    updated_at = Column(DateTime(timezone=True), onupdate=func.now())