database/crud.py

import json
import os
from datetime import datetime
from .db import get_connection


def create_session(session_id, image_path, result_paths):
    conn = get_connection()
    c = conn.cursor()
    c.execute(
        "INSERT INTO sessions (session_id, image_path, result_paths, created_at) VALUES (?, ?, ?, ?)",
        (session_id, image_path, json.dumps(result_paths), datetime.utcnow().isoformat())
    )
    conn.commit()
    conn.close()


def get_session(session_id):
    conn = get_connection()
    c = conn.cursor()
    c.execute("SELECT image_path, result_paths FROM sessions WHERE session_id = ?", (session_id,))
    row = c.fetchone()
    conn.close()

    if row:
        image_path, result_paths_json = row
        return {
            "image_path": image_path,
            "result_paths": json.loads(result_paths_json)
        }
    return None


def delete_session(session_id):
    session = get_session(session_id)
    if session:

        # Delete input image
        if os.path.exists(session["image_path"]):
            os.remove(session["image_path"])

        # Delete output images
        for path in session["result_paths"]:
            if os.path.exists(path):
                os.remove(path)

        # Delete empty temporary folder
        if session["result_paths"]:
            session_dir = os.path.dirname(session["result_paths"][0])
            if os.path.exists(session_dir) and not os.listdir(session_dir):
                os.rmdir(session_dir)

    # Delete from DB
    conn = get_connection()
    c = conn.cursor()
    c.execute("DELETE FROM sessions WHERE session_id = ?", (session_id,))
    conn.commit()
    conn.close()