""" Image table class for database operations. This module provides the ImageTable class that handles image-specific database operations for storing and retrieving Image objects. """ import sqlite3 from typing import Optional from ..Datamodel.Image import Image class ImageTable: """ Table class for managing image data in the database. Attributes: connection: SQLite database connection cursor: Database cursor for executing queries """ def __init__(self, connection: sqlite3.Connection): """ Initialize the image table with a database connection. Args: connection: SQLite database connection """ self.connection = connection self.cursor = self.connection.cursor() def create(self): """Create the images table if it doesn't exist.""" self.cursor.execute(''' CREATE TABLE IF NOT EXISTS images ( id INTEGER PRIMARY KEY AUTOINCREMENT, message_id TEXT NOT NULL, url TEXT NOT NULL, name TEXT NOT NULL, type INTEGER NOT NULL, width INTEGER NOT NULL, height INTEGER NOT NULL, size INTEGER NOT NULL, duration TEXT, caption TEXT, data BLOB, FOREIGN KEY (message_id) REFERENCES messages (id) ON DELETE CASCADE ) ''') self.connection.commit() def exists(self, image_url: str) -> bool: """ Check if an image exists in the database. Args: image_url: Image URL to check Returns: True if image exists, False otherwise """ self.cursor.execute('SELECT 1 FROM images WHERE url = ?', (image_url,)) return self.cursor.fetchone() is not None def save(self, image: Image) -> bool: """ Save an image to the database. Args: image: Image object to save message_id: Message ID that this image belongs to Returns: True if image was saved, False if it already exists """ # Check if image exists if self.exists(image.url): # Load existing image to check if it has data existing_image = self.load(image.url, include_data=False) # If existing image has no data but new image has data, update it if existing_image and existing_image.data is None and image.data is not None: self.cursor.execute(''' UPDATE images SET message_id = ?, url = ?, name = ?, type = ?, width = ?, height = ?, size = ?, duration = ?, caption = ?, data = ? WHERE url = ? ''', ( image.ref_id, image.url, image.name, image.type, image.width, image.height, image.size, image.duration, image.caption, image.data, image.url )) self.connection.commit() return True else: # Image exists with data, skip return False # Image doesn't exist, insert it self.cursor.execute(''' INSERT INTO images (message_id, url, name, type, width, height, size, duration, caption, data) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ''', ( image.ref_id, image.url, image.name, image.type, image.width, image.height, image.size, image.duration, image.caption, image.data )) self.connection.commit() return True def load(self, image_url: str, include_data: bool = True) -> Optional[Image]: """ Load an image from the database by URL. Args: image_url: Image URL to load include_data: Whether to include the blob data in the loaded image (default: True) Returns: Image object if found, None otherwise """ # Conditionally select data column based on include_blob parameter if include_data: self.cursor.execute(''' SELECT message_id, url, name, type, width, height, size, duration, caption, data FROM images WHERE url = ? ''', (image_url,)) else: self.cursor.execute(''' SELECT message_id, url, name, type, width, height, size, duration, caption FROM images WHERE url = ? ''', (image_url,)) row = self.cursor.fetchone() if row is None: return None # Conditionally include data based on include_blob parameter if include_data: return Image( ref_id=row[0], url=row[1], name=row[2], type=row[3], width=row[4], height=row[5], size=row[6], duration=row[7], caption=row[8], data=row[9] ) else: return Image( ref_id=row[0], url=row[1], name=row[2], type=row[3], width=row[4], height=row[5], size=row[6], duration=row[7], caption=row[8], data=None ) def commit(self): """Commit pending changes to the database.""" self.connection.commit() def close(self): """Close the cursor.""" self.cursor.close()