File size: 4,653 Bytes
469692c 9391632 84e6d52 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """
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()
|