Update storage_handler.py

storage_handler.py

@@ -5,6 +5,7 @@ import uuid
 import logging
 from typing import Optional, BinaryIO, Dict, Any, List
 import time
+import threading
 
 from huggingface_hub import HfFileSystem, HfApi
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
@@ -64,7 +65,6 @@ _fs = HfFileSystem(token=HF_TOKEN)
 def ensure_bucket():
     """Create the private bucket if it doesn't exist."""
     global OWNER, BUCKET_URI, BUCKET_ID
-
     try:
         _api.create_bucket(
             bucket_id=f"{OWNER}/{BUCKET_ID}",
@@ -76,9 +76,8 @@ def ensure_bucket():
         error_str = str(e).lower()
         if "already exists" in error_str:
             logger.info(f"📦 Bucket '{OWNER}/{BUCKET_ID}' already exists")
-        elif "http" in error_str and ("403" in error_str or "401" in error_str):
+        elif "403" in error_str or "401" in error_str:
             logger.warning(f"⚠️  Cannot create bucket - permission issue: {e}")
-            logger.info("The bucket may already exist or you don't have permissions")
         else:
             try:
                 _api.create_bucket(
@@ -95,31 +94,29 @@ def ensure_bucket():
                     logger.info(f"📦 Bucket '{BUCKET_ID}' already exists")
                 else:
                     logger.error(f"❌ Failed to create bucket: {e2}")
-                    logger.info("The app will still work but file uploads will fail")
 
 ensure_bucket()
 
-# storage_handler.py – add after ensure_bucket() call at bottom of file
-import threading
-
+# ------------------------------------------------------------------------
+# Heartbeat sync
+# ------------------------------------------------------------------------
 def _sync_bucket_interval():
-    """
-    Periodic sync task: runs every 60 seconds.
-    Uses HfFileSystem which already writes directly to the bucket,
-    so this function logs stats and verifies the connection is alive.
-    """
     def _sync():
         while True:
             time.sleep(60)
             try:
                 stats = get_storage_stats()
-                logger.info(f"📦 Bucket sync heartbeat — {stats['file_count']} files / {stats['total_size_mb']} MB")
+                logger.info(
+                    f"📦 Bucket heartbeat — "
+                    f"{stats['file_count']} files / {stats['total_size_mb']} MB"
+                )
             except Exception as e:
-                logger.warning(f"⚠️  Bucket sync heartbeat failed: {e}")
+                logger.warning(f"⚠️  Bucket heartbeat failed: {e}")
     t = threading.Thread(target=_sync, daemon=True)
     t.start()
 
 _sync_bucket_interval()
+
 # ------------------------------------------------------------------------
 # Encryption / Decryption
 # ------------------------------------------------------------------------
@@ -141,9 +138,6 @@ def verify_file_integrity(encrypted_blob: bytes) -> bool:
 # ------------------------------------------------------------------------
 # Path Utilities
 # ------------------------------------------------------------------------
-def _build_path(*parts: str) -> str:
-    return "/".join(p.strip("/") for p in parts if p)
-
 def _bucket_path(remote_path: str) -> str:
     return f"{BUCKET_URI}/{remote_path.lstrip('/')}"
 
@@ -163,19 +157,16 @@ def _validate_path(remote_path: str) -> None:
 def store_file(remote_path: str, data: bytes, encrypt: bool = True) -> str:
     _validate_path(remote_path)
     try:
-        if encrypt:
-            encrypted_data = encrypt_bytes(data, remote_path.encode('utf-8'))
-        else:
-            encrypted_data = data
+        payload = encrypt_bytes(data, remote_path.encode('utf-8')) if encrypt else data
         full_uri = _bucket_path(remote_path)
         with _fs.open(full_uri, "wb") as f:
-            f.write(encrypted_data)
+            f.write(payload)
         if not _fs.exists(full_uri):
             raise IOError(f"Failed to verify file was stored: {full_uri}")
-        logger.debug(f"💾 Stored file: {remote_path} ({len(data)} bytes)")
+        logger.debug(f"💾 Stored: {remote_path} ({len(data)} bytes)")
         return remote_path
     except Exception as e:
-        logger.error(f"❌ Failed to store file {remote_path}: {e}")
+        logger.error(f"❌ Failed to store {remote_path}: {e}")
         raise IOError(f"Storage failed: {e}")
 
 def retrieve_file(remote_path: str, decrypt: bool = True) -> bytes:
@@ -185,21 +176,20 @@ def retrieve_file(remote_path: str, decrypt: bool = True) -> bytes:
         if not _fs.exists(full_uri):
             raise FileNotFoundError(f"File not found: {remote_path}")
         with _fs.open(full_uri, "rb") as f:
-            encrypted_data = f.read()
-        if not encrypted_data:
+            payload = f.read()
+        if not payload:
             raise IOError(f"Empty file: {remote_path}")
         if decrypt:
-            if not verify_file_integrity(encrypted_data):
+            if not verify_file_integrity(payload):
                 raise IOError(f"Corrupted file: {remote_path}")
-            decrypted_data = decrypt_bytes(encrypted_data, remote_path.encode('utf-8'))
-            logger.debug(f"📂 Retrieved file: {remote_path} ({len(decrypted_data)} bytes)")
-            return decrypted_data
-        else:
-            return encrypted_data
+            data = decrypt_bytes(payload, remote_path.encode('utf-8'))
+            logger.debug(f"📂 Retrieved: {remote_path} ({len(data)} bytes)")
+            return data
+        return payload
     except FileNotFoundError:
         raise
     except Exception as e:
-        logger.error(f"❌ Failed to retrieve file {remote_path}: {e}")
+        logger.error(f"❌ Failed to retrieve {remote_path}: {e}")
         raise IOError(f"Retrieval failed: {e}")
 
 def delete_file(remote_path: str) -> bool:
@@ -208,26 +198,25 @@ def delete_file(remote_path: str) -> bool:
     try:
         if _fs.exists(full_uri):
             _fs.rm(full_uri)
-            logger.debug(f"🗑️  Deleted file: {remote_path}")
+            logger.debug(f"🗑️  Deleted: {remote_path}")
             return True
-        else:
-            logger.warning(f"⚠️  File not found for deletion: {remote_path}")
-            return False
+        logger.warning(f"⚠️  Not found for deletion: {remote_path}")
+        return False
     except Exception as e:
-        logger.error(f"❌ Failed to delete file {remote_path}: {e}")
+        logger.error(f"❌ Failed to delete {remote_path}: {e}")
         raise IOError(f"Deletion failed: {e}")
 
 def file_exists(remote_path: str) -> bool:
     _validate_path(remote_path)
-    return _fs.exists(_bucket_path(remote_path))
+    try:
+        return _fs.exists(_bucket_path(remote_path))
+    except Exception:
+        return False
 
 def list_files(prefix: str = "", recursive: bool = True) -> List[str]:
     search_path = _bucket_path(prefix) if prefix else BUCKET_URI
     try:
-        if recursive:
-            items = _fs.ls(search_path, detail=False, recursive=True)
-        else:
-            items = _fs.ls(search_path, detail=False)
+        items = _fs.ls(search_path, detail=False, recursive=recursive)
         prefix_len = len(BUCKET_URI) + 1
         return [item[prefix_len:] for item in items if not _fs.isdir(item)]
     except FileNotFoundError:
@@ -241,10 +230,7 @@ def get_file_size(remote_path: str) -> Optional[int]:
     try:
         info = _fs.info(_bucket_path(remote_path))
         return info.get("size")
-    except FileNotFoundError:
-        return None
-    except Exception as e:
-        logger.error(f"Failed to get file size: {e}")
+    except Exception:
         return None
 
 def get_file_info(remote_path: str) -> Optional[Dict[str, Any]]:
@@ -265,13 +251,11 @@ def get_file_info(remote_path: str) -> Optional[Dict[str, Any]]:
         return None
 
 def store_file_stream(remote_path: str, data_stream: BinaryIO, encrypt: bool = True) -> str:
-    data = data_stream.read()
-    return store_file(remote_path, data, encrypt=encrypt)
+    return store_file(remote_path, data_stream.read(), encrypt=encrypt)
 
 def retrieve_file_stream(remote_path: str, decrypt: bool = True) -> BinaryIO:
-    data = retrieve_file(remote_path, decrypt=decrypt)
     import io
-    return io.BytesIO(data)
+    return io.BytesIO(retrieve_file(remote_path, decrypt=decrypt))
 
 def get_storage_stats() -> Dict[str, Any]:
     try:
@@ -307,8 +291,7 @@ def create_backup(backup_prefix: str = "backups") -> str:
         for file_path in files:
             try:
                 data = retrieve_file(file_path)
-                backup_file_path = f"{backup_path}/{file_path}"
-                store_file(backup_file_path, data, encrypt=True)
+                store_file(f"{backup_path}/{file_path}", data, encrypt=True)
             except Exception as e:
                 logger.error(f"Failed to backup {file_path}: {e}")
         logger.info(f"💾 Backup created at: {backup_path}")
@@ -317,16 +300,89 @@ def create_backup(backup_prefix: str = "backups") -> str:
         logger.error(f"❌ Backup failed: {e}")
         raise
 
-def close():
-    if _fs:
+# ------------------------------------------------------------------------
+# Database Sync to/from Bucket
+# ------------------------------------------------------------------------
+DB_BUCKET_PATH = "database/infinitychat.db"
+_db_sync_lock = threading.Lock()
+_last_db_sync = 0
+DB_SYNC_INTERVAL = 30  # seconds
+
+def download_database(local_path: str) -> bool:
+    """
+    Download database from bucket to local path.
+    Returns True if database was found and downloaded.
+    """
+    try:
+        if file_exists(DB_BUCKET_PATH):
+            logger.info("📥 Downloading database from bucket...")
+            data = retrieve_file(DB_BUCKET_PATH)
+            os.makedirs(os.path.dirname(local_path), exist_ok=True)
+            with open(local_path, 'wb') as f:
+                f.write(data)
+            logger.info(f"✅ Database downloaded ({len(data):,} bytes)")
+            return True
+        else:
+            logger.info("📭 No existing database in bucket - will create fresh")
+            return False
+    except Exception as e:
+        logger.error(f"❌ Failed to download database: {e}")
+        return False
+
+def upload_database(local_path: str) -> bool:
+    """
+    Upload local database file to bucket.
+    Returns True on success.
+    """
+    global _last_db_sync
+    with _db_sync_lock:
         try:
-            _fs.close()
-        except:
-            pass
+            if not os.path.exists(local_path):
+                logger.warning(f"⚠️  Database not found at {local_path}, skipping upload")
+                return False
+            with open(local_path, 'rb') as f:
+                data = f.read()
+            if not data:
+                logger.warning("⚠️  Database file is empty, skipping upload")
+                return False
+            store_file(DB_BUCKET_PATH, data)
+            _last_db_sync = time.time()
+            logger.debug(f"☁️  Database synced to bucket ({len(data):,} bytes)")
+            return True
+        except Exception as e:
+            logger.error(f"❌ Failed to upload database: {e}")
+            return False
+
+def start_db_sync(local_path: str):
+    """
+    Start background thread that periodically uploads the database to bucket.
+    """
+    def _sync_loop():
+        while True:
+            time.sleep(DB_SYNC_INTERVAL)
+            try:
+                upload_database(local_path)
+            except Exception as e:
+                logger.warning(f"⚠️  Periodic DB sync failed: {e}")
+
+    t = threading.Thread(target=_sync_loop, daemon=True)
+    t.start()
+    logger.info(f"🔄 Database auto-sync started (every {DB_SYNC_INTERVAL}s)")
+
+def close():
+    try:
+        _fs.close()
+    except Exception:
+        pass
 
+# ------------------------------------------------------------------------
+# Public API
+# ------------------------------------------------------------------------
 __all__ = [
     'store_file', 'retrieve_file', 'delete_file', 'file_exists',
     'list_files', 'get_file_size', 'get_file_info',
     'store_file_stream', 'retrieve_file_stream',
-    'get_storage_stats', 'create_backup', 'close'
+    'get_storage_stats', 'create_backup', 'close',
+    'download_database', 'upload_database', 'start_db_sync',
+    'DB_BUCKET_PATH'
 ]