better backup restore

backend/scripts/backup.py

@@ -18,15 +18,13 @@ def get_last_backup_time(repo_id, hf_token):
         temp_file = hf_hub_download(
             repo_id=repo_id,
             repo_type="space",
-            filename=TIMESTAMP_FILE_PATH,  # "db_backup/last_backup_time.txt"
+            filename=TIMESTAMP_FILE_PATH,
             token=hf_token
         )
     except Exception as e:
-        # File might not exist or there's some other error
         print(f"Could not download {TIMESTAMP_FILE_PATH}: {e}")
         return None
 
-    # If we downloaded successfully, read the timestamp
     try:
         with open(temp_file, "r", encoding="utf-8") as f:
             timestamp_str = f.read().strip()
@@ -37,18 +35,22 @@ def get_last_backup_time(repo_id, hf_token):
         return None
 
 def save_timestamp_locally():
-    """
-    Save the current UTC time to db_backup/last_backup_time.txt (locally).
-    """
+    """Save the current UTC time to db_backup/last_backup_time.txt (locally)."""
     now = datetime.datetime.utcnow()
     os.makedirs("db_backup", exist_ok=True)
     with open(TIMESTAMP_FILE_PATH, "w", encoding="utf-8") as f:
         f.write(now.isoformat())
 
 def backup_db():
+    # Get environment variables
     passphrase = os.environ.get("BACKUP_PASSPHRASE")
     hf_token = os.environ.get("HF_TOKEN")
     space_id = os.environ.get("SPACE_ID")
+    
+    # Get threshold from environment variable, defaulting to 2 hours if not set
+    # For testing, set BACKUP_THRESHOLD_MINUTES=0 to disable the threshold check
+    threshold_minutes = int(os.environ.get("BACKUP_THRESHOLD_MINUTES", 120))
+    
     if not passphrase:
         raise ValueError("BACKUP_PASSPHRASE is not set.")
     if not hf_token:
@@ -56,37 +58,40 @@ def backup_db():
     if not space_id:
         raise ValueError("SPACE_ID is not set (or define repo_id manually).")
 
-    # 1) Check if last backup was too recent
-    threshold_hours = 2
-    last_backup_dt = get_last_backup_time(space_id, hf_token)
-    if last_backup_dt is not None:
-        now = datetime.datetime.utcnow()
-        elapsed = now - last_backup_dt
-        if elapsed.total_seconds() < threshold_hours * 3600:
-            print(f"Last backup was only {elapsed.total_seconds()/3600:.2f} hours ago.")
-            print("Skipping backup to avoid rebuild loop.")
-            return
+    # Check if backup is needed based on threshold
+    if threshold_minutes > 0:  # Only check if threshold is positive
+        last_backup_dt = get_last_backup_time(space_id, hf_token)
+        if last_backup_dt is not None:
+            now = datetime.datetime.utcnow()
+            elapsed = now - last_backup_dt
+            if elapsed.total_seconds() < threshold_minutes * 60:
+                print(f"Last backup was only {elapsed.total_seconds()/3600:.2f} hours ago.")
+                print(f"Threshold is {threshold_minutes} minutes.")
+                print("Skipping backup to avoid rebuild loop.")
+                return
+    else:
+        print("Backup threshold check disabled for testing")
 
-    # 2) Encrypt the DB using GPG
+    # Proceed with backup...
     print("Encrypting database with GPG...")
     encrypt_cmd = [
         "gpg", "--batch", "--yes", "--passphrase", passphrase,
         "-c", "--cipher-algo", "AES256",
-        "-o", DB_GPG_PATH,  # "db_backup/webui.db.gpg"
-        DB_FILE_PATH        # "data/webui.db"
+        "-o", DB_GPG_PATH,
+        DB_FILE_PATH
     ]
     subprocess.run(encrypt_cmd, check=True)
     print(f"Database encrypted successfully to {DB_GPG_PATH}")
 
-    # 3) Update the timestamp file locally
+    # Update the timestamp file locally
     save_timestamp_locally()
 
-    # 4) Upload to Hugging Face Spaces via huggingface_hub
-    print("Uploading to Hugging Face Spaces via huggingface_hub...")
+    # Upload to Hugging Face Spaces
+    print("Uploading to Hugging Face Spaces...")
     api = HfApi()
     repo_id = space_id
 
-    # Create a list of operations as CommitOperationAdd objects
+    # Create operations list for both files
     operations = [
         CommitOperationAdd(path_in_repo=DB_GPG_PATH, path_or_fileobj=DB_GPG_PATH),
         CommitOperationAdd(path_in_repo=TIMESTAMP_FILE_PATH, path_or_fileobj=TIMESTAMP_FILE_PATH)

backend/scripts/restore.py

@@ -6,93 +6,103 @@ import datetime
 from pathlib import Path
 from huggingface_hub import HfApi, hf_hub_download
 
-# Get the backend directory (equivalent to BACKEND_DIR in bash)
-BACKEND_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-
-# Define paths relative to BACKEND_DIR
-DB_GPG_PATH = os.path.join(BACKEND_DIR, "db_backup/webui.db.gpg")
-DB_FILE_PATH = os.path.join(BACKEND_DIR, "data/webui.db")
-TIMESTAMP_FILE_PATH = os.path.join(BACKEND_DIR, "db_backup/last_backup_time.txt")
+# Keep paths consistent with backup.py for better maintainability
+TIMESTAMP_FILE_PATH = "db_backup/last_backup_time.txt"
+DB_GPG_PATH = "db_backup/webui.db.gpg"
+DB_FILE_PATH = "data/webui.db"
 
 def check_requirements():
     """
     Check if GPG is installed and available.
-    Equivalent to check_requirements() in db_crypt.sh
     """
     try:
         subprocess.run(["gpg", "--version"], check=True, capture_output=True)
         return True
-    except subprocess.CalledProcessError:
-        print("Error: gpg is not installed")
-        return False
-    except FileNotFoundError:
+    except (subprocess.CalledProcessError, FileNotFoundError):
         print("Error: gpg is not installed")
         return False
 
 def validate_secrets():
     """
-    Validate that required environment variables are set.
-    Equivalent to validate_secrets() in db_crypt.sh
+    Ensure all required environment variables are set.
+    Returns False if any required variable is missing.
     """
-    missing = []
-    for var in ["BACKUP_PASSPHRASE", "HF_TOKEN"]:
-        if not os.environ.get(var):
-            missing.append(var)
+    required_vars = ["BACKUP_PASSPHRASE", "HF_TOKEN", "SPACE_ID"]
+    missing = [var for var in required_vars if not os.environ.get(var)]
     
     if missing:
-        print(f"Error: {', '.join(missing)} secret(s) not set")
+        print(f"Error: Missing required environment variables: {', '.join(missing)}")
         return False
     return True
 
 def ensure_directories():
     """
-    Create necessary directories with proper permissions.
-    Equivalent to mkdir -p commands in db_crypt.sh
+    Create necessary directories for database and backup files.
+    Uses relative paths consistent with backup.py.
     """
     try:
-        # Create data directory
-        data_dir = os.path.join(BACKEND_DIR, "data")
-        os.makedirs(data_dir, mode=0o755, exist_ok=True)
-        
-        # Create backup directory
-        backup_dir = os.path.join(BACKEND_DIR, "db_backup")
-        os.makedirs(backup_dir, mode=0o755, exist_ok=True)
-        
+        for directory in ["data", "db_backup"]:
+            os.makedirs(directory, mode=0o755, exist_ok=True)
         return True
     except Exception as e:
         print(f"Error creating directories: {e}")
         return False
 
-def download_latest_backup(repo_id, hf_token):
+def get_latest_backup_info(repo_id, hf_token):
     """
-    Download the latest encrypted database backup from HuggingFace Spaces.
+    Retrieve information about the latest backup from Hugging Face.
+    Returns a tuple of (backup_exists, timestamp) where timestamp might be None.
     """
+    api = HfApi()
     try:
-        print("Downloading latest encrypted database backup...")
-        temp_file = hf_hub_download(
-            repo_id=repo_id,
-            repo_type="space",
-            filename="db_backup/webui.db.gpg",  # Use relative path for HF download
-            token=hf_token,
-            local_dir=BACKEND_DIR  # Download directly to backend directory
-        )
+        # First check if the backup file exists in the repository
+        files = api.list_files_info(repo_id=repo_id, repo_type="space", token=hf_token)
+        backup_exists = any(file.rfilename == DB_GPG_PATH for file in files)
         
-        # Also download timestamp file for verification
+        if not backup_exists:
+            print("No backup file found in the repository")
+            return False, None
+            
+        # Try to get the timestamp information
         try:
             timestamp_file = hf_hub_download(
                 repo_id=repo_id,
                 repo_type="space",
-                filename="db_backup/last_backup_time.txt",
-                token=hf_token,
-                local_dir=BACKEND_DIR
+                filename=TIMESTAMP_FILE_PATH,
+                token=hf_token
             )
             
             with open(timestamp_file, "r", encoding="utf-8") as f:
                 timestamp = datetime.datetime.fromisoformat(f.read().strip())
                 print(f"Found backup from: {timestamp} UTC")
+                return True, timestamp
+                
         except Exception as e:
-            print(f"Note: Could not download timestamp file (this is okay for first run): {e}")
+            print(f"Note: Could not read timestamp (this is okay for first run): {e}")
+            return True, None
+            
+    except Exception as e:
+        print(f"Error checking repository: {e}")
+        return False, None
+
+def download_backup(repo_id, hf_token):
+    """
+    Download the encrypted database backup from Hugging Face.
+    Returns True if successful, False otherwise.
+    """
+    try:
+        print("Downloading encrypted database backup...")
+        temp_file = hf_hub_download(
+            repo_id=repo_id,
+            repo_type="space",
+            filename=DB_GPG_PATH,
+            token=hf_token
+        )
         
+        # Move the downloaded file to the correct location
+        os.makedirs(os.path.dirname(DB_GPG_PATH), exist_ok=True)
+        os.replace(temp_file, DB_GPG_PATH)
+        print("Backup downloaded successfully")
         return True
         
     except Exception as e:
@@ -101,11 +111,11 @@ def download_latest_backup(repo_id, hf_token):
 
 def decrypt_database(passphrase):
     """
-    Decrypt the downloaded database file using GPG.
-    Equivalent to decrypt_database() in db_crypt.sh
+    Decrypt the database file using GPG.
+    Returns True if successful or if no backup exists (first run).
     """
     if not os.path.exists(DB_GPG_PATH):
-        print("No encrypted backup found at db_backup/webui.db.gpg")
+        print("No encrypted backup found locally")
         return True  # Not an error, might be first run
         
     try:
@@ -120,7 +130,6 @@ def decrypt_database(passphrase):
             DB_GPG_PATH
         ]
         
-        # Run GPG decryption
         subprocess.run(decrypt_cmd, check=True, stderr=subprocess.PIPE)
         print(f"Database decrypted successfully to {DB_FILE_PATH}")
         return True
@@ -130,42 +139,72 @@ def decrypt_database(passphrase):
         print(f"GPG error output: {e.stderr.decode()}")
         return False
 
+def verify_database():
+    """
+    Verify the integrity of the restored database.
+    Returns True if the database is valid or doesn't exist.
+    """
+    if not os.path.exists(DB_FILE_PATH):
+        return True  # Not an error, might be first run
+        
+    try:
+        print("Verifying database integrity...")
+        verify_cmd = ["sqlite3", DB_FILE_PATH, "PRAGMA integrity_check;"]
+        result = subprocess.run(verify_cmd, capture_output=True, text=True, check=True)
+        
+        if "ok" in result.stdout.lower():
+            print("Database integrity verified")
+            return True
+        else:
+            print("Database integrity check failed")
+            return False
+            
+    except subprocess.CalledProcessError as e:
+        print(f"Database verification failed: {e}")
+        return False
+
 def restore_db():
     """
     Main function to handle the database restoration process.
+    Downloads the latest backup from Hugging Face and restores it.
     """
-    os.chdir(BACKEND_DIR)  # Ensure we're in the backend directory
-    
-    # Check requirements first
-    if not check_requirements():
-        return False
-        
-    # Validate environment variables
-    if not validate_secrets():
+    # Check requirements and environment
+    if not check_requirements() or not validate_secrets():
         return False
     
     # Ensure required directories exist
     if not ensure_directories():
         return False
     
-    passphrase = os.environ.get("BACKUP_PASSPHRASE")
-    hf_token = os.environ.get("HF_TOKEN")
-    space_id = os.environ.get("SPACE_ID")
+    # Get environment variables
+    passphrase = os.environ["BACKUP_PASSPHRASE"]
+    hf_token = os.environ["HF_TOKEN"]
+    space_id = os.environ["SPACE_ID"]
     
-    # Download latest backup if space_id is provided
-    if space_id:
-        if not download_latest_backup(space_id, hf_token):
-            print("Warning: Failed to download latest backup")
-            # Continue anyway - we might have a local backup
+    # Check if there's a backup in the repository
+    backup_exists, timestamp = get_latest_backup_info(space_id, hf_token)
     
-    # Decrypt the database
-    if not decrypt_database(passphrase):
-        print("Failed to decrypt database")
-        return False
+    if backup_exists:
+        if not download_backup(space_id, hf_token):
+            print("Failed to download backup")
+            return False
+        
+        if not decrypt_database(passphrase):
+            print("Failed to decrypt database")
+            return False
+            
+        if not verify_database():
+            print("Failed to verify database integrity")
+            # Remove potentially corrupted database
+            if os.path.exists(DB_FILE_PATH):
+                os.unlink(DB_FILE_PATH)
+            return False
+    else:
+        print("No backup found - starting with fresh database")
     
     print("Database restore completed successfully!")
     return True
 
 if __name__ == "__main__":
     success = restore_db()
-    sys.exit(0 if success else 1)
+    sys.exit(0 if success else 1)

backend/start.sh

@@ -3,30 +3,34 @@
 SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 cd "$SCRIPT_DIR" || exit
 
+# Validate required environment variables for backup/restore operations
+for var in "BACKUP_PASSPHRASE" "HF_TOKEN" "SPACE_ID"; do
+    if [ -z "${!var}" ]; then
+        echo "Error: $var is not set. This is required for database backup/restore operations."
+        exit 1
+    fi
+done
+
 # Restore database from backup
 echo "Restoring database from backup..."
 python scripts/restore.py
 restore_status=$?
-
 if [ $restore_status -ne 0 ]; then
     echo "Warning: Database restore failed. Starting with empty database."
 fi
 
+# Handle WebUI secret key generation/loading
 KEY_FILE=.webui_secret_key
-
 PORT="${PORT:-8080}"
 HOST="${HOST:-0.0.0.0}"
 if test "$WEBUI_SECRET_KEY $WEBUI_JWT_SECRET_KEY" = " "; then
-  echo "Loading WEBUI_SECRET_KEY from file, not provided as an environment variable."
-
-  if ! [ -e "$KEY_FILE" ]; then
-    echo "Generating WEBUI_SECRET_KEY"
-    # Generate a random value to use as a WEBUI_SECRET_KEY in case the user didn't provide one.
-    echo $(head -c 12 /dev/random | base64) > "$KEY_FILE"
-  fi
-
-  echo "Loading WEBUI_SECRET_KEY from $KEY_FILE"
-  WEBUI_SECRET_KEY=$(cat "$KEY_FILE")
+    echo "Loading WEBUI_SECRET_KEY from file, not provided as an environment variable."
+    if ! [ -e "$KEY_FILE" ]; then
+        echo "Generating WEBUI_SECRET_KEY"
+        echo $(head -c 12 /dev/random | base64) > "$KEY_FILE"
+    fi
+    echo "Loading WEBUI_SECRET_KEY from $KEY_FILE"
+    WEBUI_SECRET_KEY=$(cat "$KEY_FILE")
 fi
 
 if [[ "${USE_OLLAMA_DOCKER,,}" == "true" ]]; then
@@ -35,53 +39,58 @@ if [[ "${USE_OLLAMA_DOCKER,,}" == "true" ]]; then
 fi
 
 if [[ "${USE_CUDA_DOCKER,,}" == "true" ]]; then
-  echo "CUDA is enabled, appending LD_LIBRARY_PATH to include torch/cudnn & cublas libraries."
-  export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/python3.11/site-packages/torch/lib:/usr/local/lib/python3.11/site-packages/nvidia/cudnn/lib"
+    echo "CUDA is enabled, appending LD_LIBRARY_PATH to include torch/cudnn & cublas libraries."
+    export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/python3.11/site-packages/torch/lib:/usr/local/lib/python3.11/site-packages/nvidia/cudnn/lib"
 fi
 
-# Check if SPACE_ID is set, if so, configure for space
+# Handle HuggingFace Space deployment configuration
 if [ -n "$SPACE_ID" ]; then
-  echo "Configuring for HuggingFace Space deployment"
-  if [ -n "$ADMIN_USER_EMAIL" ] && [ -n "$ADMIN_USER_PASSWORD" ]; then
-    echo "Admin user configured, creating"
-    WEBUI_SECRET_KEY="$WEBUI_SECRET_KEY" uvicorn open_webui.main:app --host "$HOST" --port "$PORT" --forwarded-allow-ips '*' &
-    webui_pid=$!
-    echo "Waiting for webui to start..."
-    while ! curl -s http://localhost:8080/health > /dev/null; do
-      sleep 1
-    done
-    echo "Creating admin user..."
-    curl \
-      -X POST "http://localhost:8080/api/v1/auths/signup" \
-      -H "accept: application/json" \
-      -H "Content-Type: application/json" \
-      -d "{ \"email\": \"${ADMIN_USER_EMAIL}\", \"password\": \"${ADMIN_USER_PASSWORD}\", \"name\": \"Admin\" }"
-    echo "Shutting down webui..."
-    kill $webui_pid
-  fi
-
-  export WEBUI_URL=${SPACE_HOST}
+    echo "Configuring for HuggingFace Space deployment"
+    if [ -n "$ADMIN_USER_EMAIL" ] && [ -n "$ADMIN_USER_PASSWORD" ]; then
+        echo "Admin user configured, creating"
+        WEBUI_SECRET_KEY="$WEBUI_SECRET_KEY" uvicorn open_webui.main:app --host "$HOST" --port "$PORT" --forwarded-allow-ips '*' &
+        webui_pid=$!
+        echo "Waiting for webui to start..."
+        while ! curl -s http://localhost:8080/health > /dev/null; do
+            sleep 1
+        done
+        echo "Creating admin user..."
+        curl \
+            -X POST "http://localhost:8080/api/v1/auths/signup" \
+            -H "accept: application/json" \
+            -H "Content-Type: application/json" \
+            -d "{ \"email\": \"${ADMIN_USER_EMAIL}\", \"password\": \"${ADMIN_USER_PASSWORD}\", \"name\": \"Admin\" }"
+        echo "Shutting down temporary webui instance..."
+        kill $webui_pid
+    fi
+    export WEBUI_URL=${SPACE_HOST}
 fi
 
-# 6. Launch the main web server in the background (DO NOT use 'exec' here)
+# Launch the main web server in the background
 WEBUI_SECRET_KEY="$WEBUI_SECRET_KEY" uvicorn open_webui.main:app \
-  --host "$HOST" --port "$PORT" --forwarded-allow-ips '*' &
+    --host "$HOST" --port "$PORT" --forwarded-allow-ips '*' &
 WEBUI_PID=$!
 
-# 7. Background job for daily backup
+# Configure backup schedule
+BACKUP_INITIAL_WAIT="${BACKUP_INITIAL_WAIT:-500}"  # Default to 500 seconds
+BACKUP_INTERVAL="${BACKUP_INTERVAL:-21600}"        # Default to 6 hours (21600 seconds)
+
+# Start the background backup job
 (
-  # Optional: wait a few minutes before the first backup, so the DB is stable
-  sleep 500
+    echo "Waiting ${BACKUP_INITIAL_WAIT} seconds before starting backup schedule..."
+    sleep "$BACKUP_INITIAL_WAIT"
 
-  while true; do
-    echo "==============================================="
-    echo "Daily backup job: encrypting and pushing the DB"
-    echo "==============================================="
-    python "$SCRIPT_DIR/backup.py"
-    # Sleep for 8 hours
-    sleep 21600
-  done
+    while true; do
+        echo "==============================================="
+        echo "Running scheduled backup job"
+        echo "==============================================="
+        # Allow threshold override for testing
+        BACKUP_THRESHOLD_MINUTES="${BACKUP_THRESHOLD_MINUTES:-120}" python "$SCRIPT_DIR/backup.py"
+        
+        echo "Waiting ${BACKUP_INTERVAL} seconds until next backup..."
+        sleep "$BACKUP_INTERVAL"
+    done
 ) &
 
-# 8. Wait on the main web server (keep container alive)
+# Wait for the main web server process to keep container alive
 wait $WEBUI_PID