"""SQLAlchemy models for the archive database.""" from datetime import datetime, timezone from sqlalchemy import Column, Integer, String, BigInteger, DateTime, create_engine from sqlalchemy.orm import DeclarativeBase, sessionmaker from pathlib import Path class Base(DeclarativeBase): pass class ArchiveVideo(Base): """SQLite model for archived downloads.""" __tablename__ = "archive_videos" row_id = Column(Integer, primary_key=True, autoincrement=True) video_id = Column(String(50), unique=True, nullable=False, index=True) title = Column(String(500), nullable=False) url = Column(String(1000), nullable=False) description = Column(String(5000), default="") thumbnail = Column(String(1000), default="") channel = Column(String(200), default="") views = Column(BigInteger, default=0) duration = Column(String(20), default="") category = Column(String(100), default="", index=True) download_path = Column(String(2000), default="") network_share_path = Column(String(2000)) file_size = Column(BigInteger, default=0) download_date = Column(DateTime, default=datetime.utcnow) item_type = Column(String(20), default="video") # video or playlist class ArchiveDB: """Database manager for the archive.""" def __init__(self, db_path: str = None): if db_path is None: db_path = str(Path.home() / ".config" / "youtube_cli" / "archive.db") self.db_path = db_path self.engine = create_engine(f"sqlite:///{self.db_path}") self.Session = sessionmaker(bind=self.engine) self._create_tables() def _create_tables(self): """Create all tables if they don't exist.""" Base.metadata.create_all(self.engine) def add_video(self, video: "ArchiveItem") -> "ArchiveVideo": """Add a video to the archive.""" session = self.Session() try: # Check if video already exists existing = session.query(ArchiveVideo).filter_by(video_id=video.video_id).first() if existing: # Update existing record existing.title = video.title existing.url = video.url existing.description = video.description existing.thumbnail = video.thumbnail existing.channel = video.channel existing.views = video.views existing.duration = video.duration existing.category = video.category existing.download_path = video.download_path existing.network_share_path = video.network_share_path existing.file_size = video.file_size or 0 existing.item_type = video.item_type else: # Create new record archive_video = ArchiveVideo( video_id=video.video_id, title=video.title, url=video.url, description=video.description, thumbnail=video.thumbnail, channel=video.channel, views=video.views, duration=video.duration, category=video.category, download_path=video.download_path, network_share_path=video.network_share_path, file_size=video.file_size or 0, download_date=datetime.fromisoformat(video.download_date) if video.download_date else datetime.now(timezone.utc), item_type=video.item_type, ) session.add(archive_video) session.commit() return existing or archive_video except Exception as e: session.rollback() raise e finally: session.close() def get_videos(self, page: int = 1, limit: int = 24, search: str = None, category: str = None, start_date: str = None, end_date: str = None): """Get archived videos with pagination and filtering.""" session = self.Session() try: query = session.query(ArchiveVideo) if search: search_pattern = f"%{search}%" query = query.filter( (ArchiveVideo.title.like(search_pattern)) | (ArchiveVideo.channel.like(search_pattern)) ) if category: query = query.filter(ArchiveVideo.category == category) if start_date: query = query.filter(ArchiveVideo.download_date >= datetime.fromisoformat(start_date)) if end_date: query = query.filter(ArchiveVideo.download_date <= datetime.fromisoformat(end_date)) total = query.count() offset = (page - 1) * limit videos = query.order_by(ArchiveVideo.download_date.desc()).offset(offset).limit(limit).all() return videos, total finally: session.close() def get_video(self, video_id: str): """Get a single video by ID.""" session = self.Session() try: return session.query(ArchiveVideo).filter_by(video_id=video_id).first() finally: session.close() def delete_video(self, video_id: str) -> bool: """Delete a video from the archive.""" session = self.Session() try: video = session.query(ArchiveVideo).filter_by(video_id=video_id).first() if video: session.delete(video) session.commit() return True return False except Exception as e: session.rollback() raise e finally: session.close() def clear_archive(self) -> int: """Clear all videos from the archive. Returns count of deleted videos.""" session = self.Session() try: count = session.query(ArchiveVideo).count() session.query(ArchiveVideo).delete() session.commit() return count except Exception as e: session.rollback() raise e finally: session.close() def get_stats(self): """Get archive statistics.""" from sqlalchemy import func session = self.Session() try: total = session.query(ArchiveVideo).count() total_size = session.query(func.coalesce(func.sum(ArchiveVideo.file_size), 0)).scalar() categories = session.query(ArchiveVideo.category, func.count(ArchiveVideo.row_id)) \ .group_by(ArchiveVideo.category).all() return { "total": total, "totalSize": f"{total_size / (1024 * 1024):.1f} MB" if total_size else "0 MB", "categories": {cat: count for cat, count in categories if cat} } finally: session.close() def get_categories(self) -> list: """Get unique categories from the archive.""" from sqlalchemy import func session = self.Session() try: categories = session.query(ArchiveVideo.category).filter( ArchiveVideo.category != "" ).distinct().all() return [cat[0] for cat in categories] finally: session.close() def export_archive(self, fmt: str = "json"): """Export archive data as JSON or CSV.""" session = self.Session() try: videos = session.query(ArchiveVideo).order_by(ArchiveVideo.download_date.desc()).all() if fmt == "csv": return self._to_csv(videos) return self._to_json(videos) finally: session.close() def _to_json(self, videos): """Convert archive videos to JSON.""" import json items = [] for v in videos: items.append({ "videoId": v.video_id, "title": v.title, "url": v.url, "description": v.description, "thumbnail": v.thumbnail, "channel": v.channel, "views": v.views, "duration": v.duration, "category": v.category, "downloadPath": v.download_path, "networkSharePath": v.network_share_path, "fileSize": v.file_size, "downloadDate": v.download_date.isoformat() if v.download_date else "", "type": v.item_type, }) return json.dumps({"items": items, "total": len(items)}, indent=2) def _to_csv(self, videos): """Convert archive videos to CSV.""" import csv import io output = io.StringIO() writer = csv.writer(output) writer.writerow(["videoId", "title", "url", "channel", "category", "downloadDate", "fileSize"]) for v in videos: writer.writerow([v.video_id, v.title, v.url, v.channel, v.category, v.download_date.isoformat() if v.download_date else "", v.file_size or 0]) return output.getvalue()