Update sync_manager.py

sync_manager.py

@@ -1,195 +1,110 @@
 #!/usr/bin/env python3
-"""
-Sync Manager - Automatically backs up PostgreSQL to HuggingFace Xet Dataset
-"""
+"""Sync PostgreSQL backups to HuggingFace Xet Dataset"""
 
 import os
 import sys
 import time
 import subprocess
-import tempfile
 import signal
 from datetime import datetime
 from pathlib import Path
 
-# Configuration
-SYNC_INTERVAL = int(os.environ.get('SYNC_INTERVAL', 300))  # 5 minutes
+SYNC_INTERVAL = int(os.environ.get('SYNC_INTERVAL', 300))
 HF_TOKEN = os.environ.get('HF_TOKEN', '')
-XET_DATASET = os.environ.get('XET_DATASET', 'your-username/postgres-backup')
-POSTGRES_HOST = 'localhost'
-POSTGRES_PORT = 5432
-POSTGRES_USER = 'postgres'
-POSTGRES_DB = 'appdb'
-PGPASSWORD = os.environ.get('POSTGRES_PASSWORD', 'postgres123')
+XET_DATASET = os.environ.get('XET_DATASET', '')
 
-# Global flag for graceful shutdown
 running = True
 
+def log(msg):
+    print(f"[{datetime.now().strftime('%H:%M:%S')}] {msg}", flush=True)
 
-def log(message: str):
-    """Print timestamped log message."""
-    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
-    print(f"[{timestamp}] {message}", flush=True)
-
-
-def wait_for_postgres(max_retries: int = 30) -> bool:
-    """Wait for PostgreSQL to be ready."""
-    import psycopg2
-    
-    for i in range(max_retries):
-        try:
-            conn = psycopg2.connect(
-                host=POSTGRES_HOST,
-                port=POSTGRES_PORT,
-                user=POSTGRES_USER,
-                password=PGPASSWORD,
-                database='postgres'
-            )
-            conn.close()
-            log("✅ PostgreSQL is ready")
+def wait_for_postgres(retries=30):
+    for i in range(retries):
+        result = subprocess.run(
+            ['pg_isready', '-h', 'localhost', '-p', '5432'],
+            capture_output=True
+        )
+        if result.returncode == 0:
+            log("✅ PostgreSQL ready")
             return True
-        except psycopg2.OperationalError:
-            log(f"⏳ Waiting for PostgreSQL... ({i+1}/{max_retries})")
-            time.sleep(2)
-    
-    log("❌ PostgreSQL did not become ready")
+        time.sleep(2)
     return False
 
-
-def create_backup() -> str | None:
-    """Create a PostgreSQL backup using pg_dump."""
+def backup_and_upload():
+    if not HF_TOKEN or not XET_DATASET:
+        log("⚠️  HF_TOKEN or XET_DATASET not set")
+        return
+    
     try:
         backup_dir = Path('/data/backup')
         backup_dir.mkdir(parents=True, exist_ok=True)
-        
         backup_file = backup_dir / 'pg_backup.sql'
         
-        env = os.environ.copy()
-        env['PGPASSWORD'] = PGPASSWORD
-        
         # Create backup
         result = subprocess.run(
-            [
-                'pg_dump',
-                '-h', POSTGRES_HOST,
-                '-p', str(POSTGRES_PORT),
-                '-U', POSTGRES_USER,
-                '-d', POSTGRES_DB,
-                '--clean',
-                '--if-exists',
-                '-f', str(backup_file)
-            ],
-            env=env,
-            capture_output=True,
-            text=True
+            ['pg_dump', '-h', 'localhost', '-U', 'postgres', '-d', 'appdb', 
+             '--clean', '--if-exists', '-f', str(backup_file)],
+            capture_output=True, text=True
         )
         
-        if result.returncode == 0:
-            size = backup_file.stat().st_size
-            log(f"✅ Backup created: {backup_file} ({size} bytes)")
-            return str(backup_file)
-        else:
+        if result.returncode != 0:
             log(f"❌ Backup failed: {result.stderr}")
-            return None
-            
-    except Exception as e:
-        log(f"❌ Backup error: {e}")
-        return None
-
-
-def upload_to_xet(backup_file: str) -> bool:
-    """Upload backup to HuggingFace Xet dataset."""
-    if not HF_TOKEN:
-        log("⚠️  No HF_TOKEN set, skipping upload")
-        return False
-    
-    try:
+            return
+        
+        log(f"✅ Backup created ({backup_file.stat().st_size} bytes)")
+        
+        # Upload to HuggingFace
         from huggingface_hub import HfApi, create_repo
         
         api = HfApi(token=HF_TOKEN)
         
-        # Create dataset repo if it doesn't exist
         try:
-            create_repo(
-                repo_id=XET_DATASET,
-                repo_type="dataset",
-                private=True,
-                token=HF_TOKEN,
-                exist_ok=True
-            )
-        except Exception as e:
-            log(f"ℹ️  Repo creation note: {e}")
+            create_repo(XET_DATASET, repo_type="dataset", private=True, 
+                       token=HF_TOKEN, exist_ok=True)
+        except:
+            pass
         
-        # Upload backup file
         api.upload_file(
-            path_or_fileobj=backup_file,
+            path_or_fileobj=str(backup_file),
             path_in_repo="backup/pg_backup.sql",
             repo_id=XET_DATASET,
             repo_type="dataset",
             token=HF_TOKEN,
-            commit_message=f"PostgreSQL backup - {datetime.now().isoformat()}"
+            commit_message=f"Backup {datetime.now().isoformat()}"
         )
         
-        log(f"✅ Backup uploaded to {XET_DATASET}")
-        return True
+        log(f"✅ Uploaded to {XET_DATASET}")
         
     except Exception as e:
-        log(f"❌ Upload failed: {e}")
-        return False
-
-
-def sync_once():
-    """Perform one sync cycle."""
-    log("🔄 Starting sync cycle...")
-    
-    backup_file = create_backup()
-    if backup_file:
-        upload_to_xet(backup_file)
-    
-    log("✅ Sync cycle complete")
-
+        log(f"❌ Error: {e}")
 
 def signal_handler(signum, frame):
-    """Handle shutdown signals."""
     global running
-    log("🛑 Shutdown signal received, performing final backup...")
+    log("🛑 Shutting down, final backup...")
     running = False
-    sync_once()
-    log("👋 Sync manager shutting down")
+    backup_and_upload()
     sys.exit(0)
 
-
 def main():
-    """Main sync loop."""
-    log("🚀 Sync Manager starting...")
-    log(f"📦 Target dataset: {XET_DATASET}")
-    log(f"⏰ Sync interval: {SYNC_INTERVAL} seconds")
+    log("🚀 Sync Manager started")
+    log(f"📦 Dataset: {XET_DATASET or 'NOT SET'}")
+    log(f"⏰ Interval: {SYNC_INTERVAL}s")
     
-    # Set up signal handlers
     signal.signal(signal.SIGTERM, signal_handler)
     signal.signal(signal.SIGINT, signal_handler)
     
-    # Wait for PostgreSQL
     if not wait_for_postgres():
-        log("❌ Cannot start without PostgreSQL")
+        log("❌ PostgreSQL not available")
         return
     
-    # Initial delay
-    time.sleep(30)
+    time.sleep(30)  # Initial delay
     
-    # Main loop
     while running:
-        try:
-            sync_once()
-        except Exception as e:
-            log(f"❌ Sync error: {e}")
-        
-        # Sleep with interrupt handling
+        backup_and_upload()
         for _ in range(SYNC_INTERVAL):
             if not running:
                 break
             time.sleep(1)
 
-
 if __name__ == '__main__':
     main()