Update sync_storage.py

sync_storage.py

@@ -2,83 +2,231 @@
 import os
 import shutil
 import json
+import hashlib
 from pathlib import Path
-from huggingface_hub import HfApi, create_repo
+from datetime import datetime
+from huggingface_hub import HfApi, create_repo, list_repo_files
 import tarfile
 import tempfile
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
 class HFStorageSync:
-    def __init__(self, repo_id, token=None, data_dir="/tmp/open-webui-data"):
+    def __init__(self, repo_id, token=None, data_dir="/tmp/open-webui-data", 
+                 max_backups=3, compression_level=6):
         self.repo_id = repo_id
         self.data_dir = Path(data_dir)
         self.token = token
+        self.max_backups = max_backups
+        self.compression_level = compression_level
         
         # Initialize API with token directly
         self.api = HfApi(token=token) if token else HfApi()
+        
+        # File patterns for better organization
+        self.archive_pattern = "data-{timestamp}.tar.gz"
+        self.latest_link = "data-latest.tar.gz"
+        self.metadata_file = "storage-metadata.json"
+    
+    def _get_directory_hash(self):
+        """Calculate hash of directory contents for change detection"""
+        hasher = hashlib.sha256()
+        
+        if not self.data_dir.exists():
+            return hasher.hexdigest()
+        
+        for item in sorted(self.data_dir.rglob('*')):
+            if item.is_file() and item.name not in [".gitkeep", "test_write"]:
+                hasher.update(str(item.relative_to(self.data_dir)).encode())
+                hasher.update(str(item.stat().st_mtime).encode())
+                hasher.update(str(item.stat().st_size).encode())
+        
+        return hasher.hexdigest()
     
+    def _get_archive_size(self, archive_path):
+        """Get the size of an archive file"""
+        try:
+            return os.path.getsize(archive_path)
+        except:
+            return 0
+    
+    def _format_size(self, size_bytes):
+        """Format file size in human readable format"""
+        for unit in ['B', 'KB', 'MB', 'GB']:
+            if size_bytes < 1024.0:
+                return f"{size_bytes:.1f} {unit}"
+            size_bytes /= 1024.0
+        return f"{size_bytes:.1f} TB"
+
     def ensure_repo_exists(self):
         """Create repository if it doesn't exist"""
         if not self.token:
-            print("No token provided, cannot create repository")
+            logger.warning("No token provided, cannot create repository")
             return False
             
         try:
             # Check if repo exists
             repo_info = self.api.repo_info(repo_id=self.repo_id, repo_type="dataset")
-            print(f"Repository {self.repo_id} exists")
+            logger.info(f"Repository {self.repo_id} exists")
             return True
         except Exception as e:
-            print(f"Repository {self.repo_id} not found, attempting to create...")
+            logger.info(f"Repository {self.repo_id} not found, attempting to create...")
             try:
                 create_repo(
                     repo_id=self.repo_id,
                     repo_type="dataset",
                     token=self.token,
-                    private=True,  # Make it private by default
+                    private=True,
                     exist_ok=True
                 )
-                print(f"Created repository {self.repo_id}")
+                logger.info(f"Created repository {self.repo_id}")
                 
-                # Create initial README
-                readme_content = """# Open WebUI Storage
+                # Create initial README and metadata
+                self._create_initial_files()
+                return True
+            except Exception as create_error:
+                logger.error(f"Failed to create repository: {create_error}")
+                return False
+    
+    def _create_initial_files(self):
+        """Create initial repository files"""
+        readme_content = """# Open WebUI Storage
 
-This dataset stores persistent data for Open WebUI deployment.
+This dataset stores persistent data for Open WebUI deployment with automatic cleanup and versioning.
 
 ## Contents
 
-- `data.tar.gz`: Compressed archive containing all Open WebUI data including:
-  - User configurations
-  - Chat histories
-  - Uploaded files
-  - Database files
+- `data-latest.tar.gz`: Latest data archive (symlink)
+- `data-YYYYMMDD-HHMMSS.tar.gz`: Timestamped data archives
+- `storage-metadata.json`: Metadata about stored archives
+- `README.md`: This file
+
+## Features
+
+- Automatic cleanup of old backups
+- Change detection to avoid unnecessary uploads
+- Compression optimization
+- Storage usage monitoring
 
 This repository is automatically managed by the Open WebUI sync system.
 """
+        
+        metadata = {
+            "created": datetime.utcnow().isoformat(),
+            "max_backups": self.max_backups,
+            "archives": [],
+            "total_size": 0
+        }
+        
+        # Upload README
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False) as tmp:
+            tmp.write(readme_content)
+            tmp.flush()
+            
+            self.api.upload_file(
+                path_or_fileobj=tmp.name,
+                path_in_repo="README.md",
+                repo_id=self.repo_id,
+                repo_type="dataset",
+                commit_message="Initial repository setup",
+                token=self.token
+            )
+            os.unlink(tmp.name)
+        
+        # Upload metadata
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as tmp:
+            json.dump(metadata, tmp, indent=2)
+            tmp.flush()
+            
+            self.api.upload_file(
+                path_or_fileobj=tmp.name,
+                path_in_repo=self.metadata_file,
+                repo_id=self.repo_id,
+                repo_type="dataset",
+                commit_message="Initial metadata",
+                token=self.token
+            )
+            os.unlink(tmp.name)
+    
+    def _get_metadata(self):
+        """Download and parse metadata"""
+        try:
+            file_path = self.api.hf_hub_download(
+                repo_id=self.repo_id,
+                filename=self.metadata_file,
+                repo_type="dataset",
+                token=self.token
+            )
+            
+            with open(file_path, 'r') as f:
+                return json.load(f)
+        except Exception as e:
+            logger.warning(f"Could not load metadata: {e}")
+            return {
+                "created": datetime.utcnow().isoformat(),
+                "max_backups": self.max_backups,
+                "archives": [],
+                "total_size": 0
+            }
+    
+    def _update_metadata(self, metadata):
+        """Upload updated metadata"""
+        try:
+            with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as tmp:
+                json.dump(metadata, tmp, indent=2)
+                tmp.flush()
                 
-                with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False) as tmp:
-                    tmp.write(readme_content)
-                    tmp.flush()
-                    
-                    self.api.upload_file(
-                        path_or_fileobj=tmp.name,
-                        path_in_repo="README.md",
-                        repo_id=self.repo_id,
-                        repo_type="dataset",
-                        commit_message="Initial repository setup",
-                        token=self.token
-                    )
-                    
-                    os.unlink(tmp.name)
-                
-                return True
-            except Exception as create_error:
-                print(f"Failed to create repository: {create_error}")
-                return False
+                self.api.upload_file(
+                    path_or_fileobj=tmp.name,
+                    path_in_repo=self.metadata_file,
+                    repo_id=self.repo_id,
+                    repo_type="dataset",
+                    commit_message="Update metadata",
+                    token=self.token
+                )
+                os.unlink(tmp.name)
+        except Exception as e:
+            logger.error(f"Failed to update metadata: {e}")
     
+    def _cleanup_old_archives(self, metadata):
+        """Remove old archives beyond max_backups limit"""
+        if len(metadata["archives"]) <= self.max_backups:
+            return metadata
+        
+        # Sort by timestamp, keep newest
+        archives = sorted(metadata["archives"], key=lambda x: x["timestamp"], reverse=True)
+        to_keep = archives[:self.max_backups]
+        to_delete = archives[self.max_backups:]
+        
+        total_deleted_size = 0
+        for archive in to_delete:
+            try:
+                self.api.delete_file(
+                    path_in_repo=archive["filename"],
+                    repo_id=self.repo_id,
+                    repo_type="dataset",
+                    token=self.token
+                )
+                total_deleted_size += archive["size"]
+                logger.info(f"Deleted old archive: {archive['filename']} ({self._format_size(archive['size'])})")
+            except Exception as e:
+                logger.warning(f"Failed to delete {archive['filename']}: {e}")
+        
+        metadata["archives"] = to_keep
+        metadata["total_size"] -= total_deleted_size
+        
+        if total_deleted_size > 0:
+            logger.info(f"Cleaned up {self._format_size(total_deleted_size)} of storage")
+        
+        return metadata
+
     def download_data(self):
-        """Download and extract data from HF dataset repo"""
+        """Download and extract latest data from HF dataset repo"""
         try:
-            print("Downloading data from Hugging Face...")
+            logger.info("Downloading data from Hugging Face...")
             
             # Ensure data directory exists and is writable
             self.data_dir.mkdir(parents=True, exist_ok=True)
@@ -88,25 +236,26 @@ This repository is automatically managed by the Open WebUI sync system.
             try:
                 test_file.touch()
                 test_file.unlink()
-                print(f"Data directory {self.data_dir} is writable")
+                logger.info(f"Data directory {self.data_dir} is writable")
             except Exception as e:
-                print(f"Warning: Data directory may not be writable: {e}")
+                logger.warning(f"Data directory may not be writable: {e}")
                 return
             
             if not self.token:
-                print("No HF_TOKEN provided, skipping download")
+                logger.warning("No HF_TOKEN provided, skipping download")
                 return
             
             # Ensure repository exists
             if not self.ensure_repo_exists():
-                print("Could not access or create repository")
+                logger.error("Could not access or create repository")
                 return
             
-            # Try to download the data archive
+            # Try to download the latest data archive
             try:
+                # First try the latest link
                 file_path = self.api.hf_hub_download(
                     repo_id=self.repo_id,
-                    filename="data.tar.gz",
+                    filename=self.latest_link,
                     repo_type="dataset",
                     token=self.token
                 )
@@ -114,56 +263,154 @@ This repository is automatically managed by the Open WebUI sync system.
                 with tarfile.open(file_path, 'r:gz') as tar:
                     tar.extractall(self.data_dir)
                 
-                print(f"Data extracted to {self.data_dir}")
+                archive_size = self._get_archive_size(file_path)
+                logger.info(f"Data extracted to {self.data_dir} ({self._format_size(archive_size)})")
                 
             except Exception as e:
-                print(f"No existing data found (this is normal for first run): {e}")
+                logger.info(f"No existing data found (normal for first run): {e}")
             
         except Exception as e:
-            print(f"Error during download: {e}")
+            logger.error(f"Error during download: {e}")
     
-    def upload_data(self):
-        """Compress and upload data to HF dataset repo"""
+    def upload_data(self, force=False):
+        """Compress and upload data to HF dataset repo with change detection"""
         try:
             if not self.token:
-                print("No HF_TOKEN provided, skipping upload")
+                logger.warning("No HF_TOKEN provided, skipping upload")
                 return
-                
-            print("Uploading data to Hugging Face...")
             
             if not self.data_dir.exists() or not any(self.data_dir.iterdir()):
-                print("No data to upload")
+                logger.warning("No data to upload")
+                return
+            
+            # Calculate current directory hash
+            current_hash = self._get_directory_hash()
+            
+            # Get metadata to check for changes
+            metadata = self._get_metadata()
+            last_hash = metadata.get("last_hash")
+            
+            if not force and current_hash == last_hash:
+                logger.info("No changes detected, skipping upload")
                 return
             
+            logger.info("Changes detected, preparing upload...")
+            
             # Ensure repository exists
             if not self.ensure_repo_exists():
-                print("Could not access or create repository")
+                logger.error("Could not access or create repository")
                 return
             
-            # Create temporary archive
+            # Create timestamped filename
+            timestamp = datetime.utcnow().strftime("%Y%m%d-%H%M%S")
+            archive_filename = self.archive_pattern.format(timestamp=timestamp)
+            
+            # Create temporary archive with optimized compression
             with tempfile.NamedTemporaryFile(suffix='.tar.gz', delete=False) as tmp:
-                with tarfile.open(tmp.name, 'w:gz') as tar:
+                with tarfile.open(tmp.name, f'w:gz', compresslevel=self.compression_level) as tar:
+                    total_files = 0
                     for item in self.data_dir.iterdir():
-                        if item.name not in ["test_write", ".gitkeep"]:  # Skip test files
+                        if item.name not in ["test_write", ".gitkeep"]:
                             tar.add(item, arcname=item.name)
+                            if item.is_file():
+                                total_files += 1
+                            elif item.is_dir():
+                                total_files += sum(1 for _ in item.rglob('*') if _.is_file())
+                
+                archive_size = self._get_archive_size(tmp.name)
+                logger.info(f"Created archive: {self._format_size(archive_size)}, {total_files} files")
                 
-                # Upload to HF
+                # Upload timestamped archive
                 self.api.upload_file(
                     path_or_fileobj=tmp.name,
-                    path_in_repo="data.tar.gz",
+                    path_in_repo=archive_filename,
                     repo_id=self.repo_id,
                     repo_type="dataset",
-                    commit_message="Update Open WebUI data",
+                    commit_message=f"Update Open WebUI data - {timestamp}",
                     token=self.token
                 )
                 
-                # Clean up
+                # Upload as latest (overwrite)
+                self.api.upload_file(
+                    path_or_fileobj=tmp.name,
+                    path_in_repo=self.latest_link,
+                    repo_id=self.repo_id,
+                    repo_type="dataset",
+                    commit_message=f"Update latest data - {timestamp}",
+                    token=self.token
+                )
+                
+                # Clean up temp file
                 os.unlink(tmp.name)
             
-            print("Data uploaded successfully")
+            # Update metadata
+            archive_info = {
+                "filename": archive_filename,
+                "timestamp": timestamp,
+                "size": archive_size,
+                "hash": current_hash,
+                "files_count": total_files
+            }
+            
+            metadata["archives"].append(archive_info)
+            metadata["total_size"] += archive_size
+            metadata["last_hash"] = current_hash
+            metadata["last_upload"] = datetime.utcnow().isoformat()
+            
+            # Cleanup old archives
+            metadata = self._cleanup_old_archives(metadata)
+            
+            # Update metadata
+            self._update_metadata(metadata)
+            
+            logger.info(f"Upload successful: {archive_filename} ({self._format_size(archive_size)})")
+            logger.info(f"Total storage used: {self._format_size(metadata['total_size'])}")
             
         except Exception as e:
-            print(f"Error uploading data: {e}")
+            logger.error(f"Error uploading data: {e}")
+    
+    def list_archives(self):
+        """List all available archives"""
+        try:
+            metadata = self._get_metadata()
+            
+            if not metadata["archives"]:
+                logger.info("No archives found")
+                return
+            
+            logger.info("Available archives:")
+            logger.info("-" * 60)
+            
+            total_size = 0
+            for archive in sorted(metadata["archives"], key=lambda x: x["timestamp"], reverse=True):
+                size_str = self._format_size(archive["size"])
+                files_count = archive.get("files_count", "unknown")
+                logger.info(f"{archive['filename']:<30} {size_str:>10} {files_count:>6} files")
+                total_size += archive["size"]
+            
+            logger.info("-" * 60)
+            logger.info(f"Total: {len(metadata['archives'])} archives, {self._format_size(total_size)}")
+            
+        except Exception as e:
+            logger.error(f"Error listing archives: {e}")
+    
+    def cleanup_storage(self):
+        """Force cleanup of old archives"""
+        try:
+            metadata = self._get_metadata()
+            old_size = metadata["total_size"]
+            
+            metadata = self._cleanup_old_archives(metadata)
+            self._update_metadata(metadata)
+            
+            saved = old_size - metadata["total_size"]
+            if saved > 0:
+                logger.info(f"Cleanup completed. Saved {self._format_size(saved)} of storage")
+            else:
+                logger.info("No cleanup needed")
+                
+        except Exception as e:
+            logger.error(f"Error during cleanup: {e}")
 
 def main():
     import sys
@@ -171,18 +418,56 @@ def main():
     repo_id = os.getenv("HF_STORAGE_REPO", "nxdev-org/open-webui-storage")
     token = os.getenv("HF_TOKEN")
     data_dir = os.getenv("DATA_DIR", "/tmp/open-webui-data")
+    max_backups = int(os.getenv("MAX_BACKUPS", "3"))
     
-    sync = HFStorageSync(repo_id, token, data_dir)
+    sync = HFStorageSync(repo_id, token, data_dir, max_backups=max_backups)
     
     if len(sys.argv) > 1:
-        if sys.argv[1] == "download":
+        command = sys.argv[1].lower()
+        
+        if command == "download":
             sync.download_data()
-        elif sys.argv[1] == "upload":
-            sync.upload_data()
+        elif command == "upload":
+            force = len(sys.argv) > 2 and sys.argv[2] == "--force"
+            sync.upload_data(force=force)
+        elif command == "list":
+            sync.list_archives()
+        elif command == "cleanup":
+            sync.cleanup_storage()
         else:
-            print("Usage: sync_storage.py [download|upload]")
+            print("Usage: sync_storage.py [download|upload [--force]|list|cleanup]")
     else:
-        print("Usage: sync_storage.py [download|upload]")
+        print("Usage: sync_storage.py [download|upload [--force]|list|cleanup]")
+        print("\nCommands:")
+        print("  download          Download and extract latest data")
+        print("  upload            Upload data (only if changed)")
+        print("  upload --force    Force upload even if no changes detected")
+        print("  list              List all available archives")
+        print("  cleanup           Force cleanup of old archives")
+        print("\nEnvironment variables:")
+        print("  HF_STORAGE_REPO   Hugging Face repository ID")
+        print("  HF_TOKEN          Hugging Face API token")
+        print("  DATA_DIR          Local data directory")
+        print("  MAX_BACKUPS       Maximum number of backups to keep (default: 3)")
+
+"""
+# Download latest data
+python sync_storage.py download
 
+# Upload only if changed (default)
+python sync_storage.py upload
+
+# Force upload regardless of changes
+python sync_storage.py upload --force
+
+# List all archives
+python sync_storage.py list
+
+# Clean up old archives
+python sync_storage.py cleanup
+
+# Set custom backup limit
+MAX_BACKUPS=5 python sync_storage.py upload
+"""
 if __name__ == "__main__":
-    main()
+    main()