#!/usr/bin/env python3 """ Database operations module for NFOGuard """ import sqlite3 from pathlib import Path from datetime import datetime, timezone from typing import Dict, List, Optional, Tuple, Any from core.logging import _log class NFOGuardDatabase: """Database operations for TV series, episodes, movies, and their metadata""" # Database schemas TV_SCHEMA = """ PRAGMA journal_mode=WAL; CREATE TABLE IF NOT EXISTS meta ( key TEXT PRIMARY KEY, value TEXT ); CREATE TABLE IF NOT EXISTS series ( imdb_id TEXT PRIMARY KEY, series_path TEXT NOT NULL, last_applied TEXT ); CREATE TABLE IF NOT EXISTS episode_dates ( imdb_id TEXT NOT NULL, season INTEGER NOT NULL, episode INTEGER NOT NULL, aired TEXT, dateadded TEXT, source TEXT, has_video_file INTEGER DEFAULT 0, PRIMARY KEY (imdb_id, season, episode) ); """ MOVIE_SCHEMA = """ CREATE TABLE IF NOT EXISTS movies ( imdb_id TEXT PRIMARY KEY, movie_path TEXT NOT NULL, last_applied TEXT ); CREATE TABLE IF NOT EXISTS movie_dates ( imdb_id TEXT PRIMARY KEY, released TEXT, dateadded TEXT, source TEXT, has_video_file INTEGER DEFAULT 0 ); """ def __init__(self, db_path: Path): self.db_path = Path(db_path) self._ensure_database() def _ensure_database(self): """Create database and tables if they don't exist""" self.db_path.parent.mkdir(parents=True, exist_ok=True) with self.get_connection() as conn: # Check and create TV schema try: result = conn.execute("PRAGMA table_info(episode_dates);").fetchall() columns = [row[1] for row in result] required_columns = ['imdb_id', 'season', 'episode', 'aired', 'dateadded', 'source'] if not all(col in columns for col in required_columns): _log("WARNING", "TV database schema outdated, recreating tables...") conn.execute("DROP TABLE IF EXISTS episode_dates;") conn.execute("DROP TABLE IF EXISTS series;") conn.execute("DROP TABLE IF EXISTS meta;") elif 'has_video_file' not in columns: _log("INFO", "Adding has_video_file column to episode_dates table") conn.execute("ALTER TABLE episode_dates ADD COLUMN has_video_file INTEGER DEFAULT 0;") except sqlite3.OperationalError: # Tables don't exist, which is fine pass # Create TV schema conn.executescript(self.TV_SCHEMA) # Check and create movie schema try: result = conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='movies'").fetchone() if not result: _log("INFO", "Creating movie tables in database") conn.executescript(self.MOVIE_SCHEMA) except sqlite3.Error as e: _log("ERROR", f"Failed to create movie tables: {e}") conn.commit() _log("DEBUG", "Database schema verified and up to date") def get_connection(self) -> sqlite3.Connection: """Get database connection with proper settings""" conn = sqlite3.Connection(str(self.db_path)) conn.execute("PRAGMA foreign_keys=ON;") conn.execute("PRAGMA busy_timeout = 30000;") # 30 second timeout return conn # TV Series Operations def upsert_series(self, imdb_id: str, series_path: str) -> None: """Insert or update series record""" with self.get_connection() as conn: now = datetime.now(timezone.utc).isoformat(timespec="seconds") conn.execute( "INSERT INTO series(imdb_id, series_path, last_applied) VALUES(?,?,?) " "ON CONFLICT(imdb_id) DO UPDATE SET series_path=excluded.series_path, last_applied=excluded.last_applied", (imdb_id, series_path, now) ) conn.commit() def get_series_episodes(self, imdb_id: str, has_video_file_only: bool = False) -> List[Dict[str, Any]]: """Get all episodes for a series""" with self.get_connection() as conn: if has_video_file_only: query = "SELECT season, episode, aired, dateadded, source FROM episode_dates WHERE imdb_id = ? AND has_video_file = 1" else: query = "SELECT season, episode, aired, dateadded, source FROM episode_dates WHERE imdb_id = ?" rows = conn.execute(query, (imdb_id,)).fetchall() return [ { "season": row[0], "episode": row[1], "aired": row[2], "dateadded": row[3], "source": row[4] } for row in rows ] def get_episode_date(self, imdb_id: str, season: int, episode: int) -> Optional[Dict[str, Any]]: """Get episode date information for a specific episode""" with self.get_connection() as conn: row = conn.execute( "SELECT aired, dateadded, source, has_video_file FROM episode_dates WHERE imdb_id = ? AND season = ? AND episode = ?", (imdb_id, season, episode) ).fetchone() if row: return { "aired": row[0], "dateadded": row[1], "source": row[2], "has_video_file": bool(row[3]) } return None def upsert_episode_date( self, imdb_id: str, season: int, episode: int, aired: Optional[str] = None, dateadded: Optional[str] = None, source: Optional[str] = None, has_video_file: bool = False ) -> None: """Insert or update episode date record""" with self.get_connection() as conn: conn.execute( "INSERT INTO episode_dates(imdb_id, season, episode, aired, dateadded, source, has_video_file) " "VALUES(?,?,?,?,?,?,?) " "ON CONFLICT(imdb_id, season, episode) DO UPDATE SET " "aired=excluded.aired, dateadded=excluded.dateadded, source=excluded.source, has_video_file=excluded.has_video_file", (imdb_id, season, episode, aired, dateadded, source, 1 if has_video_file else 0) ) conn.commit() def cleanup_orphaned_episodes(self, existing_episodes: Dict[str, List[Tuple[int, int]]]) -> int: """ Remove episode records that no longer have video files on disk. Args: existing_episodes: Dict mapping imdb_id -> [(season, episode), ...] """ removed_count = 0 with self.get_connection() as conn: # Get all series from database series_rows = conn.execute("SELECT imdb_id FROM series").fetchall() for (imdb_id,) in series_rows: if imdb_id not in existing_episodes: # Series no longer exists on disk removed = conn.execute("DELETE FROM episode_dates WHERE imdb_id = ?", (imdb_id,)).rowcount removed_count += removed _log("INFO", f"Removed {removed} episodes for missing series {imdb_id}") continue # Get episodes in database for this series db_episodes = conn.execute( "SELECT season, episode FROM episode_dates WHERE imdb_id = ?", (imdb_id,) ).fetchall() disk_episodes = set(existing_episodes[imdb_id]) for season, episode in db_episodes: if (season, episode) not in disk_episodes: conn.execute( "DELETE FROM episode_dates WHERE imdb_id = ? AND season = ? AND episode = ?", (imdb_id, season, episode) ) removed_count += 1 _log("DEBUG", f"Removed orphaned episode {imdb_id} S{season:02d}E{episode:02d}") conn.commit() return removed_count # Movie Operations def upsert_movie(self, imdb_id: str, movie_path: str) -> None: """Insert or update movie record""" with self.get_connection() as conn: now = datetime.now(timezone.utc).isoformat(timespec="seconds") conn.execute( "INSERT INTO movies(imdb_id, movie_path, last_applied) VALUES(?,?,?) " "ON CONFLICT(imdb_id) DO UPDATE SET movie_path=excluded.movie_path, last_applied=excluded.last_applied", (imdb_id, movie_path, now) ) conn.commit() def get_movie_dates(self, imdb_id: str) -> Optional[Dict[str, Any]]: """Get stored movie dates""" with self.get_connection() as conn: row = conn.execute( "SELECT released, dateadded, source, has_video_file FROM movie_dates WHERE imdb_id = ?", (imdb_id,) ).fetchone() if row: return { "released": row[0], "dateadded": row[1], "source": row[2], "has_video_file": bool(row[3]) } return None def upsert_movie_dates( self, imdb_id: str, released: Optional[str] = None, dateadded: Optional[str] = None, source: Optional[str] = None, has_video_file: bool = False ) -> None: """Insert or update movie dates""" with self.get_connection() as conn: conn.execute( "INSERT INTO movie_dates(imdb_id, released, dateadded, source, has_video_file) VALUES(?,?,?,?,?) " "ON CONFLICT(imdb_id) DO UPDATE SET " "released=excluded.released, dateadded=excluded.dateadded, source=excluded.source, has_video_file=excluded.has_video_file", (imdb_id, released, dateadded, source, 1 if has_video_file else 0) ) conn.commit() def cleanup_orphaned_movies(self, existing_movies: List[str]) -> int: """Remove movie records that no longer exist on disk""" if not existing_movies: return 0 removed_count = 0 existing_set = set(existing_movies) with self.get_connection() as conn: # Get all movies from database db_movies = conn.execute("SELECT imdb_id FROM movies").fetchall() for (imdb_id,) in db_movies: if imdb_id not in existing_set: # Movie no longer exists on disk conn.execute("DELETE FROM movies WHERE imdb_id = ?", (imdb_id,)) conn.execute("DELETE FROM movie_dates WHERE imdb_id = ?", (imdb_id,)) removed_count += 1 _log("DEBUG", f"Removed orphaned movie {imdb_id}") conn.commit() return removed_count # Statistics and Reporting def get_stats(self) -> Dict[str, Any]: """Get database statistics""" with self.get_connection() as conn: stats = {} # TV stats stats["tv_series_count"] = conn.execute("SELECT COUNT(*) FROM series").fetchone()[0] stats["tv_episode_count"] = conn.execute("SELECT COUNT(*) FROM episode_dates").fetchone()[0] stats["tv_episodes_with_files"] = conn.execute("SELECT COUNT(*) FROM episode_dates WHERE has_video_file = 1").fetchone()[0] # Movie stats stats["movie_count"] = conn.execute("SELECT COUNT(*) FROM movies").fetchone()[0] stats["movies_with_files"] = conn.execute("SELECT COUNT(*) FROM movie_dates WHERE has_video_file = 1").fetchone()[0] # Source breakdowns tv_sources = conn.execute(""" SELECT source, COUNT(*) as count FROM episode_dates WHERE has_video_file = 1 GROUP BY source ORDER BY count DESC """).fetchall() stats["tv_sources"] = dict(tv_sources) movie_sources = conn.execute(""" SELECT source, COUNT(*) as count FROM movie_dates WHERE has_video_file = 1 GROUP BY source ORDER BY count DESC """).fetchall() stats["movie_sources"] = dict(movie_sources) return stats def vacuum_database(self) -> None: """Vacuum database to reclaim space""" with self.get_connection() as conn: conn.execute("VACUUM;") _log("INFO", "Database vacuumed") # Utility Methods def clear_all_data(self) -> None: """Clear all data (keep schema) - useful for testing""" with self.get_connection() as conn: conn.execute("DELETE FROM episode_dates;") conn.execute("DELETE FROM series;") conn.execute("DELETE FROM movie_dates;") conn.execute("DELETE FROM movies;") conn.execute("DELETE FROM meta;") conn.commit() _log("WARNING", "All database data cleared") def export_series_episodes(self, imdb_id: str) -> List[Dict[str, Any]]: """Export all episode data for a series (for debugging)""" episodes = self.get_series_episodes(imdb_id) _log("INFO", f"Series {imdb_id} has {len(episodes)} episodes in database") return episodes def check_database_health(self) -> Dict[str, Any]: """Check database health and integrity""" health = {"status": "healthy", "issues": []} try: with self.get_connection() as conn: # Check integrity integrity = conn.execute("PRAGMA integrity_check;").fetchone()[0] if integrity != "ok": health["status"] = "corrupted" health["issues"].append(f"Integrity check failed: {integrity}") # Check for orphaned records orphaned_episodes = conn.execute(""" SELECT COUNT(*) FROM episode_dates WHERE imdb_id NOT IN (SELECT imdb_id FROM series) """).fetchone()[0] if orphaned_episodes > 0: health["issues"].append(f"{orphaned_episodes} orphaned episodes") orphaned_movie_dates = conn.execute(""" SELECT COUNT(*) FROM movie_dates WHERE imdb_id NOT IN (SELECT imdb_id FROM movies) """).fetchone()[0] if orphaned_movie_dates > 0: health["issues"].append(f"{orphaned_movie_dates} orphaned movie dates") health["orphaned_episodes"] = orphaned_episodes health["orphaned_movie_dates"] = orphaned_movie_dates except Exception as e: health["status"] = "error" health["error"] = str(e) return health if __name__ == "__main__": # Test the database test_db = NFOGuardDatabase(Path("/tmp/test_nfoguard.db")) # Test TV operations test_db.upsert_series("tt1234567", "/media/tv/Test Series [imdb-tt1234567]") test_db.upsert_episode_date("tt1234567", 1, 1, "2020-01-01", "2025-07-07T12:00:00+00:00", "test:import", True) episodes = test_db.get_series_episodes("tt1234567") print(f"Test series episodes: {episodes}") # Test movie operations test_db.upsert_movie("tt7654321", "/media/movies/Test Movie [imdb-tt7654321]") test_db.upsert_movie_dates("tt7654321", "2020-01-01", "2025-07-07T12:00:00+00:00", "test:import", True) movie_dates = test_db.get_movie_dates("tt7654321") print(f"Test movie dates: {movie_dates}") # Get stats stats = test_db.get_stats() print(f"Database stats: {stats}") # Check health health = test_db.check_database_health() print(f"Database health: {health}") print("Database test completed successfully!")