Spaces:

kipeles
/

apexkit

Running

apexkit / backup.py

added a list of backups

5fd8a3c 5 months ago

5.36 kB

import os
import sys
import tarfile
import shutil
import boto3
from botocore.client import Config
from flask import Flask, jsonify
from datetime import datetime

# Cloudflare R2 Config
R2_ACCOUNT_ID = os.getenv('R2_ACCOUNT_ID')
R2_ACCESS_KEY = os.getenv('R2_ACCESS_KEY')
R2_SECRET_KEY = os.getenv('R2_SECRET_KEY')
R2_BUCKET_NAME = os.getenv('R2_BUCKET_NAME')

STORAGE_DIR = 'storage'
PREFIX = 'apexkit_backup_'
MAX_BACKUPS_TO_KEEP = 5

def get_s3_client():
    if not all([R2_ACCOUNT_ID, R2_ACCESS_KEY, R2_SECRET_KEY, R2_BUCKET_NAME]):
        print("⚠️  R2 credentials missing in Hugging Face secrets. Skipping cloud storage operations.")
        return None
    
    endpoint_url = f"https://{R2_ACCOUNT_ID}.r2.cloudflarestorage.com"
    return boto3.client('s3',
        endpoint_url=endpoint_url,
        aws_access_key_id=R2_ACCESS_KEY,
        aws_secret_access_key=R2_SECRET_KEY,
        config=Config(signature_version='s3v4')
    )

def make_tarfile(output_filename, source_dir):
    with tarfile.open(output_filename, "w:gz") as tar:
        tar.add(source_dir, arcname=os.path.basename(source_dir))

def extract_tarfile(input_filename):
    with tarfile.open(input_filename, "r:gz") as tar:
        tar.extractall(path=".")

def clean_junk_folders():
    """Removes tmp and internal backup files before zipping to save massive amounts of space"""
    for folder in ['storage/tmp', 'storage/backups']:
        if os.path.exists(folder):
            shutil.rmtree(folder, ignore_errors=True)
            os.makedirs(folder, exist_ok=True)

def perform_backup():
    s3 = get_s3_client()
    if not s3:
        return False, "S3 client not configured."
    
    clean_junk_folders()

    # Create a unique timestamped filename to bust Cloudflare's cache
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    archive_name = f"{PREFIX}{timestamp}.tar.gz"
    
    print(f"📦 Compressing {STORAGE_DIR} to {archive_name}...")
    make_tarfile(archive_name, STORAGE_DIR)
    
    print(f"☁️  Uploading {archive_name} to R2 bucket '{R2_BUCKET_NAME}'...")
    try:
        s3.upload_file(archive_name, R2_BUCKET_NAME, archive_name)
        print(f"✅ Upload successful! Size: {os.path.getsize(archive_name) / (1024*1024):.2f} MB")
        
        # Clean up the local tarball so it doesn't waste disk space
        os.remove(archive_name)

        # GET RID OF THE OLD ONES
        response = s3.list_objects_v2(Bucket=R2_BUCKET_NAME, Prefix=PREFIX)
        if 'Contents' in response:
            # Sort backups by date (newest first)
            backups = sorted(response['Contents'], key=lambda x: x['LastModified'], reverse=True)
            
            # If we have more than the max allowed, delete the oldest ones
            if len(backups) > MAX_BACKUPS_TO_KEEP:
                for old_backup in backups[MAX_BACKUPS_TO_KEEP:]:
                    old_key = old_backup['Key']
                    print(f"🗑️  Deleting old backup to save space: {old_key}")
                    s3.delete_object(Bucket=R2_BUCKET_NAME, Key=old_key)

        return True, f"Backup successful: {archive_name}"
    except Exception as e:
        print(f"❌ Backup failed: {e}")
        return False, str(e)

def perform_restore():
    s3 = get_s3_client()
    if not s3:
        return False
    
    print(f"☁️  Scanning for latest backup in R2 bucket '{R2_BUCKET_NAME}'...")
    try:
        # 1. Search for NEW timestamped backups
        response = s3.list_objects_v2(Bucket=R2_BUCKET_NAME, Prefix=PREFIX)
        
        latest_backup_key = None
        
        if 'Contents' in response and len(response['Contents']) > 0:
            # Sort by date and grab the absolute newest one
            backups = sorted(response['Contents'], key=lambda x: x['LastModified'], reverse=True)
            latest_backup_key = backups[0]['Key']
        else:
            # 2. LEGACY FALLBACK: If no new backups exist, look for the old one!
            print("⚠️  No timestamped backups found. Attempting to restore legacy 'apexkit_storage_backup.tar.gz'...")
            latest_backup_key = 'apexkit_storage_backup.tar.gz'

        print(f"📥 Downloading backup: {latest_backup_key}...")
        temp_download_name = "restore.tar.gz"
        
        # Download the file
        s3.download_file(R2_BUCKET_NAME, latest_backup_key, temp_download_name)
        
        print(f"📦 Extracting {latest_backup_key}...")
        extract_tarfile(temp_download_name)
        
        # Clean up downloaded tarball
        os.remove(temp_download_name)
        
        print("✅ Restore successful! Files are fully up to date.")
        return True

    except Exception as e:
        print(f"⚠️  Restore failed (Is the bucket empty?). Error: {e}")
        return False

# --- LOCAL API SERVER ---
app = Flask(__name__)

@app.route('/backup', methods=['POST', 'GET'])
def trigger_backup():
    success, msg = perform_backup()
    if success:
        return jsonify({"status": "success", "message": msg}), 200
    else:
        return jsonify({"status": "error", "message": msg}), 500

if __name__ == '__main__':
    if len(sys.argv) > 1:
        if sys.argv[1] == '--restore':
            perform_restore()
        elif sys.argv[1] == '--serve':
            app.run(host='127.0.0.1', port=5000)
    else:
        print("Use --restore or --serve")

	import os
	import sys
	import tarfile
	import shutil
	import boto3
	from botocore.client import Config
	from flask import Flask, jsonify
	from datetime import datetime

	# Cloudflare R2 Config
	R2_ACCOUNT_ID = os.getenv('R2_ACCOUNT_ID')
	R2_ACCESS_KEY = os.getenv('R2_ACCESS_KEY')
	R2_SECRET_KEY = os.getenv('R2_SECRET_KEY')
	R2_BUCKET_NAME = os.getenv('R2_BUCKET_NAME')

	STORAGE_DIR = 'storage'
	PREFIX = 'apexkit_backup_'
	MAX_BACKUPS_TO_KEEP = 5

	def get_s3_client():
	if not all([R2_ACCOUNT_ID, R2_ACCESS_KEY, R2_SECRET_KEY, R2_BUCKET_NAME]):
	print("⚠️ R2 credentials missing in Hugging Face secrets. Skipping cloud storage operations.")
	return None

	endpoint_url = f"https://{R2_ACCOUNT_ID}.r2.cloudflarestorage.com"
	return boto3.client('s3',
	endpoint_url=endpoint_url,
	aws_access_key_id=R2_ACCESS_KEY,
	aws_secret_access_key=R2_SECRET_KEY,
	config=Config(signature_version='s3v4')
	)

	def make_tarfile(output_filename, source_dir):
	with tarfile.open(output_filename, "w:gz") as tar:
	tar.add(source_dir, arcname=os.path.basename(source_dir))

	def extract_tarfile(input_filename):
	with tarfile.open(input_filename, "r:gz") as tar:
	tar.extractall(path=".")

	def clean_junk_folders():
	"""Removes tmp and internal backup files before zipping to save massive amounts of space"""
	for folder in ['storage/tmp', 'storage/backups']:
	if os.path.exists(folder):
	shutil.rmtree(folder, ignore_errors=True)
	os.makedirs(folder, exist_ok=True)

	def perform_backup():
	s3 = get_s3_client()
	if not s3:
	return False, "S3 client not configured."

	clean_junk_folders()

	# Create a unique timestamped filename to bust Cloudflare's cache
	timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
	archive_name = f"{PREFIX}{timestamp}.tar.gz"

	print(f"📦 Compressing {STORAGE_DIR} to {archive_name}...")
	make_tarfile(archive_name, STORAGE_DIR)

	print(f"☁️ Uploading {archive_name} to R2 bucket '{R2_BUCKET_NAME}'...")
	try:
	s3.upload_file(archive_name, R2_BUCKET_NAME, archive_name)
	print(f"✅ Upload successful! Size: {os.path.getsize(archive_name) / (1024*1024):.2f} MB")

	# Clean up the local tarball so it doesn't waste disk space
	os.remove(archive_name)

	# GET RID OF THE OLD ONES
	response = s3.list_objects_v2(Bucket=R2_BUCKET_NAME, Prefix=PREFIX)
	if 'Contents' in response:
	# Sort backups by date (newest first)
	backups = sorted(response['Contents'], key=lambda x: x['LastModified'], reverse=True)

	# If we have more than the max allowed, delete the oldest ones
	if len(backups) > MAX_BACKUPS_TO_KEEP:
	for old_backup in backups[MAX_BACKUPS_TO_KEEP:]:
	old_key = old_backup['Key']
	print(f"🗑️ Deleting old backup to save space: {old_key}")
	s3.delete_object(Bucket=R2_BUCKET_NAME, Key=old_key)

	return True, f"Backup successful: {archive_name}"
	except Exception as e:
	print(f"❌ Backup failed: {e}")
	return False, str(e)

	def perform_restore():
	s3 = get_s3_client()
	if not s3:
	return False

	print(f"☁️ Scanning for latest backup in R2 bucket '{R2_BUCKET_NAME}'...")
	try:
	# 1. Search for NEW timestamped backups
	response = s3.list_objects_v2(Bucket=R2_BUCKET_NAME, Prefix=PREFIX)

	latest_backup_key = None

	if 'Contents' in response and len(response['Contents']) > 0:
	# Sort by date and grab the absolute newest one
	backups = sorted(response['Contents'], key=lambda x: x['LastModified'], reverse=True)
	latest_backup_key = backups[0]['Key']
	else:
	# 2. LEGACY FALLBACK: If no new backups exist, look for the old one!
	print("⚠️ No timestamped backups found. Attempting to restore legacy 'apexkit_storage_backup.tar.gz'...")
	latest_backup_key = 'apexkit_storage_backup.tar.gz'

	print(f"📥 Downloading backup: {latest_backup_key}...")
	temp_download_name = "restore.tar.gz"

	# Download the file
	s3.download_file(R2_BUCKET_NAME, latest_backup_key, temp_download_name)

	print(f"📦 Extracting {latest_backup_key}...")
	extract_tarfile(temp_download_name)

	# Clean up downloaded tarball
	os.remove(temp_download_name)

	print("✅ Restore successful! Files are fully up to date.")
	return True

	except Exception as e:
	print(f"⚠️ Restore failed (Is the bucket empty?). Error: {e}")
	return False

	# --- LOCAL API SERVER ---
	app = Flask(__name__)

	@app.route('/backup', methods=['POST', 'GET'])
	def trigger_backup():
	success, msg = perform_backup()
	if success:
	return jsonify({"status": "success", "message": msg}), 200
	else:
	return jsonify({"status": "error", "message": msg}), 500

	if __name__ == '__main__':
	if len(sys.argv) > 1:
	if sys.argv[1] == '--restore':
	perform_restore()
	elif sys.argv[1] == '--serve':
	app.run(host='127.0.0.1', port=5000)
	else:
	print("Use --restore or --serve")