Add regularbm.py

regularbm.py

@@ -0,0 +1,163 @@
+#!/usr/bin/env python3
+
+import os
+import sys
+import json
+import sqlite3
+import datetime
+import shutil
+import subprocess
+import argparse
+from collections import defaultdict
+
+try:
+    import boto3
+    import lz4.frame
+except ImportError:
+    print("Error: Required Python packages 'boto3' or 'lz4' are not installed.")
+    print("Please run: pip3 install boto3 lz4")
+    sys.exit(1)
+
+def main(args):
+    """Main function to orchestrate the backup process."""
+    print(f"Starting Open WebUI backup process at {datetime.datetime.now().isoformat()}")
+
+    # 1. Setup temporary directory for staging backup files
+    timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H%M%S")
+    backup_staging_dir = os.path.join(args.tmp_dir, f"openwebui_backup_{timestamp}")
+    db_local_path = os.path.join(backup_staging_dir, "webui.db")
+
+    try:
+        os.makedirs(backup_staging_dir, exist_ok=True)
+        print(f"Created temporary staging directory: {backup_staging_dir}")
+
+        # 2. Copy database from Docker container
+        copy_db_from_docker(args.container_name, db_local_path)
+
+        # 3. Export data from the database
+        export_data_from_db(db_local_path, backup_staging_dir)
+
+        # 4. Compress the backup directory
+        archive_path = compress_backup(backup_staging_dir, timestamp)
+
+        # 5. Upload to S3
+        upload_to_s3(archive_path, args.s3_bucket, args.s3_region, args.s3_prefix)
+
+    except Exception as e:
+        print(f"An error occurred during the backup process: {e}", file=sys.stderr)
+        sys.exit(1)
+    finally:
+        # 6. Cleanup local files
+        print("Cleaning up local temporary files...")
+        if os.path.exists(backup_staging_dir):
+            shutil.rmtree(backup_staging_dir)
+        if 'archive_path' in locals() and os.path.exists(archive_path):
+            os.remove(archive_path)
+        print("Cleanup complete.")
+
+    print(f"Backup process finished successfully at {datetime.datetime.now().isoformat()}")
+
+def copy_db_from_docker(container_name, local_path):
+    """Copies the SQLite database from the Docker container to a local path."""
+    db_container_path = "/app/backend/data/webui.db"
+    command = ["docker", "cp", f"{container_name}:{db_container_path}", local_path]
+    print(f"Executing: {' '.join(command)}")
+    
+    try:
+        result = subprocess.run(command, check=True, capture_output=True, text=True)
+        print("Successfully copied database from Docker container.")
+    except FileNotFoundError:
+        print("Error: 'docker' command not found. Is Docker installed and in your PATH?", file=sys.stderr)
+        raise
+    except subprocess.CalledProcessError as e:
+        print(f"Error copying database from Docker container '{container_name}'.", file=sys.stderr)
+        print(f"Stderr: {e.stderr}", file=sys.stderr)
+        print("Please ensure the container is running and the name is correct.", file=sys.stderr)
+        raise
+
+def export_data_from_db(db_path, output_dir):
+    """Connects to the SQLite DB and exports chats and config."""
+    print("Connecting to the database and exporting data...")
+    conn = None
+    try:
+        conn = sqlite3.connect(db_path)
+        cursor = conn.cursor()
+
+        # Export all user chats
+        cursor.execute("""
+            SELECT user.email, user.name, chat.title, chat.chat 
+            FROM chat JOIN user ON chat.user_id = user.id
+        """)
+        rows = cursor.fetchall()
+
+        user_chats = defaultdict(list)
+        for email, name, title, chat_json_str in rows:
+            try:
+                chat_content = json.loads(chat_json_str)
+            except json.JSONDecodeError:
+                chat_content = {"error": "Failed to parse chat JSON."}
+            
+            user_chats[email].append({"title": title, "messages": chat_content})
+        
+        for email, chats in user_chats.items():
+            safe_email = email.replace('@', '_at_').replace('.', '_')
+            filename = os.path.join(output_dir, f"chats_{safe_email}.json")
+            with open(filename, "w") as f:
+                json.dump(chats, f, indent=2)
+        print(f"Exported chats for {len(user_chats)} users.")
+
+        # Export config
+        cursor.execute("SELECT data FROM config ORDER BY id DESC LIMIT 1")
+        config_row = cursor.fetchone()
+        if config_row:
+            config_data = json.loads(config_row[0])
+            with open(os.path.join(output_dir, "openwebui_config.json"), "w") as f:
+                json.dump(config_data, f, indent=2)
+            print("Successfully exported configuration.")
+
+    finally:
+        if conn:
+            conn.close()
+
+def compress_backup(source_dir, timestamp):
+    """Creates a .tar archive and then compresses it with lz4."""
+    print("Compressing backup files...")
+    base_name = f"openwebui_backup_{timestamp}"
+    tar_path = os.path.join(os.path.dirname(source_dir), f"{base_name}.tar")
+    lz4_path = f"{tar_path}.lz4"
+
+    # Create a tar archive first
+    shutil.make_archive(base_name=tar_path.replace('.tar',''), format='tar', root_dir=source_dir)
+    
+    # Compress the tar file with lz4
+    with open(tar_path, 'rb') as f_in, open(lz4_path, 'wb') as f_out:
+        f_out.write(lz4.frame.compress(f_in.read()))
+    
+    # Clean up the intermediate tar file
+    os.remove(tar_path)
+    print(f"Successfully created compressed archive: {lz4_path}")
+    return lz4_path
+
+def upload_to_s3(file_path, bucket, region, prefix):
+    """Uploads a file to an S3 bucket."""
+    s3_client = boto3.client("s3", region_name=region)
+    s3_key = f"{prefix}{os.path.basename(file_path)}"
+    
+    print(f"Uploading {os.path.basename(file_path)} to s3://{bucket}/{s3_key}...")
+    try:
+        s3_client.upload_file(file_path, bucket, s3_key)
+        print("Upload successful.")
+    except Exception as e:
+        print(f"Error uploading to S3: {e}", file=sys.stderr)
+        raise
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Backup Open WebUI data to an S3 bucket.")
+    parser.add_argument("--s3-bucket", required=True, help="S3 bucket name for backups.")
+    parser.add_argument("--s3-region", required=True, help="AWS region of the S3 bucket.")
+    parser.add_argument("--container-name", default="open-webui", help="Name of the Open WebUI Docker container.")
+    parser.add_argument("--s3-prefix", default="openwebui_backups/", help="Prefix (folder) in the S3 bucket.")
+    parser.add_argument("--tmp-dir", default="/tmp", help="Temporary directory for staging files.")
+    
+    args = parser.parse_args()
+    main(args)