# database.py import sqlite3 from pathlib import Path import logging logger = logging.getLogger(__name__) # Directory to store downloaded files DATA_DIR = Path(__file__).parent / "data" DATA_DIR.mkdir(parents=True, exist_ok=True) EMBEDDINGS_DB_PATH = DATA_DIR / "data.sqlite" def set_db_path(db_path: Path): """ Allows external code to override the default EMBEDDINGS_DB_PATH. Must be called before get_connection() or initialize_db() are used. """ global EMBEDDINGS_DB_PATH EMBEDDINGS_DB_PATH = db_path logger.info(f"[CONCURRENT] Overriding EMBEDDINGS_DB_PATH -> {EMBEDDINGS_DB_PATH}") def get_connection(): conn = sqlite3.connect(EMBEDDINGS_DB_PATH, timeout=2.0) return conn def initialize_db(): conn = get_connection() cursor = conn.cursor() # Create videos table cursor.execute(""" CREATE TABLE IF NOT EXISTS videos ( id INTEGER PRIMARY KEY AUTOINCREMENT, url TEXT UNIQUE NOT NULL, duration REAL NOT NULL, description TEXT, saved_up_to REAL DEFAULT 0.0, row INTEGER ) """) # Modify embeddings table to include video_id cursor.execute(""" CREATE TABLE IF NOT EXISTS embeddings ( id INTEGER PRIMARY KEY AUTOINCREMENT, video_id INTEGER NOT NULL, start_time REAL NOT NULL, FOREIGN KEY (video_id) REFERENCES videos(id) ) """) conn.commit() conn.close() logger.info("SQLite database initialized with updated schema") def add_video(url, duration, description, row): """ Adds a video to the database or retrieves it if already exists. Args: url (str): The URL of the video. duration (float): The duration of the video in seconds. description (str): A description of the video. row (int): The csv row the video came from Returns: dict: A dictionary containing the video's metadata. """ conn = get_connection() cursor = conn.cursor() try: cursor.execute(""" INSERT INTO videos (url, duration, description, row) VALUES (?, ?, ?, ?) """, (url, duration, description, row)) conn.commit() video_db_id = cursor.lastrowid logger.debug(f"Added video {url} with DB ID {video_db_id}") # Construct metadata dictionary for the newly added video video_metadata = { 'id': video_db_id, 'url': url, 'duration': duration, 'description': description, 'saved_up_to': 0.0, 'row': row } except sqlite3.IntegrityError: # Video already exists, fetch its metadata logger.debug(f"Video {url} already exists in the database.") video_metadata = get_video_by_url(url) if not video_metadata: logger.error(f"Failed to retrieve metadata for existing video {url}.") raise finally: conn.close() return video_metadata def get_video_by_url(url): """ Retrieves the video metadata from the database by URL. Args: url (str): The URL of the video. Returns: dict or None: A dictionary containing the video metadata if found, else None. """ conn = get_connection() cursor = conn.cursor() cursor.execute(""" SELECT id, url, duration, description, saved_up_to, row FROM videos WHERE url = ? """, (url,)) row = cursor.fetchone() conn.close() if row: keys = ['id', 'url', 'duration', 'description', 'saved_up_to', 'row'] video_metadata = dict(zip(keys, row)) logger.debug(f"Retrieved video metadata for URL {url}: {video_metadata}") return video_metadata else: logger.debug(f"No video found in database for URL {url}") return None def add_embedding(video_id, start_time): conn = get_connection() cursor = conn.cursor() cursor.execute(""" INSERT INTO embeddings (video_id, start_time) VALUES (?, ?) """, (video_id, start_time)) embedding_id = cursor.lastrowid conn.commit() conn.close() logger.debug(f"Added embedding ID {embedding_id} for video ID {video_id} at {start_time}s") return embedding_id def update_video_saved_up_to(video_id, new_time): conn = get_connection() cursor = conn.cursor() cursor.execute(""" UPDATE videos SET saved_up_to = ? WHERE id = ? """, (new_time, video_id)) conn.commit() conn.close() logger.debug(f"Updated video ID {video_id} saved_up_to to {new_time}s") def update_video_duration(video_id, duration): """ Updates the duration of a video in the database. Args: video_id (int): The ID of the video to update. duration (float): The new duration in seconds. """ conn = get_connection() cursor = conn.cursor() cursor.execute(""" UPDATE videos SET duration = ? WHERE id = ? """, (duration, video_id)) conn.commit() conn.close() logger.debug(f"Updated video ID {video_id} duration to {duration}s")