yottob/models.py

"""Database models for YouTube feed storage."""

from datetime import datetime
from typing import List, Optional
from enum import Enum as PyEnum

import bcrypt
from flask_login import UserMixin
from sqlalchemy import String, DateTime, ForeignKey, Index, Enum, BigInteger, Text
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship


class DownloadStatus(PyEnum):
    """Download status enumeration."""
    PENDING = "pending"
    DOWNLOADING = "downloading"
    COMPLETED = "completed"
    FAILED = "failed"


class Base(DeclarativeBase):
    """Base class for all database models."""
    pass


class User(Base, UserMixin):
    """User model for authentication."""

    __tablename__ = "users"

    id: Mapped[int] = mapped_column(primary_key=True)
    username: Mapped[str] = mapped_column(String(80), unique=True, nullable=False, index=True)
    email: Mapped[str] = mapped_column(String(120), unique=True, nullable=False, index=True)
    password_hash: Mapped[str] = mapped_column(String(255), nullable=False)
    created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, default=datetime.utcnow)

    # Relationships
    channels: Mapped[List["Channel"]] = relationship("Channel", back_populates="user", cascade="all, delete-orphan")

    def set_password(self, password: str) -> None:
        """Hash and set the password."""
        self.password_hash = bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')

    def check_password(self, password: str) -> bool:
        """Check if the provided password matches the hash."""
        return bcrypt.checkpw(password.encode('utf-8'), self.password_hash.encode('utf-8'))

    def __repr__(self) -> str:
        return f"<User(id={self.id}, username='{self.username}', email='{self.email}')>"


class Channel(Base):
    """YouTube channel model."""

    __tablename__ = "channels"

    id: Mapped[int] = mapped_column(primary_key=True)
    user_id: Mapped[int] = mapped_column(ForeignKey("users.id"), nullable=False, index=True)
    channel_id: Mapped[str] = mapped_column(String(50), nullable=False, index=True)
    title: Mapped[str] = mapped_column(String(200), nullable=False)
    link: Mapped[str] = mapped_column(String(500), nullable=False)
    rss_url: Mapped[str] = mapped_column(String(500), nullable=False)
    last_fetched_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)

    # Relationships
    user: Mapped["User"] = relationship("User", back_populates="channels")
    video_entries: Mapped[List["VideoEntry"]] = relationship("VideoEntry", back_populates="channel", cascade="all, delete-orphan")

    # Composite unique constraint on user_id and channel_id (one user can't subscribe to same channel twice)
    __table_args__ = (
        Index('idx_user_channel', 'user_id', 'channel_id', unique=True),
    )

    def __repr__(self) -> str:
        return f"<Channel(id={self.id}, channel_id='{self.channel_id}', title='{self.title}')>"


class VideoEntry(Base):
    """YouTube video entry model."""

    __tablename__ = "video_entries"

    id: Mapped[int] = mapped_column(primary_key=True)
    channel_id: Mapped[int] = mapped_column(ForeignKey("channels.id"), nullable=False)
    video_id: Mapped[str] = mapped_column(String(50), nullable=False, index=True)
    title: Mapped[str] = mapped_column(String(500), nullable=False)
    video_url: Mapped[str] = mapped_column(String(500), nullable=False, index=True)
    thumbnail_url: Mapped[Optional[str]] = mapped_column(String(500), nullable=True)
    description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    published_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, index=True)
    created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, default=datetime.utcnow)

    # Download tracking fields
    download_status: Mapped[DownloadStatus] = mapped_column(
        Enum(DownloadStatus),
        nullable=False,
        default=DownloadStatus.PENDING
    )
    download_path: Mapped[Optional[str]] = mapped_column(String(1000), nullable=True)
    download_started_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
    download_completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
    download_error: Mapped[Optional[str]] = mapped_column(String(2000), nullable=True)
    file_size: Mapped[Optional[int]] = mapped_column(BigInteger, nullable=True)

    # Relationship to channel
    channel: Mapped["Channel"] = relationship("Channel", back_populates="video_entries")

    # Indexes for faster queries
    __table_args__ = (
        Index('idx_channel_created', 'channel_id', 'created_at'),
        Index('idx_download_status', 'download_status'),
        Index('idx_published_at', 'published_at'),
        Index('idx_video_id_channel', 'video_id', 'channel_id', unique=True),
    )

    def __repr__(self) -> str:
        return f"<VideoEntry(id={self.id}, title='{self.title}', video_url='{self.video_url}', status='{self.download_status.value}')>"

    def to_dict(self) -> dict:
        """Convert to dictionary for API responses."""
        return {
            "id": self.id,
            "video_id": self.video_id,
            "title": self.title,
            "video_url": self.video_url,
            "thumbnail_url": self.thumbnail_url,
            "description": self.description,
            "published_at": self.published_at.isoformat(),
            "created_at": self.created_at.isoformat(),
            "download_status": self.download_status.value,
            "download_path": self.download_path,
            "download_started_at": self.download_started_at.isoformat() if self.download_started_at else None,
            "download_completed_at": self.download_completed_at.isoformat() if self.download_completed_at else None,
            "download_error": self.download_error,
            "file_size": self.file_size
        }