| import sqlite3
|
| from dotenv import load_dotenv
|
| import os
|
|
|
| class Database:
|
| def __init__(self):
|
| load_dotenv()
|
| base_dir = os.getenv("HF_HOME", "/home/user/app")
|
| data_dir = os.path.join(base_dir, "data")
|
| os.makedirs(data_dir, exist_ok=True)
|
| self.db_path = os.path.join(data_dir, "database.sqlite")
|
| os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
|
| self.conn = sqlite3.connect(self.db_path, check_same_thread=False)
|
| self.create_tables()
|
|
|
| def create_tables(self):
|
| with self.conn:
|
| self.conn.execute("""
|
| CREATE TABLE IF NOT EXISTS documents (
|
| id INTEGER PRIMARY KEY AUTOINCREMENT,
|
| filename TEXT NOT NULL,
|
| path TEXT NOT NULL,
|
| status VARCHAR(20) NOT NULL
|
| CHECK(status IN ('uploaded', 'ingested')),
|
| timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
|
| )
|
| """)
|
|
|
| def add_document(self, filename, path, status="uploaded"):
|
| with self.conn:
|
| self.conn.execute("""
|
| INSERT INTO documents (filename, path, status) VALUES (?, ?, ?)
|
| """, (filename, path, status))
|
|
|
| def update_document_status(self, path, status):
|
| with self.conn:
|
| self.conn.execute("""
|
| UPDATE documents SET status = ? WHERE path = ?
|
| """, (status, path))
|
|
|
| def list_documents(self):
|
| with self.conn:
|
| cursor = self.conn.execute("""
|
| SELECT filename, status, timestamp, path FROM documents ORDER BY timestamp DESC
|
| """)
|
| return cursor.fetchall()
|
|
|
| def delete_document(self, path):
|
| with self.conn:
|
| self.conn.execute("""
|
| DELETE FROM documents WHERE path = ?
|
| """, (path,))
|
|
|
| def disconnect(self):
|
| self.conn.close()
|
|
|
|
|
| if __name__ == "__main__":
|
| db = Database()
|
|
|
|
|
|
|
|
|
|
|
| db.delete_document("sample.pdf")
|
| print(db.list_documents()) |
