103 lines
3.9 KiB
Python
103 lines
3.9 KiB
Python
import sqlite3
|
|
import json
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Optional, Dict, Any, List
|
|
from agent.config import AGENT_HOME
|
|
|
|
STATE_DB_PATH = AGENT_HOME / "agent_state.db"
|
|
|
|
class StateDatabase:
|
|
"""Local SQLite database for agent offline resiliency, session resume tracking and file caching."""
|
|
|
|
def __init__(self, db_path: Path = STATE_DB_PATH):
|
|
self.db_path = db_path
|
|
self._init_db()
|
|
|
|
def _get_conn(self) -> sqlite3.Connection:
|
|
return sqlite3.connect(str(self.db_path))
|
|
|
|
def _init_db(self):
|
|
with self._get_conn() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS upload_sessions (
|
|
filepath TEXT PRIMARY KEY,
|
|
sha256 TEXT NOT NULL,
|
|
session_code TEXT NOT NULL,
|
|
total_chunks INTEGER NOT NULL,
|
|
chunk_size INTEGER NOT NULL,
|
|
completed_chunks INTEGER DEFAULT 0,
|
|
status TEXT NOT NULL,
|
|
last_updated TIMESTAMP NOT NULL
|
|
)
|
|
""")
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS completed_files (
|
|
filepath TEXT PRIMARY KEY,
|
|
sha256 TEXT NOT NULL,
|
|
file_size INTEGER NOT NULL,
|
|
job_id INTEGER,
|
|
last_backup_time TIMESTAMP NOT NULL
|
|
)
|
|
""")
|
|
conn.commit()
|
|
|
|
def save_session(
|
|
self,
|
|
filepath: str,
|
|
sha256: str,
|
|
session_code: str,
|
|
total_chunks: int,
|
|
chunk_size: int,
|
|
status: str = "UPLOADING"
|
|
):
|
|
with self._get_conn() as conn:
|
|
cursor = conn.cursor()
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
cursor.execute("""
|
|
INSERT OR REPLACE INTO upload_sessions
|
|
(filepath, sha256, session_code, total_chunks, chunk_size, status, last_updated)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?)
|
|
""", (filepath, sha256, session_code, total_chunks, chunk_size, status, now))
|
|
conn.commit()
|
|
|
|
def get_session(self, filepath: str) -> Optional[Dict[str, Any]]:
|
|
with self._get_conn() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT filepath, sha256, session_code, total_chunks, chunk_size, completed_chunks, status FROM upload_sessions WHERE filepath = ?", (filepath,))
|
|
row = cursor.fetchone()
|
|
if row:
|
|
return {
|
|
"filepath": row[0],
|
|
"sha256": row[1],
|
|
"session_code": row[2],
|
|
"total_chunks": row[3],
|
|
"chunk_size": row[4],
|
|
"completed_chunks": row[5],
|
|
"status": row[6]
|
|
}
|
|
return None
|
|
|
|
def mark_session_completed(self, filepath: str, sha256: str, file_size: int, job_id: Optional[int] = None):
|
|
with self._get_conn() as conn:
|
|
cursor = conn.cursor()
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
cursor.execute("DELETE FROM upload_sessions WHERE filepath = ?", (filepath,))
|
|
cursor.execute("""
|
|
INSERT OR REPLACE INTO completed_files (filepath, sha256, file_size, job_id, last_backup_time)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
""", (filepath, sha256, file_size, job_id, now))
|
|
conn.commit()
|
|
|
|
def is_file_already_backed_up(self, filepath: str, current_sha256: str) -> bool:
|
|
with self._get_conn() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT sha256 FROM completed_files WHERE filepath = ?", (filepath,))
|
|
row = cursor.fetchone()
|
|
if row and row[0] == current_sha256:
|
|
return True
|
|
return False
|
|
|
|
state_db = StateDatabase()
|