database_handler.py

import sqlite3
from contextlib import contextmanager
import os


class DatabaseHandler:
    def __init__(self, db_name="bulk_api_app.db"):
        self.db_name = db_name
        self.conn = None
        self.create_tables()

    @contextmanager
    def get_cursor(self):
        conn = sqlite3.connect(self.db_name)
        try:
            yield conn.cursor()
            conn.commit()
        finally:
            conn.close()

    def create_tables(self):
        with self.get_cursor() as cursor:
            cursor.execute(
                """
                CREATE TABLE IF NOT EXISTS files (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL,
                    path TEXT NOT NULL,
                    status TEXT DEFAULT 'Waiting for shipment',
                    openai_id TEXT,
                    batch_id TEXT,
                    result_path TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """
            )
            cursor.execute(
                """
                CREATE TABLE IF NOT EXISTS settings (
                    key TEXT PRIMARY KEY,
                    value TEXT
                )
            """
            )

            cursor.execute("PRAGMA table_info(files)")
            columns = [column[1] for column in cursor.fetchall()]
            if "created_at" not in columns:
                cursor.execute(
                    "ALTER TABLE files ADD COLUMN created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP"
                )

    def add_file(self, name, path, status="Waiting for shipment", openai_id=None):
        with self.get_cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO files (name, path, status, openai_id)
                VALUES (?, ?, ?, ?)
            """,
                (name, path, status, openai_id),
            )
            return cursor.lastrowid

    def get_all_files(self):
        with self.get_cursor() as cursor:
            cursor.execute("SELECT id, name, status, openai_id, created_at FROM files")
            return [
                {
                    "id": row[0],
                    "name": row[1],
                    "status": row[2],
                    "openai_id": row[3],
                    "created_at": row[4],
                }
                for row in cursor.fetchall()
            ]

    def delete_file(self, file_id):
        with self.get_cursor() as cursor:
            cursor.execute("SELECT path FROM files WHERE id = ?", (file_id,))
            result = cursor.fetchone()
            if result:
                file_path = result[0]
                cursor.execute("DELETE FROM files WHERE id = ?", (file_id,))
                if os.path.exists(file_path):
                    os.remove(file_path)
                return True
            return False

    def get_file_path(self, file_id):
        with self.get_cursor() as cursor:
            cursor.execute("SELECT path FROM files WHERE id = ?", (file_id,))
            result = cursor.fetchone()
            return result[0] if result else None

    def get_openai_id(self, file_id):
        with self.get_cursor() as cursor:
            cursor.execute("SELECT openai_id FROM files WHERE id = ?", (file_id,))
            result = cursor.fetchone()
            return result[0] if result else None

    def update_file_status(self, file_id, status):
        with self.get_cursor() as cursor:
            cursor.execute(
                "UPDATE files SET status = ? WHERE id = ?", (status, file_id)
            )

    def update_result_path(self, file_id, result_path):
        with self.get_cursor() as cursor:
            cursor.execute(
                "UPDATE files SET result_path = ? WHERE id = ?", (result_path, file_id)
            )

    def get_file_name(self, file_id):
        with self.get_cursor() as cursor:
            cursor.execute("SELECT name FROM files WHERE id = ?", (file_id,))
            result = cursor.fetchone()
            return result[0] if result else None

    def save_api_token(self, token):
        with self.get_cursor() as cursor:
            cursor.execute(
                """
                INSERT OR REPLACE INTO settings (key, value)
                VALUES ('api_token', ?)
            """,
                (token,),
            )

    def get_api_token(self):
        with self.get_cursor() as cursor:
            cursor.execute("SELECT value FROM settings WHERE key = 'api_token'")
            result = cursor.fetchone()
            return result[0] if result else None

    def set_api_token(self, token):
        with self.get_cursor() as cursor:
            cursor.execute(
                """
                INSERT OR REPLACE INTO settings (key, value)
                VALUES ('api_token', ?)
            """,
                (token,),
            )

    def update_openai_id(self, file_id, openai_id):
        with self.get_cursor() as cursor:
            cursor.execute(
                "UPDATE files SET openai_id = ? WHERE id = ?", (openai_id, file_id)
            )

    def update_batch_id(self, file_id, batch_id):
        with self.get_cursor() as cursor:
            cursor.execute(
                "UPDATE files SET batch_id = ? WHERE id = ?", (batch_id, file_id)
            )

    def get_batch_id(self, file_id):
        with self.get_cursor() as cursor:
            cursor.execute("SELECT batch_id FROM files WHERE id = ?", (file_id,))
            result = cursor.fetchone()
            return result[0] if result else None

    def get_result_path(self, file_id):
        with self.get_cursor() as cursor:
            cursor.execute("SELECT result_path FROM files WHERE id = ?", (file_id,))
            result = cursor.fetchone()
            return result[0] if result else None