Upload hf_sync.py with huggingface_hub

hf_sync.py

@@ -2,20 +2,64 @@ import os
 import sys
 import shutil
 import sqlite3
+import json
+import time
+from datetime import datetime
 from huggingface_hub import snapshot_download, HfApi
 
 # Configuration
 REPO_ID = os.environ.get("DATASET_REPO_ID")
 HF_TOKEN = os.environ.get("HF_TOKEN")
+DATA_DIR = "data_repo"
+DB_FILE = os.path.join(DATA_DIR, "database.db")
+METADATA_FILE = os.path.join(DATA_DIR, "sync_metadata.json")
+LOCK_FILE = "/tmp/hf_sync.lock"
+
+def get_metadata():
+    if os.path.exists(METADATA_FILE):
+        try:
+            with open(METADATA_FILE, 'r') as f:
+                return json.load(f)
+        except:
+            pass
+    return {"version": 0, "last_sync": None, "source": None}
+
+def update_metadata(action):
+    meta = get_metadata()
+    meta["version"] += 1
+    meta["last_sync"] = datetime.now().isoformat()
+    meta["source"] = os.environ.get("HOSTNAME", "local")
+    meta["last_action"] = action
+    with open(METADATA_FILE, 'w') as f:
+        json.dump(meta, f, indent=2)
+
+def safe_db_backup():
+    """Create a consistent snapshot of the SQLite database."""
+    if not os.path.exists(DB_FILE):
+        return
+    
+    print("Creating consistent database snapshot...")
+    backup_db = DB_FILE + ".tmp"
+    try:
+        source_conn = sqlite3.connect(DB_FILE)
+        dest_conn = sqlite3.connect(backup_db)
+        with dest_conn:
+            source_conn.backup(dest_conn)
+        source_conn.close()
+        dest_conn.close()
+        shutil.move(backup_db, DB_FILE)
+    except Exception as e:
+        print(f"Database backup failed: {e}")
+        if os.path.exists(backup_db):
+            os.remove(backup_db)
 
 def verify_data():
-    db_path = "data_repo/database.db"
-    if not os.path.exists(db_path):
-        print(f"VERIFICATION FAILED: {db_path} does not exist.")
+    if not os.path.exists(DB_FILE):
+        print(f"VERIFICATION FAILED: {DB_FILE} does not exist.")
         return False
     
     try:
-        conn = sqlite3.connect(db_path)
+        conn = sqlite3.connect(DB_FILE)
         cursor = conn.cursor()
         cursor.execute("SELECT username FROM users WHERE username = ?", ("akshit",))
         row = cursor.fetchone()
@@ -25,7 +69,7 @@ def verify_data():
             print("VERIFICATION SUCCESS: User 'akshit' found in database.")
             return True
         else:
-            print("VERIFICATION FAILED: User 'akshit' NOT found in database.")
+            print("VERIFICATION FAILED: User 'akshit' NOT found.")
             return False
     except Exception as e:
         print(f"VERIFICATION ERROR: {e}")
@@ -36,61 +80,63 @@ def download():
         print("DATASET_REPO_ID not set, skipping download.")
         return
     
-    print(f"Downloading data from {REPO_ID}...")
+    print(f"Downloading data from {REPO_ID} (Restore)...")
     try:
-        # snapshot_download is more efficient for many files than the CLI
         snapshot_download(
             repo_id=REPO_ID,
             repo_type="dataset",
-            local_dir="data_repo",
+            local_dir=DATA_DIR,
             token=HF_TOKEN,
             max_workers=8
         )
         print("Download successful.")
+        update_metadata("restore")
         verify_data()
     except Exception as e:
         print(f"Download failed: {e}")
 
 def upload():
-    if not REPO_ID:
-        print("DATASET_REPO_ID not set, skipping upload.")
-        return
-    if not HF_TOKEN:
-        print("HF_TOKEN not set, skipping upload.")
+    if not REPO_ID or not HF_TOKEN:
         return
 
-    print(f"Uploading data to {REPO_ID}...")
+    # Simple lock check
+    if os.path.exists(LOCK_FILE):
+        # If lock is older than 10 mins, assume it's stale
+        if time.time() - os.path.getmtime(LOCK_FILE) < 600:
+            print("Upload already in progress, skipping...")
+            return
+    
     try:
+        with open(LOCK_FILE, 'w') as f: f.write(str(os.getpid()))
+        
+        print(f"Starting scheduled backup to {REPO_ID}...")
+        safe_db_backup()
+        update_metadata("backup")
+        
         api = HfApi(token=HF_TOKEN)
         api.upload_folder(
-            folder_path="data_repo",
+            folder_path=DATA_DIR,
             repo_id=REPO_ID,
             repo_type="dataset",
-            delete_patterns="*", # Optional: sync deletion if needed
+            commit_message=f"Automated backup v{get_metadata()['version']}"
         )
         print("Upload successful.")
     except Exception as e:
         print(f"Upload failed: {e}")
+    finally:
+        if os.path.exists(LOCK_FILE):
+            os.remove(LOCK_FILE)
 
 def init_local():
-    """Ensure data_repo has the necessary structure if download failed or it's new."""
-    os.makedirs("data_repo/output", exist_ok=True)
-    os.makedirs("data_repo/processed", exist_ok=True)
-    os.makedirs("data_repo/uploads", exist_ok=True)
+    os.makedirs(f"{DATA_DIR}/output", exist_ok=True)
+    os.makedirs(f"{DATA_DIR}/processed", exist_ok=True)
+    os.makedirs(f"{DATA_DIR}/uploads", exist_ok=True)
 
 if __name__ == "__main__":
-    if len(sys.argv) < 2:
-        print("Usage: python hf_sync.py [download|upload|init|verify]")
-        sys.exit(1)
-    
-    action = sys.argv[1]
-    if action == "download":
-        download()
-    elif action == "upload":
-        upload()
-    elif action == "init":
-        init_local()
-    elif action == "verify":
-        verify_data()
+    action = sys.argv[1] if len(sys.argv) > 1 else "help"
+    if action == "download": download()
+    elif action == "upload": upload()
+    elif action == "init": init_local()
+    elif action == "verify": verify_data()
     else:
-        print(f"Unknown action: {action}")
+        print("Usage: python hf_sync.py [download|upload|init|verify]")