Update storage_handler.py

storage_handler.py

@@ -1,105 +1,215 @@
+# storage_handler.py
 import os
 import base64
-import tempfile
-from typing import Optional, BinaryIO
 import logging
+from typing import Optional, BinaryIO, Dict, Any
 
-from huggingface_hub import HfFileSystem, HfApi, create_bucket
+from huggingface_hub import HfFileSystem, HfApi
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 
 logger = logging.getLogger("StorageHandler")
 
 # ------------------------------------------------------------------------
-# Configuration
+# Configuration – all from environment variables
 # ------------------------------------------------------------------------
-HF_TOKEN = os.environ["HF_TOKEN"]  # required
-BUCKET_NAME = os.environ.get("INFINITY_CHAT_BUCKET", "infinitychat-data")
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("HF_TOKEN environment variable is required")
+
+BUCKET_ID = os.environ.get("INFINITY_CHAT_BUCKET", "infinitychat-data")
 FILE_ENCRYPTION_KEY_B64 = os.environ.get("FILE_ENCRYPTION_KEY", None)
 if FILE_ENCRYPTION_KEY_B64 is None:
-    # Auto‑generate (for production, set env var permanently)
     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!")
 FILE_ENCRYPTION_KEY = base64.urlsafe_b64decode(FILE_ENCRYPTION_KEY_B64.encode())
-assert len(FILE_ENCRYPTION_KEY) == 32, "AES‑256 requires a 32‑byte key"
+assert len(FILE_ENCRYPTION_KEY) == 32, "AES‑256 requires exactly 32 bytes"
 
-# Fetch the bucket owner (username) from the token
+# Get the bucket owner (your HF username)
 _api = HfApi(token=HF_TOKEN)
-OWNER = _api.whoami()["name"]
-del _api
+try:
+    OWNER = _api.whoami()["name"]
+except Exception:
+    # Fallback: use the bucket ID prefix (if it contains a slash)
+    if "/" in BUCKET_ID:
+        OWNER, BUCKET_ID = BUCKET_ID.split("/", 1)
+    else:
+        raise ValueError("Cannot determine bucket owner. Set INFINITY_CHAT_BUCKET as 'owner/bucket-name'")
+
+# Build the proper bucket URI prefix (Xet-native)
+BUCKET_URI_PREFIX = f"hf://buckets/{OWNER}/{BUCKET_ID}"
 
-# Bucket URI prefix
-BUCKET_URI_PREFIX = f"hf://buckets/{OWNER}/{BUCKET_NAME}"
+# ------------------------------------------------------------------------
+# HfFileSystem for low-level bucket operations (Xet-native, NO Git/LFS)
+# ------------------------------------------------------------------------
+_fs = HfFileSystem(token=HF_TOKEN)
 
 # ------------------------------------------------------------------------
-# Bucket lifecycle
+# Bucket initialization (using the modern Hugging Face Hub API)
 # ------------------------------------------------------------------------
 def ensure_bucket():
-    """Create the private bucket if it does not exist (idempotent)."""
+    """
+    Ensure the bucket exists using the native HF Bucket API.
+    The HfApi.create_bucket call does NOT use 'name' – it uses 'bucket_id'.
+    """
+    api = HfApi(token=HF_TOKEN)
+    
+    # Modern API: create_bucket(bucket_id=..., private=True, exist_ok=True)
     try:
-        create_bucket(name=BUCKET_NAME, token=HF_TOKEN, private=True, exist_ok=True)
-        logger.info(f"Bucket '{BUCKET_NAME}' ready (owner: {OWNER}).")
+        api.create_bucket(
+            bucket_id=f"{OWNER}/{BUCKET_ID}",
+            private=True,
+            exist_ok=True
+        )
+        logger.info(f"✅ Bucket '{OWNER}/{BUCKET_ID}' is ready (private).")
     except Exception as e:
-        logger.error(f"Failed to ensure bucket: {e}")
-        raise
+        # Some versions require the full ID
+        if "already exists" in str(e).lower():
+            logger.info(f"Bucket '{OWNER}/{BUCKET_ID}' already exists.")
+        # Try without the owner prefix
+        else:
+            try:
+                api.create_bucket(
+                    bucket_id=BUCKET_ID,
+                    private=True,
+                    exist_ok=True
+                )
+                logger.info(f"✅ Bucket '{BUCKET_ID}' ready.")
+                global BUCKET_URI_PREFIX
+                BUCKET_URI_PREFIX = f"hf://buckets/{BUCKET_ID}"
+            except Exception as e2:
+                if "already exists" in str(e2).lower():
+                    logger.info(f"Bucket '{BUCKET_ID}' already exists.")
+                else:
+                    logger.error(f"Failed to create bucket: {e2}")
+                    raise
 
-# Initialise on import
+# Initialize bucket
 ensure_bucket()
 
 # ------------------------------------------------------------------------
-# Encryption helpers
+# Encryption / Decryption
 # ------------------------------------------------------------------------
 def encrypt_file(data: bytes) -> bytes:
-    """AES‑256‑GCM: nonce (12 B) || ciphertext (incl. tag)"""
+    """Encrypt binary data using AES‑256‑GCM. Output: nonce (12 B) || ciphertext."""
     aesgcm = AESGCM(FILE_ENCRYPTION_KEY)
     nonce = os.urandom(12)
     ciphertext = aesgcm.encrypt(nonce, data, None)
     return nonce + ciphertext
 
 def decrypt_file(encrypted_blob: bytes) -> bytes:
+    """Decrypt data produced by encrypt_file()."""
     nonce = encrypted_blob[:12]
     ciphertext = encrypted_blob[12:]
     aesgcm = AESGCM(FILE_ENCRYPTION_KEY)
     return aesgcm.decrypt(nonce, ciphertext, None)
 
 # ------------------------------------------------------------------------
-# Core file operations via HfFileSystem
+# Core Bucket Operations (Xet-native via HfFileSystem)
 # ------------------------------------------------------------------------
-_fs = HfFileSystem(token=HF_TOKEN)
-
 def _bucket_path(remote_path: str) -> str:
-    """Constructs the full hf://buckets/... URI."""
+    """Construct the full hf://buckets/... URI."""
     return f"{BUCKET_URI_PREFIX}/{remote_path}"
 
 def store_file(remote_path: str, data: bytes) -> str:
-    """Encrypt and write data directly into the bucket. Returns the remote path."""
-    encrypted = encrypt_file(data)
+    """
+    Encrypt data and write it to the bucket using HfFileSystem.
+    Returns: the remote path (for DB storage).
+    """
+    encrypted_data = encrypt_file(data)
     full_uri = _bucket_path(remote_path)
+    
     with _fs.open(full_uri, "wb") as f:
-        f.write(encrypted)
+        f.write(encrypted_data)
+    
+    logger.debug(f"Stored file: {full_uri}")
     return remote_path
 
 def retrieve_file(remote_path: str) -> bytes:
-    """Read from bucket and decrypt."""
+    """
+    Read from bucket using HfFileSystem and decrypt.
+    """
     full_uri = _bucket_path(remote_path)
+    
     with _fs.open(full_uri, "rb") as f:
-        encrypted = f.read()
-    return decrypt_file(encrypted)
+        encrypted_data = f.read()
+    
+    return decrypt_file(encrypted_data)
 
 def delete_file(remote_path: str):
-    """Remove a file from the bucket."""
+    """
+    Delete a file from the bucket using HfFileSystem.
+    """
     full_uri = _bucket_path(remote_path)
     try:
         _fs.rm(full_uri)
+        logger.debug(f"Deleted file: {full_uri}")
     except FileNotFoundError:
-        pass
+        logger.warning(f"File not found for deletion: {full_uri}")
+    except Exception as e:
+        logger.error(f"Failed to delete file {full_uri}: {e}")
+
+def list_files(prefix: str = "") -> list:
+    """
+    List files in the bucket under a given prefix.
+    Returns list of full URIs.
+    """
+    search_path = f"{BUCKET_URI_PREFIX}/{prefix}" if prefix else BUCKET_URI_PREFIX
+    try:
+        return _fs.ls(search_path)
+    except FileNotFoundError:
+        return []
+
+def file_exists(remote_path: str) -> bool:
+    """Check if a file exists in the bucket."""
+    full_uri = _bucket_path(remote_path)
+    return _fs.exists(full_uri)
+
+def get_file_size(remote_path: str) -> Optional[int]:
+    """Get the size of a file in bytes (compressed)."""
+    full_uri = _bucket_path(remote_path)
+    try:
+        info = _fs.info(full_uri)
+        return info.get("size")
+    except Exception:
+        return None
 
 # ------------------------------------------------------------------------
-# Helper for reassembled chunked uploads (used by app.py)
+# Stream-based store (for reassembled chunks)
 # ------------------------------------------------------------------------
 def store_file_stream(remote_path: str, data_stream: BinaryIO) -> str:
     """
     Read the entire stream, encrypt, and write to bucket.
-    Useful for handling reassembled chunks without a single byte array.
     """
     data = data_stream.read()
-    return store_file(remote_path, data)
+    return store_file(remote_path, data)
+
+# ------------------------------------------------------------------------
+# Bucket info (for monitoring / health checks)
+# ------------------------------------------------------------------------
+def get_bucket_info() -> Dict[str, Any]:
+    """Return information about the bucket."""
+    api = HfApi(token=HF_TOKEN)
+    try:
+        # Try the modern API first
+        info = api.bucket_info(bucket_id=f"{OWNER}/{BUCKET_ID}")
+        return {
+            "id": f"{OWNER}/{BUCKET_ID}",
+            "exists": True,
+            "private": info.get("private", True),
+            "total_files": len(list_files()),
+        }
+    except Exception:
+        return {
+            "id": f"{OWNER}/{BUCKET_ID}",
+            "exists": False,
+            "note": "Cannot fetch bucket info"
+        }
+
+# ------------------------------------------------------------------------
+# Clean shutdown
+# ------------------------------------------------------------------------
+def close():
+    """Clean up resources (HfFileSystem)."""
+    if _fs:
+        _fs.close()