246 lines
9.0 KiB
Python
246 lines
9.0 KiB
Python
"""SQLAlchemy models for the archive database."""
|
|
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
from sqlalchemy import (
|
|
BigInteger,
|
|
Column,
|
|
DateTime,
|
|
Integer,
|
|
String,
|
|
create_engine,
|
|
)
|
|
from sqlalchemy.orm import DeclarativeBase, sessionmaker
|
|
|
|
|
|
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."""
|
|
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()
|