test.py

#!/usr/bin/env python3
"""
Test script for optimized S3 upload (direct for small files, parallel multipart for large files).
Tests uploading ComfyUI_0001_.png using the optimized upload function.
"""

import os
import sys
import boto3
from botocore.config import Config
from boto3.s3.transfer import TransferConfig
from botocore.exceptions import ClientError, NoCredentialsError
import time
import urllib.parse

def optimized_upload(file_path, bucket_name, object_name, s3_client, content_type=None):
    """
    Upload a file to S3 using optimized upload (direct for small files, 
    parallel multipart for large files).
    
    Args:
        file_path (str): Path to the file to upload
        bucket_name (str): Name of the S3 bucket
        object_name (str): S3 object key (path)
        s3_client: Configured boto3 S3 client
        content_type (str, optional): Content type/MIME type for the file
    
    Returns:
        None
    
    Raises:
        Exception: If upload fails
    """
    try:
        file_size = os.path.getsize(file_path)
        file_size_mb = file_size / (1024 * 1024)
        
        # Configure transfer settings for optimal performance
        # - multipart_threshold: Use multipart for files > 16MB
        # - max_concurrency: Upload 10 parts in parallel
        # - multipart_chunksize: 8MB chunks (good balance)
        # - use_threads: Enable parallel uploads
        transfer_config = TransferConfig(
            multipart_threshold=16 * 1024 * 1024,  # 16MB threshold
            max_concurrency=10,  # Upload 10 parts in parallel
            multipart_chunksize=8 * 1024 * 1024,  # 8MB chunks
            use_threads=True  # Enable threading for parallel uploads
        )
        
        # Prepare upload parameters
        upload_params = {
            'Filename': file_path,
            'Bucket': bucket_name,
            'Key': object_name,
            'Config': transfer_config
        }
        
        # Add content type if provided
        extra_args = {}
        if content_type:
            extra_args['ContentType'] = content_type
        
        if extra_args:
            upload_params['ExtraArgs'] = extra_args
        
        upload_type = "parallel multipart" if file_size > 16 * 1024 * 1024 else "direct"
        print(f"📤 Uploading file using {upload_type} ({file_size_mb:.2f} MB)...")
        
        # Perform optimized upload
        s3_client.upload_file(**upload_params)
        
        print(f"✅ File uploaded successfully via {upload_type} upload")
        
    except NoCredentialsError:
        error_msg = "Credentials not available"
        print(f"❌ ERROR: {error_msg}")
        raise Exception(error_msg)
    except ClientError as e:
        error_code = e.response.get('Error', {}).get('Code', '')
        error_message = e.response.get('Error', {}).get('Message', str(e))
        error_msg = f"S3 upload error ({error_code}): {error_message}"
        print(f"❌ ERROR: {error_msg}")
        raise Exception(error_msg)
    except Exception as e:
        error_msg = f"Unexpected error during upload: {e}"
        print(f"❌ ERROR: {error_msg}")
        raise Exception(error_msg)

def test_upload(file_path=None, s3_key=None):
    """
    Test uploading ComfyUI_0001_.png to S3 with optimized upload.
    
    Args:
        file_path: Optional path to file. If not provided, searches for ComfyUI_0001_.png
        s3_key: Optional S3 key (path). If not provided, uses test-uploads/filename
    """
    
    # Get S3 configuration from environment variables
    endpoint_url = os.environ.get("BUCKET_ENDPOINT_URL")
    bucket_name = os.environ.get("BUCKET_NAME")
    access_key = os.environ.get("BUCKET_ACCESS_KEY_ID")
    secret_key = os.environ.get("BUCKET_SECRET_ACCESS_KEY")
    region = os.environ.get("S3_REGION", "auto")
    
    if not all([endpoint_url, bucket_name, access_key, secret_key]):
        print("❌ S3 not configured. Please set:")
        print("   - BUCKET_ENDPOINT_URL")
        print("   - BUCKET_NAME")
        print("   - BUCKET_ACCESS_KEY_ID")
        print("   - BUCKET_SECRET_ACCESS_KEY")
        return False
    
    # Find the file if not provided
    if not file_path:
        # Try different possible filenames
        possible_names = [
            "ComfyUI_0001_.png",
            "ComfyUI_00001_.png",
            "./ComfyUI_0001_.png",
            "./ComfyUI_00001_.png"
        ]
        
        file_path = None
        for name in possible_names:
            if os.path.exists(name):
                file_path = name
                break
        
        if not file_path:
            print("❌ File not found. Please provide file path or ensure ComfyUI_0001_.png exists in current directory")
            return False
    
    # Check if file exists
    if not os.path.exists(file_path):
        print(f"❌ File not found: {file_path}")
        return False
    
    # Get file info
    file_size = os.path.getsize(file_path)
    file_size_mb = file_size / (1024 * 1024)
    filename = os.path.basename(file_path)
    
    print("=" * 60)
    print("🧪 S3 Optimized Upload Test (Direct + Parallel Multipart)")
    print("=" * 60)
    print(f"📁 File: {file_path}")
    print(f"📊 Size: {file_size_mb:.2f} MB ({file_size:,} bytes)")
    print(f"🔌 Endpoint: {endpoint_url}")
    print(f"🪣 Bucket: {bucket_name}")
    print(f"🌍 Region: {region}")
    
    # Determine S3 key
    if not s3_key:
        s3_key = f"test-uploads/{filename}"
    
    print(f"🔑 S3 Key: {s3_key}")
    
    # Determine ContentType from extension
    file_extension = os.path.splitext(filename)[1].lower()
    mime_types = {
        '.png': 'image/png',
        '.jpg': 'image/jpeg',
        '.jpeg': 'image/jpeg',
        '.webp': 'image/webp',
        '.gif': 'image/gif',
        '.mp4': 'video/mp4',
        '.webm': 'video/webm',
    }
    content_type = mime_types.get(file_extension, 'application/octet-stream')
    print(f"📄 ContentType: {content_type}")
    
    # Show upload strategy
    if file_size > 16 * 1024 * 1024:
        num_parts = (file_size + 8 * 1024 * 1024 - 1) // (8 * 1024 * 1024)
        print(f"⚡ Strategy: Parallel multipart (10 concurrent parts, ~{num_parts} total parts)")
    else:
        print(f"⚡ Strategy: Direct upload (single request)")
    
    # Create S3 client
    print("\n🔌 Creating S3 client...")
    boto3_config = Config(signature_version='s3v4')
    s3_client = boto3.client(
        's3',
        endpoint_url=endpoint_url,
        aws_access_key_id=access_key,
        aws_secret_access_key=secret_key,
        region_name=region,
        config=boto3_config
    )
    
    # Upload file
    print(f"\n⬆️  Uploading {filename} ({file_size_mb:.2f} MB) to S3...")
    print(f"   Destination: s3://{bucket_name}/{s3_key}")
    
    start_time = time.time()
    
    try:
        optimized_upload(
            file_path=file_path,
            bucket_name=bucket_name,
            object_name=s3_key,
            s3_client=s3_client,
            content_type=content_type
        )
        
        elapsed_time = time.time() - start_time
        upload_speed_mbps = file_size_mb / elapsed_time if elapsed_time > 0 else 0
        
        print(f"\n✅ Upload completed successfully!")
        print(f"⏱️  Time: {elapsed_time:.2f} seconds")
        print(f"🚀 Speed: {upload_speed_mbps:.2f} MB/s")
        
        # Construct public URL
        public_url = os.environ.get("BUCKET_PUBLIC_URL")
        if public_url:
            public_url = public_url.rstrip("/")
            s3_url = f"{public_url}/{bucket_name}/{s3_key}"
        else:
            # Fallback: construct from endpoint URL
            parsed = urllib.parse.urlparse(endpoint_url)
            hostname = parsed.hostname or ""
            
            if 'storage.supabase.co' in hostname:
                project_ref = hostname.split('.')[0]
                s3_url = f"https://{project_ref}.supabase.co/storage/v1/object/public/{bucket_name}/{s3_key}"
            elif 'r2.cloudflarestorage.com' in hostname:
                account_id = hostname.split('.')[0]
                s3_url = f"https://{account_id}.r2.cloudflarestorage.com/{bucket_name}/{s3_key}"
            else:
                base_url = endpoint_url.rstrip("/")
                s3_url = f"{base_url}/{bucket_name}/{s3_key}"
        
        print(f"\n🔗 Public URL: {s3_url}")
        
        # Verify upload
        print("\n🔍 Verifying upload...")
        try:
            response = s3_client.head_object(Bucket=bucket_name, Key=s3_key)
            uploaded_size = response.get('ContentLength', 0)
            uploaded_size_mb = uploaded_size / (1024 * 1024)
            print(f"✅ Verified: File exists in S3 ({uploaded_size_mb:.2f} MB)")
            print(f"   Content-Type: {response.get('ContentType', 'N/A')}")
            print(f"   Last Modified: {response.get('LastModified', 'N/A')}")
        except ClientError as e:
            print(f"⚠️  Warning: Could not verify upload: {e}")
        
        return True
        
    except Exception as e:
        elapsed_time = time.time() - start_time
        print(f"\n❌ Upload failed after {elapsed_time:.2f} seconds")
        print(f"   Error: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    print("🧪 S3 Optimized Upload Test Script")
    print("   (Direct upload for small files, parallel multipart for large files)")
    print()
    
    # Get optional file path from command line
    file_path = sys.argv[1] if len(sys.argv) > 1 else None
    s3_key = sys.argv[2] if len(sys.argv) > 2 else None
    
    if file_path == "--help" or file_path == "-h":
        print(f"\nUsage: python {sys.argv[0]} [file_path] [s3_key]")
        print("\nExamples:")
        print(f"  python {sys.argv[0]}")
        print(f"  python {sys.argv[0]} ComfyUI_0001_.png")
        print(f"  python {sys.argv[0]} ComfyUI_0001_.png test/my-image.png")
        sys.exit(0)
    
    success = test_upload(file_path, s3_key)
    sys.exit(0 if success else 1)