forum-scrapper/Programm/Database/image_table.py

"""
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()