Update data-sync.sh

data-sync.sh

@@ -1,91 +1,112 @@
-#!/bin/bash
-# Note: Ensure that date, git, and other tools are installed in the system when the script is executed.
+# This script is injected to handle data backups to Hugging Face.
 
-# Check necessary environment variables
-if [ -z "$G_NAME" ] || [ -z "$G_TOKEN" ]; then
-    echo "Missing required environment variables G_NAME or G_TOKEN"
-    exit 1
-fi
-
-# Build GitHub repository clone URL with token
-REPO_URL="https://${G_TOKEN}@github.com/${G_NAME}.git"
-REPO_DIR="./data/github_data"
-mkdir -p "$REPO_DIR"
-
-# Clone repository
-echo "Cloning repository..."
-git clone --depth 1 "$REPO_URL" "$REPO_DIR" || {
-    echo "Clone failed, please check if G_NAME and G_TOKEN are correct."
-    exit 1
-}
-
-# Check if the initial webui.db exists and copy it
-if [ -f "$REPO_DIR/webui.db" ]; then
-    cp "$REPO_DIR/webui.db" ./data/webui.db
-    echo "Successfully pulled webui.db from GitHub repository"
+if [ -z "$HF_TOKEN" ] || [ -z "$DATASET_ID" ]; then
+    echo "Warning: HF_TOKEN or DATASET_ID not set. Backup functionality is disabled."
 else
-    echo "webui.db not found in GitHub repository, will push during sync"
-fi
+    # Create the Python sync utility
+    cat > /tmp/hf_sync.py << 'EOL'
+from huggingface_hub import HfApi
+import sys
+import os
 
-# Define sync function, performs a sync every 12 hours.
-sync_data() {
-    while true; do
-        # Output current London time during sync
-        CURRENT_TIME=$(TZ=Europe/London date '+%Y-%m-%d %H:%M:%S')
-        echo "Current time $CURRENT_TIME"
+# Keeps the latest 'max_files' backups and deletes older ones.
+def manage_backups(api, repo_id, max_files=25):
+    try:
+        files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")
+        backup_files = sorted([f for f in files if f.startswith('webui_backup_') and f.endswith('.db')])
+        
+        if len(backup_files) >= max_files:
+            num_to_delete = len(backup_files) - max_files + 1
+            files_to_delete = backup_files[:num_to_delete]
+            print(f"Pruning {num_to_delete} old backup(s)...")
+            for file_to_delete in files_to_delete:
+                try:
+                    api.delete_file(path_in_repo=file_to_delete, repo_id=repo_id, repo_type="dataset")
+                except Exception as e:
+                    print(f'Error deleting {file_to_delete}: {str(e)}')
+    except Exception as e:
+        print(f"An error occurred during backup management: {str(e)}")
+
+# Uploads a single backup file.
+def upload_backup(file_path, file_name, token, repo_id):
+    api = HfApi(token=token)
+    try:
+        print(f"Uploading backup: {file_name}...")
+        api.upload_file(
+            path_or_fileobj=file_path,
+            path_in_repo=file_name,
+            repo_id=repo_id,
+            repo_type="dataset"
+        )
+        print("Upload successful.")
+        manage_backups(api, repo_id)
+    except Exception as e:
+        print(f"File upload failed: {str(e)}")
 
-        # ---- Start sync process ----
-        echo "Starting GitHub sync..."
-        cd "$REPO_DIR" || { echo "Failed to change directory to $REPO_DIR"; exit 1; }
+# Downloads the most recent backup file.
+def download_latest_backup(token, repo_id):
+    try:
+        api = HfApi(token=token)
+        files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")
+        backup_files = sorted([f for f in files if f.startswith('webui_backup_') and f.endswith('.db')])
+        
+        if not backup_files:
+            print("No existing backups found.")
+            return
+            
+        latest_backup = backup_files[-1]
+        print(f"Restoring from latest backup: {latest_backup}")
+        
+        filepath = api.hf_hub_download(
+            repo_id=repo_id,
+            filename=latest_backup,
+            repo_type="dataset"
+        )
         
-        git config user.name "AutoSync Bot"
-        git config user.email "autosync@bot.com"
+        if filepath and os.path.exists(filepath):
+            os.makedirs('./data', exist_ok=True)
+            os.system(f'cp "{filepath}" ./data/webui.db')
+            print("Restore successful.")
+                
+    except Exception as e:
+        print(f"Could not download backup: {str(e)}")
 
-        if [ -z "$(git branch --list main)" ]; then
-            echo "Branch 'main' not found. Creating it for the initial commit."
-            git checkout -b main
-        else
-            git switch main
-        fi
+if __name__ == "__main__":
+    action = sys.argv[1]
+    token = sys.argv[2]
+    repo_id = sys.argv[3]
+    
+    if action == "upload":
+        upload_backup(sys.argv[4], sys.argv[5], token, repo_id)
+    elif action == "download":
+        download_latest_backup(token, repo_id)
+EOL
 
-        # Copy latest database file to repository directory
-        if [ -f "../webui.db" ]; then
-            cp ../webui.db ./webui.db
-        else
-            echo "Database file ../webui.db not yet initialized by the application."
-            if [ ! -f "./webui.db" ] && [ -z "$(git ls-files)" ]; then
-                echo "Creating .gitkeep to initialize the repository."
-                touch .gitkeep
-                git add .gitkeep
-            fi
-        fi
+    # On startup, restore the latest backup
+    echo "Searching for existing backup to restore..."
+    python3 /tmp/hf_sync.py download "${HF_TOKEN}" "${DATASET_ID}"
 
-        # Check if there are changes
-        if [[ -n $(git status --porcelain) ]]; then
-            git add webui.db .gitkeep
-            git commit -m "Auto sync webui.db: $(TZ=Europe/London date '+%Y-%m-%d %H:%M:%S')"
-            
-            echo "Pushing changes to GitHub..."
-            git push --set-upstream origin main && {
-                echo "GitHub push successful"
-            } || {
-                echo "Push failed, will retry once..."
-                sleep 10
-                git push origin main || {
-                    echo "Retry failed, abandoning GitHub push for this cycle."
-                }
-            }
-        else
-            echo "GitHub: No database changes detected"
-        fi
-        # Return to the original directory
-        cd ../..
+    # Define the continuous sync function
+    sync_data() {
+        while true; do
+            SYNC_INTERVAL=${SYNC_INTERVAL:-7200}
+            sleep $SYNC_INTERVAL
 
-        # ---- Sync process complete ----
-        echo "Sync complete. Waiting 12 hours for the next sync."
-        sleep 43200
-    done
-}
+            echo "Starting data sync cycle..."
+            
+            if [ -f "./data/webui.db" ]; then
+                timestamp=$(date +%Y%m%d_%H%M%S)
+                backup_file="webui_backup_${timestamp}.db"
+                temp_path="/tmp/${backup_file}"
+                
+                cp ./data/webui.db "${temp_path}"
+                python3 /tmp/hf_sync.py upload "${HF_TOKEN}" "${DATASET_ID}" "${temp_path}" "${backup_file}"
+                rm -f "${temp_path}"
+            fi
+        done
+    }
 
-# Start sync process in background
-sync_data &
+    # Start the sync process in the background
+    sync_data &
+    echo "Data backup process started in background."
+fi