beavergram/backend/src/utils/db/models.py

from datetime import datetime
from enum import StrEnum
from typing import Any

from sqlalchemy import BigInteger, Column, DateTime, LargeBinary, Text, func
from sqlalchemy.dialects.postgresql import JSONB
from sqlmodel import Field, SQLModel


class JobStatus(StrEnum):
    PENDING = "pending"
    RUNNING = "running"
    DONE = "done"
    FAILED = "failed"


class Account(SQLModel, table=True):
    __tablename__ = "accounts"

    account_id: int | None = Field(default=None, primary_key=True)
    tg_user_id: int | None = Field(
        default=None, sa_column=Column(BigInteger, unique=True)
    )
    label: str | None = None
    phone: str | None = None
    session_name: str
    is_active: bool = True
    raw: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )
    created_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), nullable=False, server_default=func.now()
        )
    )
    updated_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True),
            nullable=False,
            server_default=func.now(),
            onupdate=func.now(),
        )
    )


class Message(SQLModel, table=True):
    __tablename__ = "messages"

    account_id: int = Field(primary_key=True)
    chat_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    message_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    date: datetime = Field(sa_column=Column(DateTime(timezone=True), primary_key=True))
    sender_id: int | None = Field(default=None, sa_column=Column(BigInteger))
    text: str | None = None
    raw: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )
    has_media: bool = False
    is_self_destruct: bool = False
    edited_at: datetime | None = Field(
        default=None, sa_column=Column(DateTime(timezone=True))
    )
    deleted_at: datetime | None = Field(
        default=None, sa_column=Column(DateTime(timezone=True))
    )


class Folder(SQLModel, table=True):
    __tablename__ = "folders"

    account_id: int = Field(primary_key=True)
    folder_id: int = Field(primary_key=True)
    title: str
    order_index: int
    is_chatlist: bool = False
    raw: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )
    updated_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True),
            nullable=False,
            server_default=func.now(),
            onupdate=func.now(),
        )
    )


class MessageVersion(SQLModel, table=True):
    __tablename__ = "message_versions"

    account_id: int = Field(primary_key=True)
    chat_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    message_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    observed_at: datetime = Field(
        sa_column=Column(DateTime(timezone=True), primary_key=True)
    )
    edit_date: datetime | None = Field(
        default=None, sa_column=Column(DateTime(timezone=True))
    )
    text: str | None = None
    raw: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )


class Media(SQLModel, table=True):
    __tablename__ = "media"

    id: int | None = Field(default=None, primary_key=True)
    account_id: int
    chat_id: int = Field(sa_column=Column(BigInteger, nullable=False))
    message_id: int = Field(sa_column=Column(BigInteger, nullable=False))
    kind: str
    storage_key: str | None = None
    file_size: int | None = Field(default=None, sa_column=Column(BigInteger))
    mime: str | None = None
    ttl_seconds: int | None = None
    downloaded: bool = False
    extracted_text: str | None = None
    created_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), nullable=False, server_default=func.now()
        )
    )


class Callback(SQLModel, table=True):
    __tablename__ = "callbacks"

    id: int | None = Field(default=None, primary_key=True)
    account_id: int
    chat_id: int = Field(sa_column=Column(BigInteger, nullable=False))
    message_id: int = Field(sa_column=Column(BigInteger, nullable=False))
    position: int
    label: str | None = None
    data: bytes | None = Field(default=None, sa_column=Column(LargeBinary))
    created_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), nullable=False, server_default=func.now()
        )
    )


class Reaction(SQLModel, table=True):
    __tablename__ = "reactions"

    id: int | None = Field(default=None, primary_key=True)
    account_id: int
    chat_id: int = Field(sa_column=Column(BigInteger, nullable=False))
    message_id: int = Field(sa_column=Column(BigInteger, nullable=False))
    peer_id: int = Field(sa_column=Column(BigInteger, nullable=False))
    reaction: str
    added_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), nullable=False, server_default=func.now()
        )
    )
    removed_at: datetime | None = Field(
        default=None, sa_column=Column(DateTime(timezone=True))
    )
    created_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), nullable=False, server_default=func.now()
        )
    )


class Job(SQLModel, table=True):
    __tablename__ = "jobs"

    id: int | None = Field(default=None, sa_column=Column(BigInteger, primary_key=True))
    account_id: int
    kind: str
    status: str = JobStatus.PENDING
    params: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )
    cursor: dict[str, Any] | None = Field(default=None, sa_column=Column(JSONB))
    progress: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )
    flood_waits: int = 0
    attempts: int = 0
    error: str | None = None
    created_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), nullable=False, server_default=func.now()
        )
    )
    updated_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True),
            nullable=False,
            server_default=func.now(),
            onupdate=func.now(),
        )
    )
    started_at: datetime | None = Field(
        default=None, sa_column=Column(DateTime(timezone=True))
    )
    finished_at: datetime | None = Field(
        default=None, sa_column=Column(DateTime(timezone=True))
    )


class CapturePolicy(SQLModel, table=True):
    __tablename__ = "capture_policy"

    id: int | None = Field(default=None, primary_key=True)
    account_id: int | None = None
    scope_type: str
    scope_id: int | None = Field(default=None, sa_column=Column(BigInteger))
    messages: bool = False
    media: bool = False
    self_destruct_media: bool = False
    stt: bool = False
    reactions: bool = False
    track_edits_deletes: bool = False
    profile_history: bool = False
    stories: bool = False
    presence: bool = False
    backfill: bool = False
    created_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), nullable=False, server_default=func.now()
        )
    )
    updated_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True),
            nullable=False,
            server_default=func.now(),
            onupdate=func.now(),
        )
    )


class Presence(SQLModel, table=True):
    __tablename__ = "presence"

    account_id: int = Field(primary_key=True)
    peer_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    ts: datetime = Field(sa_column=Column(DateTime(timezone=True), primary_key=True))
    status: str
    last_online_date: datetime | None = Field(
        default=None, sa_column=Column(DateTime(timezone=True))
    )
    next_offline_date: datetime | None = Field(
        default=None, sa_column=Column(DateTime(timezone=True))
    )
    raw: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )


class ReadReceipt(SQLModel, table=True):
    __tablename__ = "read_receipts"

    account_id: int = Field(primary_key=True)
    chat_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    message_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    ts: datetime = Field(sa_column=Column(DateTime(timezone=True), primary_key=True))
    reader_id: int = Field(sa_column=Column(BigInteger, nullable=False))
    kind: str
    raw: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )


class Peer(SQLModel, table=True):
    __tablename__ = "peers"

    account_id: int = Field(primary_key=True)
    peer_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    first_name: str | None = None
    last_name: str | None = None
    username: str | None = None
    phone: str | None = None
    photo_unique_id: str | None = None
    is_deleted_account: bool = False
    raw: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )
    updated_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True),
            nullable=False,
            server_default=func.now(),
            onupdate=func.now(),
        )
    )


class PeerHistory(SQLModel, table=True):
    __tablename__ = "peer_history"

    account_id: int = Field(primary_key=True)
    peer_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    observed_at: datetime = Field(
        sa_column=Column(DateTime(timezone=True), primary_key=True)
    )
    first_name: str | None = None
    last_name: str | None = None
    username: str | None = None
    phone: str | None = None
    photo_unique_id: str | None = None
    is_deleted_account: bool = False
    raw: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )


class Avatar(SQLModel, table=True):
    __tablename__ = "avatars"

    id: int | None = Field(default=None, primary_key=True)
    account_id: int
    owner_id: int = Field(sa_column=Column(BigInteger, nullable=False))
    owner_kind: str
    unique_id: str
    storage_key: str | None = None
    file_size: int | None = Field(default=None, sa_column=Column(BigInteger))
    mime: str | None = None
    downloaded: bool = False
    raw: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )
    first_seen_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), nullable=False, server_default=func.now()
        )
    )


class ChatHistory(SQLModel, table=True):
    __tablename__ = "chat_history"

    account_id: int = Field(primary_key=True)
    chat_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    message_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    event: str
    title: str | None = None
    photo_unique_id: str | None = None
    actor_id: int | None = Field(default=None, sa_column=Column(BigInteger))
    ts: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=False))
    raw: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )


class Membership(SQLModel, table=True):
    __tablename__ = "memberships"

    account_id: int = Field(primary_key=True)
    chat_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    message_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    user_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    event: str
    actor_id: int | None = Field(default=None, sa_column=Column(BigInteger))
    ts: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=False))
    raw: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )


class Story(SQLModel, table=True):
    __tablename__ = "stories"

    account_id: int = Field(primary_key=True)
    peer_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    story_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    date: datetime | None = Field(
        default=None, sa_column=Column(DateTime(timezone=True))
    )
    expire_date: datetime | None = Field(
        default=None, sa_column=Column(DateTime(timezone=True))
    )
    caption: str | None = None
    media_kind: str | None = None
    storage_key: str | None = None
    file_size: int | None = Field(default=None, sa_column=Column(BigInteger))
    downloaded: bool = False
    views: int | None = None
    pinned: bool = False
    deleted: bool = False
    raw: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )
    observed_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), nullable=False, server_default=func.now()
        )
    )


class Link(SQLModel, table=True):
    __tablename__ = "links"

    account_id: int = Field(primary_key=True)
    chat_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    message_id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    position: int = Field(primary_key=True)
    url: str
    kind: str
    web_url: str | None = None
    web_title: str | None = None
    web_description: str | None = None
    web_site_name: str | None = None
    raw: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )


class Annotation(SQLModel, table=True):
    __tablename__ = "annotations"

    id: int | None = Field(default=None, primary_key=True)
    account_id: int
    chat_id: int = Field(sa_column=Column(BigInteger, nullable=False, index=True))
    message_id: int = Field(sa_column=Column(BigInteger, nullable=False))
    text: str
    created_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), nullable=False, server_default=func.now()
        )
    )
    updated_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True),
            nullable=False,
            server_default=func.now(),
            onupdate=func.now(),
        )
    )


class Watch(SQLModel, table=True):
    __tablename__ = "watches"

    id: int | None = Field(default=None, primary_key=True)
    account_id: int
    kind: str
    params: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )
    enabled: bool = True
    created_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), nullable=False, server_default=func.now()
        )
    )
    updated_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True),
            nullable=False,
            server_default=func.now(),
            onupdate=func.now(),
        )
    )


class Alert(SQLModel, table=True):
    __tablename__ = "alerts"

    id: int | None = Field(default=None, primary_key=True)
    account_id: int
    watch_id: int = Field(foreign_key="watches.id")
    ts: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=False))
    payload: dict[str, Any] = Field(
        default_factory=dict, sa_column=Column(JSONB, nullable=False)
    )
    seen: bool = False
    dedup_key: str | None = Field(default=None, sa_column=Column(Text, nullable=True))
    created_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True), nullable=False, server_default=func.now()
        )
    )


class Dialog(SQLModel, table=True):
    __tablename__ = "dialogs"

    account_id: int = Field(primary_key=True)
    chat_id: int = Field(primary_key=True)
    updated_at: datetime = Field(
        sa_column=Column(
            DateTime(timezone=True),
            nullable=False,
            server_default=func.now(),
            onupdate=func.now(),
        )
    )