lambda_function.py

import os
import json
import boto3
import logging
import requests
import time
from typing import Dict, Any, Optional
from supabase import create_client, Client
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from botocore.exceptions import ClientError
from dotenv import load_dotenv

load_dotenv()

# Configure logging with structured format
logger = logging.getLogger()
logger.setLevel(logging.INFO)

# Environment Variables
S3_BUCKET = os.getenv('S3_BUCKET', 'oriane-videos')
SUPABASE_URL = os.getenv('SUPABASE_URL')
SUPABASE_KEY = os.getenv('SUPABASE_KEY')
MAX_RETRIES = int(os.getenv('MAX_RETRIES', '3'))
DOWNLOAD_TIMEOUT = int(os.getenv('DOWNLOAD_TIMEOUT', '300'))  # 5 minutes
CHUNK_SIZE = int(os.getenv('CHUNK_SIZE', '1048576'))  # 1MB chunks

# AWS Clients
s3_client = boto3.client('s3')
cloudwatch = boto3.client('cloudwatch')

# Supabase Client
supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)

# Configure requests session with retries
session = requests.Session()
retry_strategy = Retry(
    total=MAX_RETRIES,
    backoff_factor=1,
    status_forcelist=[500, 502, 503, 504],
)
adapter = HTTPAdapter(max_retries=retry_strategy)
session.mount("http://", adapter)
session.mount("https://", adapter)

def put_metric(name: str, value: float, unit: str = 'Count') -> None:
    """Send metric to CloudWatch."""
    try:
        cloudwatch.put_metric_data(
            Namespace='InstagramCollector',
            MetricData=[
                {
                    'MetricName': name,
                    'Value': value,
                    'Unit': unit,
                    'Timestamp': time.time()
                }
            ]
        )
    except Exception as e:
        logger.error(f"Failed to put metric {name}: {str(e)}")

def get_video_url_from_supabase(shortcode: str, resolution: int) -> Optional[str]:
    """
    Fetch the video URL from Supabase using the shortcode.
    If multiple resolutions are available, select the requested resolution.
    """
    try:
        response = supabase.table("watched_content").select("video_url").eq("code", shortcode).single().execute()
        
        if not response.data or "video_url" not in response.data:
            logger.error(f"Video URL not found in Supabase for shortcode: {shortcode}")
            put_metric('SupabaseVideoNotFound', 1)
            return None

        video_url = response.data["video_url"]
        if "{resolution}" in video_url:
            video_url = video_url.replace("{resolution}", str(resolution))

        put_metric('SupabaseVideoFound', 1)
        return video_url

    except Exception as e:
        logger.error(f"Failed to fetch video URL from Supabase: {str(e)}")
        put_metric('SupabaseError', 1)
        return None

def download_video_directly(video_url: str, shortcode: str) -> Optional[str]:
    """
    Downloads the video using `requests` and saves it locally.
    Returns the local file path.
    """
    local_path = f"/tmp/{shortcode}.mp4"
    start_time = time.time()

    try:
        logger.info(f"Starting download from {video_url}")
        response = session.get(video_url, stream=True, timeout=DOWNLOAD_TIMEOUT)
        response.raise_for_status()

        total_size = int(response.headers.get('content-length', 0))
        block_size = CHUNK_SIZE
        downloaded = 0

        with open(local_path, "wb") as video_file:
            for chunk in response.iter_content(chunk_size=block_size):
                if chunk:
                    video_file.write(chunk)
                    downloaded += len(chunk)
                    if total_size:
                        progress = (downloaded / total_size) * 100
                        logger.info(f"Download progress: {progress:.2f}%")

        download_time = time.time() - start_time
        logger.info(f"Video downloaded successfully in {download_time:.2f} seconds")
        put_metric('DownloadTime', download_time, 'Seconds')
        put_metric('DownloadSize', total_size, 'Bytes')
        return local_path

    except requests.Timeout:
        logger.error(f"Download timeout for {video_url}")
        put_metric('DownloadTimeout', 1)
        return None
    except requests.RequestException as e:
        logger.error(f"Error downloading video: {str(e)}")
        put_metric('DownloadError', 1)
        return None

def upload_to_s3(local_path: str, shortcode: str) -> Optional[str]:
    """
    Uploads the downloaded video to S3 and returns the S3 URL.
    """
    s3_key = f"instagram/{shortcode}.mp4"
    start_time = time.time()

    try:
        # Get file size for progress tracking
        file_size = os.path.getsize(local_path)
        
        # Configure S3 transfer config for better performance
        config = boto3.s3.transfer.TransferConfig(
            multipart_threshold=8 * 1024 * 1024,  # 8MB
            max_concurrency=10,
            multipart_chunksize=8 * 1024 * 1024,  # 8MB
            use_threads=True
        )

        s3_client.upload_file(
            local_path, 
            S3_BUCKET, 
            s3_key,
            Config=config
        )
        
        upload_time = time.time() - start_time
        s3_url = f"s3://{S3_BUCKET}/{s3_key}"
        
        logger.info(f"Successfully uploaded video to S3 in {upload_time:.2f} seconds")
        put_metric('UploadTime', upload_time, 'Seconds')
        put_metric('UploadSize', file_size, 'Bytes')
        return s3_url

    except ClientError as e:
        logger.error(f"S3 upload error: {str(e)}")
        put_metric('S3UploadError', 1)
        return None
    except Exception as e:
        logger.error(f"Failed to upload video to S3: {str(e)}")
        put_metric('S3UploadError', 1)
        return None

def update_supabase_after_upload(shortcode: str) -> bool:
    """
    Updates Supabase to mark the video as fetched.
    Returns True if successful, False otherwise.
    """
    try:
        supabase.table("watched_content").update({"is_fetched": True}).eq("code", shortcode).execute()
        logger.info(f"Updated Supabase: is_fetched set to True for shortcode {shortcode}")
        put_metric('SupabaseUpdateSuccess', 1)
        return True

    except Exception as e:
        logger.error(f"Failed to update Supabase: {str(e)}")
        put_metric('SupabaseUpdateError', 1)
        return False

def validate_input(body: Dict[str, Any]) -> bool:
    """
    Validates the input message body.
    """
    required_fields = ["shortcode"]
    if not all(field in body for field in required_fields):
        logger.error(f"Missing required fields in message body: {required_fields}")
        put_metric('InvalidInput', 1)
        return False
    
    if not isinstance(body.get("resolution", 720), int):
        logger.error("Resolution must be an integer")
        put_metric('InvalidInput', 1)
        return False
        
    return True

def process_video(shortcode: str, resolution: int) -> Dict[str, Any]:
    """
    Fetches the video URL, downloads it, uploads to S3, and updates Supabase.
    """
    start_time = time.time()
    
    try:
        video_url = get_video_url_from_supabase(shortcode, resolution)
        if not video_url:
            return {"shortcode": shortcode, "error": "Video URL not found in Supabase."}

        local_path = download_video_directly(video_url, shortcode)
        if not local_path:
            return {"shortcode": shortcode, "error": "Failed to download video."}

        s3_url = upload_to_s3(local_path, shortcode)
        if not s3_url:
            return {"shortcode": shortcode, "error": "Failed to upload video to S3."}

        if not update_supabase_after_upload(shortcode):
            logger.warning(f"Failed to update Supabase status for {shortcode}")

        # Cleanup local file
        try:
            os.remove(local_path)
        except Exception as e:
            logger.error(f"Failed to delete local file {local_path}: {str(e)}")

        total_time = time.time() - start_time
        put_metric('TotalProcessingTime', total_time, 'Seconds')
        
        return {"shortcode": shortcode, "s3_url": s3_url}

    except Exception as e:
        logger.error(f"Error processing video {shortcode}: {str(e)}")
        put_metric('ProcessingError', 1)
        return {"shortcode": shortcode, "error": str(e)}

def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
    """
    AWS Lambda Handler: Processes Instagram video downloads.
    """
    try:
        # Validate environment variables
        if not all([SUPABASE_URL, SUPABASE_KEY]):
            raise ValueError("Missing required environment variables")

        # Validate event structure
        if not event.get("Records"):
            raise ValueError("No records found in event")

        message = event["Records"][0]
        if "body" not in message:
            raise ValueError("No body found in message")

        body = json.loads(message["body"])
        
        # Validate input data
        if not validate_input(body):
            return {
                "status": "error",
                "message": "Invalid input data"
            }

        shortcode = body["shortcode"]
        resolution = body.get("resolution", 720)

        logger.info(f"Processing video: {shortcode} at {resolution}p")
        put_metric('ProcessingStarted', 1)

        result = process_video(shortcode, resolution)

        if "error" in result:
            logger.error(f"Failed to process video: {result['error']}")
            put_metric('ProcessingFailed', 1)
            return {"status": "error", "result": result}

        logger.info(f"Successfully processed video: {result['s3_url']}")
        put_metric('ProcessingSuccess', 1)
        return {"status": "success", "result": result}

    except json.JSONDecodeError as e:
        logger.error(f"Failed to parse message body: {str(e)}")
        put_metric('JSONParseError', 1)
        return {"status": "error", "message": "Invalid JSON in message body"}
    except Exception as e:
        logger.error(f"Unexpected error: {str(e)}")
        put_metric('UnexpectedError', 1)
        return {"status": "error", "message": str(e)}