src/database_sync.py

"""
Database Sync to Hugging Face Hub
Syncs SQLite database to HF Hub for persistence on ephemeral storage
"""

import os
import shutil
from pathlib import Path
from typing import Optional

class DatabaseSync:
    """Sync SQLite database to Hugging Face Hub"""
    
    def __init__(self, db_path: str, hf_repo_id: Optional[str] = None, hf_token: Optional[str] = None):
        """
        Initialize database sync
        
        Args:
            db_path: Path to SQLite database file
            hf_repo_id: HF Hub repository ID (e.g., 'username/dataset-name')
            hf_token: HF Hub token (if None, tries to get from environment)
        """
        self.db_path = Path(db_path)
        self.hf_repo_id = hf_repo_id or os.environ.get('HF_DATASET_REPO', 'hamza-ksr/bbPlease-db')
        self.hf_token = hf_token or os.environ.get('HF_TOKEN')
        self.enabled = bool(self.hf_token)
        
        if self.enabled:
            try:
                from huggingface_hub import HfApi, login
                self.HfApi = HfApi
                self.login = login
                login(token=self.hf_token)
                self.hf_api = HfApi(token=self.hf_token)
            except ImportError:
                print("⚠️  huggingface_hub not installed. Database sync disabled.")
                self.enabled = False
            except Exception as e:
                print(f"⚠️  Failed to initialize HF Hub: {e}. Database sync disabled.")
                self.enabled = False
    
    def sync_to_hub(self) -> bool:
        """Upload database to HF Hub"""
        if not self.enabled:
            return False
        
        if not self.db_path.exists():
            print("⚠️  Database file not found, skipping sync")
            return False
        
        try:
            # Upload database file
            self.hf_api.upload_file(
                path_or_fileobj=str(self.db_path),
                path_in_repo="app.db",
                repo_id=self.hf_repo_id,
                repo_type="dataset",
                commit_message=f"Auto-sync database at {Path(self.db_path).stat().st_mtime}"
            )
            print(f"✅ Database synced to HF Hub: {self.hf_repo_id}")
            return True
        except Exception as e:
            print(f"⚠️  Failed to sync database to HF Hub: {e}")
            return False
    
    def load_from_hub(self) -> bool:
        """Download database from HF Hub if it doesn't exist locally"""
        if not self.enabled:
            return False
        
        # If database already exists, don't overwrite (could cause data loss)
        if self.db_path.exists():
            return False
        
        try:
            from huggingface_hub import hf_hub_download
            
            # Ensure parent directory exists
            self.db_path.parent.mkdir(parents=True, exist_ok=True)
            
            # Try to download database file
            try:
                downloaded_path = hf_hub_download(
                    repo_id=self.hf_repo_id,
                    filename="app.db",
                    repo_type="dataset",
                    token=self.hf_token,
                    local_dir=str(self.db_path.parent),
                    local_dir_use_symlinks=False
                )
                # Move to expected location if needed
                if downloaded_path != str(self.db_path):
                    shutil.move(downloaded_path, str(self.db_path))
                print(f"✅ Database loaded from HF Hub: {self.hf_repo_id}")
                return True
            except Exception as e:
                # Database doesn't exist in Hub yet, that's okay
                print(f"ℹ️  No existing database in HF Hub (this is normal for first deployment): {e}")
                return False
        except ImportError:
            print("⚠️  huggingface_hub not installed. Cannot load from HF Hub.")
            return False
        except Exception as e:
            print(f"⚠️  Failed to load database from HF Hub: {e}")
            return False
    
    def sync_if_needed(self, force: bool = False) -> bool:
        """Sync database to Hub (can be called periodically)"""
        if not self.enabled:
            return False
        
        if force or self._should_sync():
            return self.sync_to_hub()
        return False
    
    def _should_sync(self) -> bool:
        """Determine if sync is needed based on file modification time"""
        # Simple strategy: sync if database exists
        # For production, you might want to sync based on time or change tracking
        return self.db_path.exists()