Spaces:
Sleeping
Sleeping
| """ | |
| S.T.E.W Database Models — SQLAlchemy ORM. | |
| Uses String for plan field to support both PostgreSQL and SQLite. | |
| """ | |
| import uuid | |
| from datetime import datetime | |
| from typing import Optional | |
| from sqlalchemy import ( | |
| String, Text, Integer, DateTime, Boolean, JSON, | |
| ForeignKey, func | |
| ) | |
| from sqlalchemy.orm import Mapped, mapped_column, relationship | |
| from server.database import Base | |
| VALID_PLANS = ("free", "pro", "business", "enterprise") | |
| def _uuid() -> str: | |
| return str(uuid.uuid4()) | |
| class User(Base): | |
| __tablename__ = "users" | |
| id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid) | |
| name: Mapped[str] = mapped_column(String(255), nullable=False) | |
| email: Mapped[str] = mapped_column(String(255), unique=True, nullable=False, index=True) | |
| password_hash: Mapped[Optional[str]] = mapped_column(String(255), nullable=True) | |
| plan: Mapped[str] = mapped_column(String(20), default="free", nullable=False) | |
| api_key: Mapped[str] = mapped_column(String(80), unique=True, nullable=False, index=True) | |
| is_active: Mapped[bool] = mapped_column(Boolean, default=True) | |
| created_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now()) | |
| updated_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now(), onupdate=func.now()) | |
| conversations: Mapped[list["Conversation"]] = relationship(back_populates="user", cascade="all, delete-orphan") | |
| api_calls: Mapped[list["APICall"]] = relationship(back_populates="user", cascade="all, delete-orphan") | |
| documents: Mapped[list["Document"]] = relationship(back_populates="user", cascade="all, delete-orphan") | |
| class Conversation(Base): | |
| __tablename__ = "conversations" | |
| id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid) | |
| user_id: Mapped[str] = mapped_column(String(36), ForeignKey("users.id", ondelete="CASCADE"), nullable=False) | |
| title: Mapped[Optional[str]] = mapped_column(String(255), nullable=True) | |
| messages: Mapped[list] = mapped_column(JSON, default=list) | |
| created_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now()) | |
| updated_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now(), onupdate=func.now()) | |
| user: Mapped["User"] = relationship(back_populates="conversations") | |
| class APICall(Base): | |
| __tablename__ = "api_calls" | |
| id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid) | |
| user_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("users.id", ondelete="SET NULL"), nullable=True) | |
| endpoint: Mapped[str] = mapped_column(String(255), nullable=False) | |
| method: Mapped[str] = mapped_column(String(10), default="POST") | |
| tokens_used: Mapped[int] = mapped_column(Integer, default=0) | |
| status_code: Mapped[int] = mapped_column(Integer, default=200) | |
| timestamp: Mapped[datetime] = mapped_column(DateTime, server_default=func.now()) | |
| user: Mapped[Optional["User"]] = relationship(back_populates="api_calls") | |
| class Document(Base): | |
| __tablename__ = "documents" | |
| id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid) | |
| user_id: Mapped[str] = mapped_column(String(36), ForeignKey("users.id", ondelete="CASCADE"), nullable=False) | |
| filename: Mapped[str] = mapped_column(String(255), nullable=False) | |
| file_type: Mapped[str] = mapped_column(String(50), nullable=False) | |
| content: Mapped[Optional[str]] = mapped_column(Text, nullable=True) | |
| file_size: Mapped[int] = mapped_column(Integer, default=0) | |
| created_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now()) | |
| user: Mapped["User"] = relationship(back_populates="documents") | |
| class PaymentTransaction(Base): | |
| __tablename__ = "payment_transactions" | |
| id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid) | |
| user_id: Mapped[str] = mapped_column(String(36), ForeignKey("users.id", ondelete="CASCADE"), nullable=False) | |
| reference: Mapped[str] = mapped_column(String(255), unique=True, nullable=False) | |
| plan: Mapped[str] = mapped_column(String(50), nullable=False) | |
| amount: Mapped[int] = mapped_column(Integer, nullable=False) | |
| status: Mapped[str] = mapped_column(String(50), default="pending") | |
| created_at: Mapped[datetime] = mapped_column(DateTime, server_default=func.now()) | |