app/utils/aws_storage.py

#!/usr/bin/env python3
"""
Utility for interacting with AWS S3 storage.
"""

import json
import os
import boto3
from botocore.exceptions import ClientError
from datetime import datetime
import uuid

class AWSS3Storage:
    """AWS S3 Storage utility for saving and retrieving data."""
    
    def __init__(self, bucket_name=None, region_name=None):
        """
        Initialize the AWS S3 Storage utility.
        
        Args:
            bucket_name (str, optional): S3 bucket name. Defaults to env var AWS_S3_BUCKET_NAME.
            region_name (str, optional): AWS region. Defaults to env var AWS_REGION or 'us-east-1'.
        """
        self.bucket_name = bucket_name or os.getenv('AWS_S3_BUCKET_NAME')
        self.region_name = region_name or os.getenv('AWS_REGION', 'us-east-1')
        
        # Initialize S3 client
        self.s3 = boto3.client(
            's3',
            aws_access_key_id=os.getenv('AWS_ACCESS_KEY_ID'),
            aws_secret_access_key=os.getenv('AWS_SECRET_ACCESS_KEY'),
            region_name=self.region_name
        )
    
    def save_pitch_data(self, script_id, elevator_pitch, full_pitch, 
                       competitors_data=None, market_insights=None):
        """
        Save pitch data to S3 bucket.
        
        Args:
            script_id (str): Unique identifier for the script
            elevator_pitch (str): Short elevator pitch
            full_pitch (str): Full pitch content
            competitors_data (str, optional): Competitor analysis markdown
            market_insights (str, optional): Market insights markdown
            
        Returns:
            dict: Response data including success status and file path
        """
        if not script_id:
            script_id = f"pitch_{datetime.now().strftime('%Y%m%d')}_{uuid.uuid4().hex[:8]}"
        
        # Prepare data structure
        pitch_data = {
            "script_id": script_id,
            "created_at": datetime.now().isoformat(),
            "elevator_pitch": elevator_pitch,
            "full_pitch": full_pitch,
            "competitors_data": competitors_data,
            "market_insights": market_insights
        }
        
        # Convert to JSON
        json_data = json.dumps(pitch_data, indent=2)
        
        # Define the object key (file path in S3)
        object_key = f"pitches/{script_id}.json"
        
        try:
            # Upload to S3
            self.s3.put_object(
                Bucket=self.bucket_name,
                Key=object_key,
                Body=json_data,
                ContentType='application/json'
            )
            
            # Generate a presigned URL for temporary access (expires in 1 hour)
            presigned_url = self.s3.generate_presigned_url(
                'get_object',
                Params={'Bucket': self.bucket_name, 'Key': object_key},
                ExpiresIn=3600  # URL expires in 1 hour
            )
            
            return {
                "success": True,
                "message": "Pitch data saved successfully",
                "script_id": script_id,
                "s3_path": f"s3://{self.bucket_name}/{object_key}",
                "presigned_url": presigned_url
            }
            
        except ClientError as e:
            error_message = e.response.get('Error', {}).get('Message', str(e))
            return {
                "success": False,
                "message": f"Failed to save pitch data: {error_message}",
                "script_id": script_id
            }
    
    def get_pitch_data(self, script_id):
        """
        Retrieve pitch data from S3 bucket.
        
        Args:
            script_id (str): Unique identifier for the script
            
        Returns:
            dict: Pitch data or error information
        """
        object_key = f"pitches/{script_id}.json"
        
        try:
            # Get object from S3
            response = self.s3.get_object(
                Bucket=self.bucket_name,
                Key=object_key
            )
            
            # Read and parse JSON data
            json_data = response['Body'].read().decode('utf-8')
            pitch_data = json.loads(json_data)
            
            return {
                "success": True,
                "data": pitch_data
            }
            
        except ClientError as e:
            error_code = e.response.get('Error', {}).get('Code')
            error_message = e.response.get('Error', {}).get('Message', str(e))
            
            return {
                "success": False,
                "message": f"Failed to retrieve pitch data: {error_message}",
                "error_code": error_code,
                "script_id": script_id
            }
    
    def list_pitches(self, limit=50):
        """
        List available pitches in the S3 bucket.
        
        Args:
            limit (int, optional): Maximum number of results. Defaults to 50.
            
        Returns:
            dict: List of pitch metadata or error information
        """
        try:
            # List objects with prefix
            response = self.s3.list_objects_v2(
                Bucket=self.bucket_name,
                Prefix="pitches/",
                MaxKeys=limit
            )
            
            pitches = []
            if 'Contents' in response:
                for item in response['Contents']:
                    # Extract script_id from the key
                    key = item['Key']
                    if key.endswith('.json'):
                        script_id = key.split('/')[-1].replace('.json', '')
                        
                        # Add basic metadata
                        pitches.append({
                            "script_id": script_id,
                            "last_modified": item['LastModified'].isoformat(),
                            "size_bytes": item['Size']
                        })
            
            return {
                "success": True,
                "count": len(pitches),
                "pitches": pitches
            }
            
        except ClientError as e:
            error_message = e.response.get('Error', {}).get('Message', str(e))
            return {
                "success": False,
                "message": f"Failed to list pitches: {error_message}"
            }