403 lines
16 KiB
Python
403 lines
16 KiB
Python
#!/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
|
|
|
|
|
|
def _log(level: str, msg: str):
|
|
"""Placeholder logging function - replace with your actual logger"""
|
|
print(f"[{datetime.now().isoformat()}] {level}: {msg}")
|
|
|
|
|
|
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 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!") |