125 lines
3.9 KiB
Python
125 lines
3.9 KiB
Python
import typing
|
|
import sqlite3
|
|
import logging
|
|
from contextlib import contextmanager
|
|
from utils import log
|
|
|
|
|
|
DB_NAME = "image.db"
|
|
|
|
|
|
@contextmanager
|
|
def db_connect() -> typing.Generator[sqlite3.Cursor, None, typing.Any]:
|
|
connection = sqlite3.connect(database=DB_NAME)
|
|
cursor = connection.cursor()
|
|
|
|
try:
|
|
yield cursor
|
|
except sqlite3.Error as e:
|
|
error_msg = str(e)
|
|
log(f"Error when try to connecting the database: {error_msg}", logging.ERROR)
|
|
finally:
|
|
cursor.close()
|
|
connection.close()
|
|
|
|
|
|
def create_table() -> None:
|
|
query_sql = """CREATE TABLE IF NOT EXISTS images (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
path TEXT NOT NULL,
|
|
output_path TEXT NOT NULL,
|
|
target_size INTEGER NOT NULL,
|
|
initial_quality INTEGER NOT NULL,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL
|
|
)"""
|
|
|
|
with db_connect() as cursor:
|
|
try:
|
|
cursor.execute(query_sql)
|
|
log("Database successfully created")
|
|
cursor.connection.commit()
|
|
except sqlite3.Error as e:
|
|
error_msg = str(e)
|
|
log(f"Error while creating the database table: {error_msg}", logging.ERROR)
|
|
|
|
|
|
def insert_image_db(
|
|
path: str, output_path: str, target_size: int, quality: int
|
|
) -> None:
|
|
with db_connect() as cursor:
|
|
try:
|
|
cursor.execute(
|
|
"""
|
|
INSERT INTO images (path, output_path, target_size, initial_quality) VALUES (?, ?, ?, ?)
|
|
""",
|
|
(path, output_path, target_size, quality),
|
|
)
|
|
cursor.connection.commit()
|
|
log("Successfully inserting the new records")
|
|
except sqlite3.Error as e:
|
|
error_msg = str(e)
|
|
log(
|
|
f"Error while inserting the image to the database: {error_msg}",
|
|
logging.ERROR,
|
|
)
|
|
|
|
|
|
def get_image_db(image_id: int) -> typing.Any:
|
|
with db_connect() as cursor:
|
|
try:
|
|
cursor.execute("SELECT * FROM images WHERE id = ?", (image_id,))
|
|
image_data = cursor.fetchone()
|
|
if image_data is not None:
|
|
return image_data
|
|
else:
|
|
return None
|
|
except sqlite3.Error as e:
|
|
error_msg = str(e)
|
|
log(
|
|
f"Error while fetching the image from the database: {error_msg}",
|
|
logging.ERROR,
|
|
)
|
|
|
|
|
|
def get_all_images_db(limit: int = None, order: str = None) -> list:
|
|
with db_connect() as cursor:
|
|
try:
|
|
if limit is None and order is None:
|
|
cursor.execute("SELECT * FROM images ORDER BY id DESC")
|
|
else:
|
|
cursor.execute(
|
|
"SELECT * FROM images ORDER BY ? LIMIT ?",
|
|
(order, limit),
|
|
)
|
|
all_images = cursor.fetchall()
|
|
return [image for image in all_images]
|
|
except sqlite3.Error as e:
|
|
error_msg = str(e)
|
|
log(
|
|
f"Error while fetch all of the images from the database: {error_msg}",
|
|
logging.ERROR,
|
|
)
|
|
|
|
|
|
def delete_image_db(image_id: int) -> bool:
|
|
with db_connect() as cursor:
|
|
try:
|
|
cursor.execute("SELECT output_path FROM images WHERE id = ?", image_id)
|
|
result = cursor.fetchone()
|
|
print(result)
|
|
if result:
|
|
cursor.execute("DELETE FROM images WHERE id = ?", image_id)
|
|
cursor.connection.commit()
|
|
log("Successfully delete the record")
|
|
return True
|
|
else:
|
|
log(
|
|
"Image is not found either on the database or local",
|
|
logging.WARNING,
|
|
)
|
|
return False
|
|
except sqlite3.Error as e:
|
|
error_msg = str(e)
|
|
log(f"Unable to delete the record: {error_msg}", logging.ERROR)
|
|
return False
|