more fixes

backend/scripts/restore.py

@@ -8,21 +8,28 @@ from pathlib import Path
 
 from huggingface_hub import HfApi, hf_hub_download
 
-# ------------------------------------------------------------------------------
-# 1. Define directories outside of /app to avoid read-only file system issues
-# ------------------------------------------------------------------------------
-RESTORE_BACKUP_DIR = "/tmp/open_webui/db_backup"
-RESTORE_DATA_DIR = "/tmp/open_webui/data"
+###############################################################################
+# 1) Determine the *final* database file location.
+#    - We read DATA_DIR from env, defaulting to "/app/backend/data".
+#    - The final DB (decrypted) should be "/app/backend/data/webui.db" so
+#      the app will use the restored DB on startup.
+###############################################################################
+DATA_DIR = os.environ.get("DATA_DIR", "/app/backend/data")
+DB_FILE_PATH = os.path.join(DATA_DIR, "webui.db")
 
-# Actual paths on the local (runtime) filesystem
-TIMESTAMP_FILE_PATH = os.path.join(RESTORE_BACKUP_DIR, "last_backup_time.txt")
+###############################################################################
+# 2) Where we put the encrypted backup and timestamp locally.
+#    - A writable directory on Hugging Face is "/tmp/open_webui/db_backup".
+###############################################################################
+RESTORE_BACKUP_DIR = "/tmp/open_webui/db_backup"
 DB_GPG_PATH = os.path.join(RESTORE_BACKUP_DIR, "webui.db.gpg")
-DB_FILE_PATH = os.path.join(RESTORE_DATA_DIR, "webui.db")
+TIMESTAMP_FILE_PATH = os.path.join(RESTORE_BACKUP_DIR, "last_backup_time.txt")
 
-# Paths used in the Hugging Face Space repository
+# Paths in the Hugging Face repo
 REPO_TIMESTAMP_FILE = "db_backup/last_backup_time.txt"
 REPO_DB_GPG_FILE = "db_backup/webui.db.gpg"
 
+
 def check_requirements():
     """
     Verify that GPG is installed and available in the system.
@@ -31,9 +38,10 @@ def check_requirements():
         subprocess.run(["gpg", "--version"], check=True, capture_output=True)
         return True
     except (subprocess.CalledProcessError, FileNotFoundError):
-        print("Error: gpg is not installed")
+        print("Error: gpg is not installed or not in PATH.")
         return False
 
+
 def validate_secrets():
     """
     Ensure all required environment variables are set: BACKUP_PASSPHRASE, HF_TOKEN, SPACE_ID.
@@ -46,21 +54,27 @@ def validate_secrets():
         return False
     return True
 
+
 def ensure_directories():
     """
-    Create necessary directories for database and backup files (in /tmp).
+    Create necessary directories for database (DATA_DIR) and backup files (RESTORE_BACKUP_DIR).
+    We must ensure the final DB directory is writable so the decrypted DB can be placed there.
     """
     try:
-        for directory in [RESTORE_DATA_DIR, RESTORE_BACKUP_DIR]:
-            os.makedirs(directory, mode=0o755, exist_ok=True)
+        # Create the final DB folder (if not already existing)
+        os.makedirs(DATA_DIR, mode=0o755, exist_ok=True)
+
+        # Create the backup folder in /tmp
+        os.makedirs(RESTORE_BACKUP_DIR, mode=0o755, exist_ok=True)
         return True
     except Exception as e:
         print(f"Error creating directories: {e}")
         return False
 
+
 def get_latest_backup_info(repo_id, hf_token):
     """
-    Check if a backup exists in the HF Space. If so, return True and a timestamp (or None).
+    Check if a backup (webui.db.gpg) exists in the HF Space. If so, return True + timestamp (or None).
     """
     api = HfApi()
     try:
@@ -69,12 +83,12 @@ def get_latest_backup_info(repo_id, hf_token):
         # Check if the encrypted DB is present
         backup_exists = (REPO_DB_GPG_FILE in files)
         if not backup_exists:
-            print("No backup file found in the repository")
+            print("No backup file found in the repository.")
             return False, None
         
         # Attempt to fetch the timestamp file
-        try:
-            if REPO_TIMESTAMP_FILE in files:
+        if REPO_TIMESTAMP_FILE in files:
+            try:
                 timestamp_file = hf_hub_download(
                     repo_id=repo_id,
                     repo_type="space",
@@ -82,22 +96,24 @@ def get_latest_backup_info(repo_id, hf_token):
                     token=hf_token
                 )
                 with open(timestamp_file, "r", encoding="utf-8") as f:
-                    timestamp = datetime.datetime.fromisoformat(f.read().strip())
+                    timestamp_str = f.read().strip()
+                timestamp = datetime.datetime.fromisoformat(timestamp_str)
                 print(f"Found backup from: {timestamp} UTC")
                 return True, timestamp
-            else:
-                print("No timestamp file found, but backup file exists")
+            except Exception as e:
+                print(f"Could not read timestamp (possibly first run): {e}")
                 return True, None
-        except Exception as e:
-            print(f"Could not read timestamp (possibly first run): {e}")
+        else:
+            print("No timestamp file found, but backup file exists.")
             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 into /tmp/open_webui/db_backup.
+    Download the encrypted database backup from Hugging Face to /tmp/open_webui/db_backup.
     """
     try:
         print("Downloading encrypted database backup...")
@@ -107,19 +123,19 @@ def download_backup(repo_id, hf_token):
             filename=REPO_DB_GPG_FILE,
             token=hf_token
         )
-        
         # Move the downloaded file to DB_GPG_PATH
-        os.makedirs(os.path.dirname(DB_GPG_PATH), exist_ok=True)
         os.replace(temp_file, DB_GPG_PATH)
-        print("Backup downloaded successfully")
+        print("Backup downloaded successfully.")
         return True
     except Exception as e:
         print(f"Error downloading backup: {e}")
         return False
 
+
 def decrypt_database(passphrase):
     """
-    Decrypt the database file using GPG, writing the decrypted DB into /tmp/open_webui/data.
+    Decrypt the database into DATA_DIR/webui.db so the app can use it.
+    The encrypted file is at /tmp/open_webui/db_backup/webui.db.gpg.
     """
     if not os.path.exists(DB_GPG_PATH):
         print("No encrypted backup found locally. Proceeding with fresh DB.")
@@ -150,12 +166,13 @@ def decrypt_database(passphrase):
             print(f"GPG error output: {e.stderr.decode(errors='ignore')}")
         return False
 
+
 def verify_database():
     """
-    Verify the integrity of the restored database using SQLite's built-in integrity check.
+    Verify the integrity of the restored DB at DATA_DIR/webui.db using SQLite's built-in checks.
     """
     if not os.path.exists(DB_FILE_PATH):
-        # If there's no DB yet, that's okay (fresh start)
+        # If there's no DB yet, might be fresh start, not an immediate error
         return True
     
     try:
@@ -168,15 +185,15 @@ def verify_database():
             tables = cursor.fetchall()
             
             if result.lower() == "ok" and len(tables) > 0:
-                print("Database integrity verified successfully")
-                print(f"Found {len(tables)} tables in database")
+                print("Database integrity verified successfully.")
+                print(f"Found {len(tables)} tables in database.")
                 return True
             else:
-                print("Database integrity check failed")
+                print("Database integrity check failed.")
                 if result.lower() != "ok":
                     print(f"Integrity check result: {result}")
                 if len(tables) == 0:
-                    print("No tables found in database")
+                    print("No tables found in database.")
                 return False
     except sqlite3.Error as e:
         print(f"Database verification failed: {e}")
@@ -185,13 +202,14 @@ def verify_database():
         print(f"Unexpected error during database verification: {e}")
         return False
 
+
 def restore_db():
     """
     Main restore function:
-    1. Check GPG installation & environment variables
-    2. Create directories
-    3. Check if remote backup exists, download & decrypt
-    4. Verify DB
+      1. Check GPG + environment
+      2. Create directories
+      3. Check if remote backup exists, download + decrypt
+      4. Verify DB
     """
     if not check_requirements() or not validate_secrets():
         return False
@@ -206,25 +224,26 @@ def restore_db():
     backup_exists, timestamp = get_latest_backup_info(space_id, hf_token)
     if backup_exists:
         if not download_backup(space_id, hf_token):
-            print("Failed to download backup")
+            print("Failed to download backup.")
             return False
         
         if not decrypt_database(passphrase):
-            print("Failed to decrypt database")
+            print("Failed to decrypt database.")
             return False
         
         if not verify_database():
-            print("Database integrity verification failed")
-            # Remove the corrupted DB so we don't keep a broken file
+            print("Database integrity verification failed.")
+            # Remove the corrupted DB so we don’t keep a broken file
             if os.path.exists(DB_FILE_PATH):
                 os.unlink(DB_FILE_PATH)
             return False
     else:
-        print("No backup found - starting with an empty/fresh database")
+        print("No backup found - starting with an empty/fresh database.")
     
     print("Database restore completed successfully!")
     return True
 
+
 if __name__ == "__main__":
     success = restore_db()
     sys.exit(0 if success else 1)