155 lines
4.6 KiB
Python
155 lines
4.6 KiB
Python
"""
|
|
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,
|
|
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
|
|
"""
|
|
if self.exists(image.url):
|
|
return False
|
|
|
|
self.cursor.execute('''
|
|
INSERT INTO images (message_id, url, name, type, width, height, size, duration, data)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
''', (
|
|
image.ref_id,
|
|
image.url,
|
|
image.name,
|
|
image.type,
|
|
image.width,
|
|
image.height,
|
|
image.size,
|
|
image.duration,
|
|
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, data
|
|
FROM images WHERE url = ?
|
|
''', (image_url,))
|
|
else:
|
|
self.cursor.execute('''
|
|
SELECT message_id, url, name, type, width, height, size, duration
|
|
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],
|
|
data=row[8]
|
|
)
|
|
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],
|
|
data=None
|
|
)
|
|
|
|
def commit(self):
|
|
"""Commit pending changes to the database."""
|
|
self.connection.commit()
|
|
|
|
def close(self):
|
|
"""Close the cursor."""
|
|
self.cursor.close() |