Update storage_handler.py

storage_handler.py

@@ -1,7 +1,7 @@
 import os
 import base64
-import hashlib
-import uuid
+import sqlite3
+import tempfile
 import logging
 from typing import Optional, BinaryIO, Dict, Any, List
 import time
@@ -16,17 +16,25 @@ logger = logging.getLogger("InfinityChat.Storage")
 # Configuration
 # ------------------------------------------------------------------------
 HF_TOKEN = os.environ.get("HF_TOKEN")
-if not HF_TOKEN:
-    raise ValueError("❌ HF_TOKEN environment variable is required")
-
 BUCKET_NAME = os.environ.get("INFINITY_CHAT_BUCKET", "infinitychat-data")
 FILE_ENCRYPTION_KEY_B64 = os.environ.get("FILE_ENCRYPTION_KEY", None)
 
+# Offline mode: lets the app run without a Hugging Face token (dev/tests).
+# File storage endpoints will return a clear error instead of crashing boot.
+_OFFLINE = not HF_TOKEN or HF_TOKEN.strip() in ("", "none", "offline", "0")
+
 if FILE_ENCRYPTION_KEY_B64 is None:
-    from cryptography.fernet import Fernet
-    FILE_ENCRYPTION_KEY_B64 = Fernet.generate_key().decode()
-    logger.warning("⚠️  FILE_ENCRYPTION_KEY not set - generated random key. Set it permanently!")
-    print(f"Generated FILE_ENCRYPTION_KEY: {FILE_ENCRYPTION_KEY_B64}")
+    if _OFFLINE:
+        # Deterministic dev-only key so local databases stay readable between runs.
+        FILE_ENCRYPTION_KEY_B64 = base64.urlsafe_b64encode(b"0" * 32).decode()
+        logger.warning("⚠️  FILE_ENCRYPTION_KEY not set and running offline - using dev-only key")
+    else:
+        from cryptography.fernet import Fernet
+        FILE_ENCRYPTION_KEY_B64 = Fernet.generate_key().decode()
+        logger.warning(
+            "⚠️  FILE_ENCRYPTION_KEY not set - generated a random key for this run. "
+            "Files uploaded now will NOT be readable after restart unless you persist it!"
+        )
 
 try:
     FILE_ENCRYPTION_KEY = base64.urlsafe_b64decode(FILE_ENCRYPTION_KEY_B64.encode())
@@ -36,37 +44,64 @@ except Exception as e:
 assert len(FILE_ENCRYPTION_KEY) == 32, "FILE_ENCRYPTION_KEY must decode to exactly 32 bytes for AES-256"
 
 # ------------------------------------------------------------------------
-# Get bucket owner (HF username)
+# Storage availability
 # ------------------------------------------------------------------------
-_api = HfApi(token=HF_TOKEN)
-try:
-    owner_info = _api.whoami()
-    OWNER = owner_info["name"]
-except Exception as e:
-    logger.warning(f"Could not determine HF username: {e}")
-    OWNER = "infinitychat"
+class StorageUnavailableError(IOError):
+    """Raised when a storage operation is attempted without a valid HF_TOKEN."""
 
-if "/" in BUCKET_NAME:
-    OWNER, BUCKET_ID = BUCKET_NAME.split("/", 1)
-else:
-    BUCKET_ID = BUCKET_NAME
 
-BUCKET_URI = f"hf://buckets/{OWNER}/{BUCKET_ID}"
-logger.info(f"📦 Bucket URI: {BUCKET_URI}")
+def _require_available():
+    if _OFFLINE:
+        raise StorageUnavailableError(
+            "File storage is unavailable: HF_TOKEN is not configured."
+        )
+
+
+def _get_api():
+    _require_available()
+    return HfApi(token=HF_TOKEN)
+
+
+def _get_fs():
+    _require_available()
+    return HfFileSystem(token=HF_TOKEN)
+
 
 # ------------------------------------------------------------------------
-# HfFileSystem instance
+# Get bucket owner (HF username)
 # ------------------------------------------------------------------------
-_fs = HfFileSystem(token=HF_TOKEN)
+OWNER = None
+BUCKET_ID = BUCKET_NAME
+BUCKET_URI = None
+
+if not _OFFLINE:
+    try:
+        owner_info = _get_api().whoami()
+        OWNER = owner_info["name"]
+    except Exception as e:
+        logger.warning(f"Could not determine HF username: {e}")
+        OWNER = None
+
+    if OWNER:
+        if "/" in BUCKET_NAME:
+            OWNER, BUCKET_ID = BUCKET_NAME.split("/", 1)
+        else:
+            BUCKET_ID = BUCKET_NAME
+        BUCKET_URI = f"hf://buckets/{OWNER}/{BUCKET_ID}"
+        logger.info(f"📦 Bucket URI: {BUCKET_URI}")
+
 
 # ------------------------------------------------------------------------
 # Bucket Initialization
 # ------------------------------------------------------------------------
 def ensure_bucket():
     """Create the private bucket if it doesn't exist."""
+    if _OFFLINE:
+        return
     global OWNER, BUCKET_URI, BUCKET_ID
+    api = _get_api()
     try:
-        _api.create_bucket(
+        api.create_bucket(
             bucket_id=f"{OWNER}/{BUCKET_ID}",
             private=True,
             exist_ok=True
@@ -80,7 +115,7 @@ def ensure_bucket():
             logger.warning(f"⚠️  Cannot create bucket - permission issue: {e}")
         else:
             try:
-                _api.create_bucket(
+                api.create_bucket(
                     bucket_id=BUCKET_ID,
                     private=True,
                     exist_ok=True
@@ -95,12 +130,17 @@ def ensure_bucket():
                 else:
                     logger.error(f"❌ Failed to create bucket: {e2}")
 
-ensure_bucket()
+
+if not _OFFLINE:
+    ensure_bucket()
 
 # ------------------------------------------------------------------------
 # Heartbeat sync
 # ------------------------------------------------------------------------
 def _sync_bucket_interval():
+    if _OFFLINE:
+        return
+
     def _sync():
         while True:
             time.sleep(60)
@@ -115,6 +155,7 @@ def _sync_bucket_interval():
     t = threading.Thread(target=_sync, daemon=True)
     t.start()
 
+
 _sync_bucket_interval()
 
 # ------------------------------------------------------------------------
@@ -126,12 +167,14 @@ def encrypt_bytes(data: bytes, aad: Optional[bytes] = None) -> bytes:
     ciphertext = aesgcm.encrypt(nonce, data, aad or b"")
     return nonce + ciphertext
 
+
 def decrypt_bytes(encrypted_blob: bytes, aad: Optional[bytes] = None) -> bytes:
     aesgcm = AESGCM(FILE_ENCRYPTION_KEY)
     nonce = encrypted_blob[:12]
     ciphertext = encrypted_blob[12:]
     return aesgcm.decrypt(nonce, ciphertext, aad or b"")
 
+
 def verify_file_integrity(encrypted_blob: bytes) -> bool:
     return len(encrypted_blob) >= 29
 
@@ -139,8 +182,12 @@ def verify_file_integrity(encrypted_blob: bytes) -> bool:
 # Path Utilities
 # ------------------------------------------------------------------------
 def _bucket_path(remote_path: str) -> str:
+    _require_available()
+    if not BUCKET_URI:
+        raise StorageUnavailableError("Bucket is not configured (offline mode).")
     return f"{BUCKET_URI}/{remote_path.lstrip('/')}"
 
+
 def _validate_path(remote_path: str) -> None:
     if not remote_path:
         raise ValueError("Path cannot be empty")
@@ -156,12 +203,14 @@ def _validate_path(remote_path: str) -> None:
 # ------------------------------------------------------------------------
 def store_file(remote_path: str, data: bytes, encrypt: bool = True) -> str:
     _validate_path(remote_path)
+    _require_available()
+    fs = _get_fs()
     try:
         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:
+        with fs.open(full_uri, "wb") as f:
             f.write(payload)
-        if not _fs.exists(full_uri):
+        if not fs.exists(full_uri):
             raise IOError(f"Failed to verify file was stored: {full_uri}")
         logger.debug(f"💾 Stored: {remote_path} ({len(data)} bytes)")
         return remote_path
@@ -169,13 +218,16 @@ def store_file(remote_path: str, data: bytes, encrypt: bool = True) -> str:
         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:
     _validate_path(remote_path)
+    _require_available()
+    fs = _get_fs()
     full_uri = _bucket_path(remote_path)
     try:
-        if not _fs.exists(full_uri):
+        if not fs.exists(full_uri):
             raise FileNotFoundError(f"File not found: {remote_path}")
-        with _fs.open(full_uri, "rb") as f:
+        with fs.open(full_uri, "rb") as f:
             payload = f.read()
         if not payload:
             raise IOError(f"Empty file: {remote_path}")
@@ -192,12 +244,16 @@ def retrieve_file(remote_path: str, decrypt: bool = True) -> bytes:
         logger.error(f"❌ Failed to retrieve {remote_path}: {e}")
         raise IOError(f"Retrieval failed: {e}")
 
+
 def delete_file(remote_path: str) -> bool:
     _validate_path(remote_path)
+    if _OFFLINE:
+        return False
+    fs = _get_fs()
     full_uri = _bucket_path(remote_path)
     try:
-        if _fs.exists(full_uri):
-            _fs.rm(full_uri)
+        if fs.exists(full_uri):
+            fs.rm(full_uri)
             logger.debug(f"🗑️  Deleted: {remote_path}")
             return True
         logger.warning(f"⚠️  Not found for deletion: {remote_path}")
@@ -206,37 +262,51 @@ def delete_file(remote_path: str) -> bool:
         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)
+    if _OFFLINE:
+        return False
     try:
-        return _fs.exists(_bucket_path(remote_path))
+        _validate_path(remote_path)
+        return _get_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
+    if _OFFLINE:
+        return []
     try:
-        items = _fs.ls(search_path, detail=False, recursive=recursive)
+        fs = _get_fs()
+        search_path = _bucket_path(prefix) if prefix else BUCKET_URI
+        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)]
+        return [item[prefix_len:] for item in items if not fs.isdir(item)]
     except FileNotFoundError:
         return []
     except Exception as e:
         logger.error(f"Failed to list files: {e}")
         return []
 
+
 def get_file_size(remote_path: str) -> Optional[int]:
-    _validate_path(remote_path)
+    if _OFFLINE:
+        return None
     try:
-        info = _fs.info(_bucket_path(remote_path))
+        _validate_path(remote_path)
+        info = _get_fs().info(_bucket_path(remote_path))
         return info.get("size")
     except Exception:
         return None
 
+
 def get_file_info(remote_path: str) -> Optional[Dict[str, Any]]:
-    _validate_path(remote_path)
+    if _OFFLINE:
+        return None
     try:
-        info = _fs.info(_bucket_path(remote_path))
+        _validate_path(remote_path)
+        fs = _get_fs()
+        info = fs.info(_bucket_path(remote_path))
         return {
             "name": remote_path,
             "size": info.get("size"),
@@ -247,25 +317,33 @@ def get_file_info(remote_path: str) -> Optional[Dict[str, Any]]:
     except FileNotFoundError:
         return None
     except Exception as e:
-        logger.error(f"Failed to get file info: {e}")
+        logger.error(f"Failed to get file info: {remote_path}: {e}")
         return None
 
+
 def store_file_stream(remote_path: str, data_stream: BinaryIO, encrypt: bool = True) -> str:
     return store_file(remote_path, data_stream.read(), encrypt=encrypt)
 
+
 def retrieve_file_stream(remote_path: str, decrypt: bool = True) -> BinaryIO:
     import io
     return io.BytesIO(retrieve_file(remote_path, decrypt=decrypt))
 
+
 def get_storage_stats() -> Dict[str, Any]:
+    if _OFFLINE:
+        return {
+            "bucket": None, "file_count": 0, "total_size": 0,
+            "total_size_mb": 0, "error": "Storage offline (HF_TOKEN not configured)"
+        }
     try:
-        files = list_files()
+        fs = _get_fs()
+        items = fs.ls(BUCKET_URI, detail=True, recursive=True)
         total_size = 0
         file_count = 0
-        for f in files:
-            size = get_file_size(f)
-            if size is not None:
-                total_size += size
+        for item in items:
+            if item.get("type") == "file":
+                total_size += item.get("size") or 0
                 file_count += 1
         return {
             "bucket": f"{OWNER}/{BUCKET_ID}",
@@ -283,7 +361,9 @@ def get_storage_stats() -> Dict[str, Any]:
             "error": str(e)
         }
 
+
 def create_backup(backup_prefix: str = "backups") -> str:
+    _require_available()
     timestamp = int(time.time())
     backup_path = f"{backup_prefix}/backup_{timestamp}"
     try:
@@ -308,11 +388,15 @@ _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.
     """
+    if _OFFLINE:
+        logger.info("📭 Storage offline - skipping database download")
+        return False
     try:
         if file_exists(DB_BUCKET_PATH):
             logger.info("📥 Downloading database from bucket...")
@@ -329,21 +413,44 @@ def download_database(local_path: str) -> bool:
         logger.error(f"❌ Failed to download database: {e}")
         return False
 
+
 def upload_database(local_path: str) -> bool:
     """
-    Upload local database file to bucket.
+    Upload a consistent snapshot of the local database to the bucket.
+
+    The app runs SQLite in WAL mode, so the main .db file can lag behind
+    committed transactions (they live in the -wal file). Copying the raw file
+    could persist stale or torn data - instead we take a proper snapshot with
+    the sqlite backup API, which is safe even while writers are active.
     Returns True on success.
     """
     global _last_db_sync
+    if _OFFLINE:
+        return False
     with _db_sync_lock:
+        tmp_path = None
         try:
             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:
+
+            # snapshot to a temp file via the online backup API
+            fd, tmp_path = tempfile.mkstemp(suffix=".db",
+                                            dir=os.path.dirname(os.path.abspath(local_path)) or ".")
+            os.close(fd)
+            os.unlink(tmp_path)  # backup() needs the destination to not exist
+            src = sqlite3.connect(f"file:{local_path}?mode=ro", uri=True)
+            dst = sqlite3.connect(tmp_path)
+            try:
+                src.backup(dst)
+            finally:
+                src.close()
+                dst.close()
+
+            with open(tmp_path, 'rb') as f:
                 data = f.read()
             if not data:
-                logger.warning("⚠️  Database file is empty, skipping upload")
+                logger.warning("⚠️  Database snapshot is empty, skipping upload")
                 return False
             store_file(DB_BUCKET_PATH, data)
             _last_db_sync = time.time()
@@ -352,11 +459,22 @@ def upload_database(local_path: str) -> bool:
         except Exception as e:
             logger.error(f"❌ Failed to upload database: {e}")
             return False
+        finally:
+            if tmp_path and os.path.exists(tmp_path):
+                try:
+                    os.unlink(tmp_path)
+                except Exception:
+                    pass
+
 
 def start_db_sync(local_path: str):
     """
     Start background thread that periodically uploads the database to bucket.
     """
+    if _OFFLINE:
+        logger.info("🔄 Database auto-sync skipped (storage offline)")
+        return
+
     def _sync_loop():
         while True:
             time.sleep(DB_SYNC_INTERVAL)
@@ -369,9 +487,11 @@ def start_db_sync(local_path: str):
     t.start()
     logger.info(f"🔄 Database auto-sync started (every {DB_SYNC_INTERVAL}s)")
 
+
 def close():
     try:
-        _fs.close()
+        if not _OFFLINE:
+            _get_fs().close()
     except Exception:
         pass
 
@@ -384,5 +504,5 @@ __all__ = [
     'store_file_stream', 'retrieve_file_stream',
     'get_storage_stats', 'create_backup', 'close',
     'download_database', 'upload_database', 'start_db_sync',
-    'DB_BUCKET_PATH'
+    'DB_BUCKET_PATH', 'StorageUnavailableError'
 ]