worker.py

import os
from utils import Log
from models import VideoStatus
import traceback
from exceptions import TranscodingException
from tasks import (
    VideoStatusHandler,
    VideoDownloader,
    VideoTranscoder,
    S3Uploader
)

# Prevent module imports
if __name__ != "__main__":
    print("Module not meant to be imported. Please run as a script.")
    exit(1)

# Load environment variables
DEBUG = os.getenv("DEBUG", "false") == "true"
VIDEO_S3_KEY: str = os.getenv("VIDEO_S3_KEY").strip('"')
VIDEO_ID: str = VIDEO_S3_KEY.split("/", maxsplit=3)[1]

# Initialize logger
log = Log()

# Initialize video status handler
status_handler = VideoStatusHandler(debug=DEBUG, video_id=VIDEO_ID)

try:
    # Initialize tasks
    downloader = VideoDownloader(log=log)
    transcoder = VideoTranscoder(log=log)
    uploader = S3Uploader(log=log)

    """
    Start executing tasks...
    """

    # Mark video status as 'processing'
    status_handler.update_status(VideoStatus.PROCESSING)

    # Download video to local storage
    log.info("Downloading video from S3...")
    source_video_path: str = downloader.download(key=VIDEO_S3_KEY)
    log.info(f"Video downloaded to: {source_video_path}")
    
    # Delete original video from S3
    log.info("Deleting original video from S3...")
    downloader.delete_original_object(key=VIDEO_S3_KEY)

    # Transcode video to HLS format
    log.info("Transcoding video to HLS format...")

    try: video_duration: float = transcoder.transcode(source_video_path)
    except TranscodingException as e:  # Handle custom transcoding exception
        log.error(f"Failed to transcode video: {e}")

        status_handler.update_status(
            VideoStatus.FAILED,
            processing_error=e.error_code
        )
        
        exit(1)
        
    log.info("Video transcoding completed successfully")

    # Delete source video from local storage
    os.remove(source_video_path)

    # Upload segments and manifest to S3
    log.info("Uploading segment files to S3...")
    
    try:
        s3_upload_root_key: str = VIDEO_S3_KEY.rsplit("/", maxsplit=1)[0] + "/"
        
        uploader.upload_hls_dir(
            hls_dir=transcoder.OUTPUT_DIR,
            root_object_key=s3_upload_root_key
        )
    
    except Exception as upload_error:
        log.error(f"Failed to upload hls files to S3: {upload_error}")
        
        # Reflect failure in video status
        status_handler.update_status(
            VideoStatus.FAILED,
            processing_error="processing_error"
        )
        
        # Rollback any uploaded files in S3
        log.info("Rolling back uploaded files in S3...")
        uploader.rollback_uploads(
            root_object_key=s3_upload_root_key
        )
        
        exit(1)
    
    log.info("HLS files uploaded to S3 successfully")
    
    # Update video status to 'ready'
    status_handler.update_status(
        status=VideoStatus.READY,
        manifest_s3_key=s3_upload_root_key + transcoder.MANIFEST_FILENAME,
        video_duration=video_duration
    )
    
    log.info("Video processing complete. Exiting...")

# Handle unexpected exceptions
except Exception:
    log.error(f"An unexpected error occurred")
    traceback.print_exc()
    
    # Attempt to update status to 'failed' if possible
    try:
        status_handler.update_status(
            VideoStatus.FAILED,
            processing_error="processing_error"
        )
        
    except Exception as status_update_error:
        log.error(f"Failed to update video status: {status_update_error}")
    
    exit(1)