db.py

import sqlite3
from env_vars import DB_PATH

# Initialize SQLite database for file tracking
def init_db(DB_PATH = DB_PATH):
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute('''CREATE TABLE IF NOT EXISTS files
                 (file_id TEXT PRIMARY KEY, 
                  file_name TEXT,
                  file_type TEXT,
                  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)''')
    conn.commit()
    conn.close()

def add_file_to_db(file_id, file_name, file_type,DB_PATH = DB_PATH):
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute('INSERT INTO files (file_id, file_name, file_type) VALUES (?, ?, ?)',
              (file_id, file_name, file_type))
    conn.commit()
    conn.close()

def remove_file_from_db(file_id,DB_PATH = DB_PATH):
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute('DELETE FROM files WHERE file_id = ?', (file_id,))
    conn.commit()
    conn.close()

def get_all_files(DB_PATH = DB_PATH):
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute('SELECT file_id, file_name, file_type, created_at FROM files ORDER BY created_at DESC')
    files = [{'file_id': row[0], 
              'file_name': row[1], 
              'file_type': row[2],
              'created_at': row[3]} for row in c.fetchall()]
    conn.close()
    return files