diff --git a/regularbm.py b/regularbm.py new file mode 100644 index 0000000..3a2e7ec --- /dev/null +++ b/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) \ No newline at end of file