This commit is contained in:
@@ -0,0 +1,582 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Direct Sonarr Database Client for NFOGuard
|
||||
Provides high-performance access to Sonarr's SQLite/PostgreSQL database
|
||||
"""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
import psycopg2
|
||||
import psycopg2.extras
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Dict, Any, List, Optional, Tuple, Union
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from core.logging import _log
|
||||
|
||||
|
||||
class SonarrDbClient:
|
||||
"""Direct database client for Sonarr's SQLite or PostgreSQL database"""
|
||||
|
||||
def __init__(self,
|
||||
db_type: str = "sqlite",
|
||||
db_path: Optional[str] = None,
|
||||
db_host: Optional[str] = None,
|
||||
db_port: Optional[int] = None,
|
||||
db_name: Optional[str] = None,
|
||||
db_user: Optional[str] = None,
|
||||
db_password: Optional[str] = None):
|
||||
"""
|
||||
Initialize Sonarr database client
|
||||
|
||||
Args:
|
||||
db_type: "sqlite" or "postgresql"
|
||||
db_path: Path to SQLite database file
|
||||
db_host: PostgreSQL host
|
||||
db_port: PostgreSQL port
|
||||
db_name: PostgreSQL database name
|
||||
db_user: PostgreSQL username
|
||||
db_password: PostgreSQL password
|
||||
"""
|
||||
self.db_type = db_type.lower()
|
||||
self.db_path = db_path
|
||||
self.db_host = db_host
|
||||
self.db_port = db_port or 5432
|
||||
self.db_name = db_name
|
||||
self.db_user = db_user
|
||||
self.db_password = db_password
|
||||
|
||||
self._test_connection()
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> Optional['SonarrDbClient']:
|
||||
"""Create client from environment variables"""
|
||||
db_type = os.environ.get("SONARR_DB_TYPE", "").lower()
|
||||
|
||||
if not db_type:
|
||||
return None
|
||||
|
||||
if db_type == "sqlite":
|
||||
db_path = os.environ.get("SONARR_DB_PATH")
|
||||
if not db_path or not Path(db_path).exists():
|
||||
_log("WARNING", f"SONARR_DB_PATH not found or invalid: {db_path}")
|
||||
return None
|
||||
return cls(db_type="sqlite", db_path=db_path)
|
||||
|
||||
elif db_type == "postgresql":
|
||||
# Support both individual vars and connection string
|
||||
db_url = os.environ.get("SONARR_DB_URL")
|
||||
if db_url:
|
||||
parsed = urlparse(db_url)
|
||||
return cls(
|
||||
db_type="postgresql",
|
||||
db_host=parsed.hostname,
|
||||
db_port=parsed.port or 5432,
|
||||
db_name=parsed.path.lstrip('/'),
|
||||
db_user=parsed.username,
|
||||
db_password=parsed.password
|
||||
)
|
||||
else:
|
||||
return cls(
|
||||
db_type="postgresql",
|
||||
db_host=os.environ.get("SONARR_DB_HOST"),
|
||||
db_port=int(os.environ.get("SONARR_DB_PORT", "5432")),
|
||||
db_name=os.environ.get("SONARR_DB_NAME"),
|
||||
db_user=os.environ.get("SONARR_DB_USER"),
|
||||
db_password=os.environ.get("SONARR_DB_PASSWORD")
|
||||
)
|
||||
else:
|
||||
_log("ERROR", f"Unsupported database type: {db_type}")
|
||||
return None
|
||||
|
||||
def _test_connection(self) -> None:
|
||||
"""Test database connection on initialization"""
|
||||
try:
|
||||
conn = self._get_connection()
|
||||
if conn:
|
||||
conn.close()
|
||||
_log("INFO", f"Connected to Sonarr {self.db_type} database successfully")
|
||||
else:
|
||||
raise Exception("Failed to create connection")
|
||||
except Exception as e:
|
||||
_log("ERROR", f"Failed to connect to Sonarr database: {e}")
|
||||
raise
|
||||
|
||||
def _get_connection(self) -> Union[sqlite3.Connection, psycopg2.extensions.connection]:
|
||||
"""Get database connection"""
|
||||
if self.db_type == "sqlite":
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
conn.row_factory = sqlite3.Row
|
||||
return conn
|
||||
elif self.db_type == "postgresql":
|
||||
conn = psycopg2.connect(
|
||||
host=self.db_host,
|
||||
port=self.db_port,
|
||||
database=self.db_name,
|
||||
user=self.db_user,
|
||||
password=self.db_password
|
||||
)
|
||||
return conn
|
||||
else:
|
||||
raise ValueError(f"Unsupported database type: {self.db_type}")
|
||||
|
||||
def get_series_by_imdb(self, imdb_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Find series by IMDb ID using database query
|
||||
|
||||
Returns:
|
||||
Dictionary with series info including id, imdb_id, title, path
|
||||
"""
|
||||
imdb_id = imdb_id if imdb_id.startswith("tt") else f"tt{imdb_id}"
|
||||
|
||||
query = """
|
||||
SELECT
|
||||
"Id" as id,
|
||||
"ImdbId" as imdb_id,
|
||||
"TvdbId" as tvdb_id,
|
||||
"Title" as title,
|
||||
"Path" as path,
|
||||
"Added" as added
|
||||
FROM "Series"
|
||||
WHERE "ImdbId" = %s
|
||||
"""
|
||||
|
||||
if self.db_type == "sqlite":
|
||||
query = query.replace("%s", "?")
|
||||
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
if self.db_type == "postgresql":
|
||||
cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
|
||||
else:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute(query, (imdb_id,))
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row:
|
||||
return dict(row) if self.db_type == "sqlite" else row
|
||||
|
||||
except Exception as e:
|
||||
_log("ERROR", f"Database query error for IMDb {imdb_id}: {e}")
|
||||
|
||||
return None
|
||||
|
||||
def get_all_episodes_for_series(self, series_id: int) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get all episodes for a series
|
||||
|
||||
Args:
|
||||
series_id: Sonarr series ID
|
||||
|
||||
Returns:
|
||||
List of episode dictionaries with season, episode, air_date
|
||||
"""
|
||||
query = """
|
||||
SELECT
|
||||
"Id" as id,
|
||||
"SeasonNumber" as season,
|
||||
"EpisodeNumber" as episode,
|
||||
"Title" as title,
|
||||
"AirDate" as air_date,
|
||||
"EpisodeFileId" as episode_file_id
|
||||
FROM "Episodes"
|
||||
WHERE "SeriesId" = %s
|
||||
ORDER BY "SeasonNumber", "EpisodeNumber"
|
||||
"""
|
||||
|
||||
if self.db_type == "sqlite":
|
||||
query = query.replace("%s", "?")
|
||||
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
if self.db_type == "postgresql":
|
||||
cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
|
||||
else:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute(query, (series_id,))
|
||||
rows = cursor.fetchall()
|
||||
|
||||
if self.db_type == "sqlite":
|
||||
return [dict(row) for row in rows]
|
||||
else:
|
||||
return rows
|
||||
|
||||
except Exception as e:
|
||||
_log("ERROR", f"Database query error for series {series_id}: {e}")
|
||||
return []
|
||||
|
||||
def get_episode_import_date(self, episode_id: int) -> Tuple[Optional[str], str]:
|
||||
"""
|
||||
Get earliest import date for an episode from History table
|
||||
|
||||
Args:
|
||||
episode_id: Sonarr episode ID
|
||||
|
||||
Returns:
|
||||
(date_iso, source_description)
|
||||
"""
|
||||
# Query for earliest import event (EventType 3)
|
||||
import_query = """
|
||||
SELECT
|
||||
"Date" as event_date,
|
||||
"EventType" as event_type
|
||||
FROM "History"
|
||||
WHERE "EpisodeId" = %s
|
||||
AND "EventType" = 3
|
||||
ORDER BY "Date" ASC
|
||||
LIMIT 1
|
||||
"""
|
||||
|
||||
if self.db_type == "sqlite":
|
||||
import_query = import_query.replace("%s", "?")
|
||||
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
if self.db_type == "postgresql":
|
||||
cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
|
||||
else:
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Try import events first
|
||||
cursor.execute(import_query, (episode_id,))
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row:
|
||||
event_date = row['event_date'] if self.db_type == "postgresql" else row[0]
|
||||
if isinstance(event_date, str):
|
||||
dt = datetime.fromisoformat(event_date.replace("Z", "+00:00"))
|
||||
else:
|
||||
dt = event_date.replace(tzinfo=timezone.utc)
|
||||
|
||||
date_iso = dt.astimezone(timezone.utc).isoformat(timespec="seconds")
|
||||
return date_iso, "sonarr:db.history.import"
|
||||
|
||||
except Exception as e:
|
||||
_log("ERROR", f"Database query error for episode {episode_id}: {e}")
|
||||
|
||||
return None, "sonarr:db.no_import_found"
|
||||
|
||||
def get_episode_file_date(self, series_id: int, season: int, episode: int) -> Optional[str]:
|
||||
"""
|
||||
Get episode file DateAdded as fallback
|
||||
|
||||
Args:
|
||||
series_id: Sonarr series ID
|
||||
season: Season number
|
||||
episode: Episode number
|
||||
|
||||
Returns:
|
||||
ISO date string or None
|
||||
"""
|
||||
# First get the episode ID to find the file
|
||||
episode_query = """
|
||||
SELECT "EpisodeFileId"
|
||||
FROM "Episodes"
|
||||
WHERE "SeriesId" = %s
|
||||
AND "SeasonNumber" = %s
|
||||
AND "EpisodeNumber" = %s
|
||||
"""
|
||||
|
||||
file_query = """
|
||||
SELECT "DateAdded"
|
||||
FROM "EpisodeFiles"
|
||||
WHERE "Id" = %s
|
||||
"""
|
||||
|
||||
if self.db_type == "sqlite":
|
||||
episode_query = episode_query.replace("%s", "?")
|
||||
file_query = file_query.replace("%s", "?")
|
||||
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
if self.db_type == "postgresql":
|
||||
cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
|
||||
else:
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Get episode file ID
|
||||
cursor.execute(episode_query, (series_id, season, episode))
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row:
|
||||
file_id = row['EpisodeFileId'] if self.db_type == "postgresql" else row[0]
|
||||
|
||||
if file_id:
|
||||
# Get file date
|
||||
cursor.execute(file_query, (file_id,))
|
||||
file_row = cursor.fetchone()
|
||||
|
||||
if file_row:
|
||||
date_value = file_row['DateAdded'] if self.db_type == "postgresql" else file_row[0]
|
||||
|
||||
if isinstance(date_value, str):
|
||||
dt = datetime.fromisoformat(date_value.replace("Z", "+00:00"))
|
||||
else:
|
||||
dt = date_value.replace(tzinfo=timezone.utc)
|
||||
|
||||
return dt.astimezone(timezone.utc).isoformat(timespec="seconds")
|
||||
|
||||
except Exception as e:
|
||||
_log("ERROR", f"Database query error for episode file: {e}")
|
||||
|
||||
return None
|
||||
|
||||
def bulk_import_dates_for_series(self, series_id: int) -> Dict[Tuple[int, int], Tuple[Optional[str], str]]:
|
||||
"""
|
||||
Get import dates for all episodes in a series in a single query
|
||||
|
||||
Args:
|
||||
series_id: Sonarr series ID
|
||||
|
||||
Returns:
|
||||
Dictionary mapping (season, episode) -> (date_iso, source)
|
||||
"""
|
||||
query = """
|
||||
SELECT
|
||||
e."SeasonNumber" as season,
|
||||
e."EpisodeNumber" as episode,
|
||||
e."Id" as episode_id,
|
||||
MIN(h."Date") as earliest_import
|
||||
FROM "Episodes" e
|
||||
LEFT JOIN "History" h ON e."Id" = h."EpisodeId" AND h."EventType" = 3
|
||||
WHERE e."SeriesId" = %s
|
||||
GROUP BY e."SeasonNumber", e."EpisodeNumber", e."Id"
|
||||
ORDER BY e."SeasonNumber", e."EpisodeNumber"
|
||||
"""
|
||||
|
||||
if self.db_type == "sqlite":
|
||||
query = query.replace("%s", "?")
|
||||
|
||||
results = {}
|
||||
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
if self.db_type == "postgresql":
|
||||
cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
|
||||
else:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute(query, (series_id,))
|
||||
rows = cursor.fetchall()
|
||||
|
||||
for row in rows:
|
||||
if self.db_type == "postgresql":
|
||||
season, episode, episode_id, earliest_import = row['season'], row['episode'], row['episode_id'], row['earliest_import']
|
||||
else:
|
||||
season, episode, episode_id, earliest_import = row[0], row[1], row[2], row[3]
|
||||
|
||||
if earliest_import:
|
||||
if isinstance(earliest_import, str):
|
||||
dt = datetime.fromisoformat(earliest_import.replace("Z", "+00:00"))
|
||||
else:
|
||||
dt = earliest_import.replace(tzinfo=timezone.utc)
|
||||
|
||||
date_iso = dt.astimezone(timezone.utc).isoformat(timespec="seconds")
|
||||
results[(season, episode)] = (date_iso, "sonarr:db.bulk.import")
|
||||
else:
|
||||
results[(season, episode)] = (None, "sonarr:db.bulk.no_import")
|
||||
|
||||
except Exception as e:
|
||||
_log("ERROR", f"Bulk query error for series {series_id}: {e}")
|
||||
|
||||
return results
|
||||
|
||||
def get_database_stats(self) -> Dict[str, Any]:
|
||||
"""Get basic statistics about the Sonarr database"""
|
||||
stats = {}
|
||||
|
||||
queries = {
|
||||
"total_series": 'SELECT COUNT(*) FROM "Series"',
|
||||
"total_episodes": 'SELECT COUNT(*) FROM "Episodes"',
|
||||
"total_episode_files": 'SELECT COUNT(*) FROM "EpisodeFiles"',
|
||||
"total_history_events": 'SELECT COUNT(*) FROM "History"',
|
||||
"import_events": 'SELECT COUNT(*) FROM "History" WHERE "EventType" = 3',
|
||||
"grab_events": 'SELECT COUNT(*) FROM "History" WHERE "EventType" = 1'
|
||||
}
|
||||
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
for stat_name, query in queries.items():
|
||||
cursor.execute(query)
|
||||
result = cursor.fetchone()
|
||||
stats[stat_name] = result[0] if result else 0
|
||||
|
||||
except Exception as e:
|
||||
_log("ERROR", f"Stats query error: {e}")
|
||||
stats["error"] = str(e)
|
||||
|
||||
return stats
|
||||
|
||||
def health_check(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Comprehensive health check for the Sonarr database connection
|
||||
|
||||
Returns:
|
||||
Dictionary with health status, connection info, and basic functionality tests
|
||||
"""
|
||||
health = {
|
||||
"status": "healthy",
|
||||
"database_type": self.db_type,
|
||||
"connection": "ok",
|
||||
"readable": False,
|
||||
"tables_exist": False,
|
||||
"sample_data": False,
|
||||
"issues": [],
|
||||
"tested_at": datetime.now(timezone.utc).isoformat(timespec="seconds")
|
||||
}
|
||||
|
||||
try:
|
||||
with self._get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Test 1: Basic read
|
||||
try:
|
||||
cursor.execute('SELECT 1')
|
||||
result = cursor.fetchone()
|
||||
if result and result[0] == 1:
|
||||
health["readable"] = True
|
||||
health["connection"] = "readable"
|
||||
else:
|
||||
health["issues"].append("Basic SELECT query failed")
|
||||
except Exception as e:
|
||||
health["issues"].append(f"Read test failed: {e}")
|
||||
health["status"] = "degraded"
|
||||
|
||||
# Test 2: Check required tables
|
||||
required_tables = ["Series", "Episodes", "EpisodeFiles", "History"]
|
||||
existing_tables = []
|
||||
|
||||
try:
|
||||
if self.db_type == "postgresql":
|
||||
cursor.execute("""
|
||||
SELECT table_name
|
||||
FROM information_schema.tables
|
||||
WHERE table_schema = 'public'
|
||||
AND table_name IN ('Series', 'Episodes', 'EpisodeFiles', 'History')
|
||||
""")
|
||||
else: # SQLite
|
||||
cursor.execute("""
|
||||
SELECT name
|
||||
FROM sqlite_master
|
||||
WHERE type='table'
|
||||
AND name IN ('Series', 'Episodes', 'EpisodeFiles', 'History')
|
||||
""")
|
||||
|
||||
rows = cursor.fetchall()
|
||||
existing_tables = [row[0] for row in rows]
|
||||
|
||||
if len(existing_tables) == len(required_tables):
|
||||
health["tables_exist"] = True
|
||||
else:
|
||||
missing = set(required_tables) - set(existing_tables)
|
||||
health["issues"].append(f"Missing tables: {list(missing)}")
|
||||
health["status"] = "degraded"
|
||||
|
||||
health["existing_tables"] = existing_tables
|
||||
|
||||
except Exception as e:
|
||||
health["issues"].append(f"Table check failed: {e}")
|
||||
health["status"] = "degraded"
|
||||
|
||||
# Test 3: Check for sample data
|
||||
if health["tables_exist"]:
|
||||
try:
|
||||
cursor.execute('SELECT COUNT(*) FROM "Series"')
|
||||
series_count = cursor.fetchone()[0]
|
||||
|
||||
cursor.execute('SELECT COUNT(*) FROM "Episodes"')
|
||||
episode_count = cursor.fetchone()[0]
|
||||
|
||||
cursor.execute('SELECT COUNT(*) FROM "History"')
|
||||
history_count = cursor.fetchone()[0]
|
||||
|
||||
if series_count > 0 and episode_count > 0:
|
||||
health["sample_data"] = True
|
||||
health["series_count"] = series_count
|
||||
health["episode_count"] = episode_count
|
||||
health["history_count"] = history_count
|
||||
else:
|
||||
health["issues"].append(f"Low data counts - Series: {series_count}, Episodes: {episode_count}")
|
||||
|
||||
except Exception as e:
|
||||
health["issues"].append(f"Sample data check failed: {e}")
|
||||
|
||||
# Test 4: Test a real query
|
||||
if health["sample_data"]:
|
||||
try:
|
||||
cursor.execute("""
|
||||
SELECT COUNT(*)
|
||||
FROM "Series"
|
||||
WHERE "ImdbId" IS NOT NULL
|
||||
""")
|
||||
imdb_series = cursor.fetchone()[0]
|
||||
health["series_with_imdb"] = imdb_series
|
||||
|
||||
if imdb_series > 0:
|
||||
health["functional"] = True
|
||||
else:
|
||||
health["issues"].append("No series with IMDb IDs found")
|
||||
|
||||
except Exception as e:
|
||||
health["issues"].append(f"Functional test failed: {e}")
|
||||
health["status"] = "degraded"
|
||||
|
||||
except Exception as e:
|
||||
health["status"] = "error"
|
||||
health["connection"] = "failed"
|
||||
health["issues"].append(f"Connection failed: {e}")
|
||||
_log("ERROR", f"Database health check failed: {e}")
|
||||
|
||||
# Overall status
|
||||
if health["issues"]:
|
||||
if health["status"] == "healthy":
|
||||
health["status"] = "degraded"
|
||||
|
||||
# Add connection details (safe info only)
|
||||
health["connection_info"] = {
|
||||
"type": self.db_type,
|
||||
"host": self.db_host if self.db_type == "postgresql" else None,
|
||||
"port": self.db_port if self.db_type == "postgresql" else None,
|
||||
"database": self.db_name if self.db_type == "postgresql" else None,
|
||||
"path": self.db_path if self.db_type == "sqlite" else None
|
||||
}
|
||||
|
||||
return health
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Test the database client
|
||||
print("Testing SonarrDbClient...")
|
||||
|
||||
# Test with environment variables
|
||||
client = SonarrDbClient.from_env()
|
||||
if client:
|
||||
print("✅ Connected to Sonarr database")
|
||||
|
||||
# Test stats
|
||||
stats = client.get_database_stats()
|
||||
print(f"Database stats: {stats}")
|
||||
|
||||
# Test series lookup
|
||||
test_series = client.get_series_by_imdb("tt1628033") # Top Gear from your data
|
||||
if test_series:
|
||||
print(f"Found test series: {test_series}")
|
||||
|
||||
# Test episodes
|
||||
series_id = test_series['id']
|
||||
episodes = client.get_all_episodes_for_series(series_id)
|
||||
print(f"Found {len(episodes)} episodes")
|
||||
|
||||
# Test bulk import dates
|
||||
if episodes:
|
||||
import_dates = client.bulk_import_dates_for_series(series_id)
|
||||
print(f"Import dates for {len(import_dates)} episodes")
|
||||
else:
|
||||
print("Test series not found")
|
||||
else:
|
||||
print("❌ Could not connect to database - check environment variables")
|
||||
Reference in New Issue
Block a user