app.py

#!/usr/bin/env python3
"""
Azure Storage Blob Upload Test with Pod Identity
This application tests uploading a file to Azure Storage using Kubernetes pod identity.
"""

import os
import sys
import logging
import asyncio
from datetime import datetime
from azure.identity import DefaultAzureCredential, AzureAuthorityHosts
from azure.storage.blob.aio import BlobServiceClient
from azure.core.exceptions import AzureError, ClientAuthenticationError

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class StorageUploadTester:
    """
    Test class for uploading files to Azure Storage using pod identity.
    Implements retry logic, proper error handling, and managed identity authentication.
    """
    
    def __init__(self, storage_account_name: str, container_name: str, cloud_environment: str = "AzurePublic"):
        """
        Initialize the storage upload tester.
        
        Args:
            storage_account_name: Name of the Azure Storage account
            container_name: Name of the blob container
            cloud_environment: Azure cloud environment (AzurePublic or AzureGovernment)
        """
        self.storage_account_name = storage_account_name
        self.container_name = container_name
        self.cloud_environment = cloud_environment

        if cloud_environment == "AzureGovernment":
            self.storage_url = f"https://{storage_account_name}.blob.core.usgovcloudapi.net"
            authority = AzureAuthorityHosts.AZURE_GOVERNMENT
        else:
            self.storage_url = f"https://{storage_account_name}.blob.core.windows.net"
            authority = AzureAuthorityHosts.AZURE_PUBLIC_CLOUD

        # Use DefaultAzureCredential for pod identity
        # This will automatically use the managed identity assigned to the pod
        self.credential = DefaultAzureCredential(authority=authority)
        
        # Initialize blob service client
        self.blob_service_client = BlobServiceClient(
            account_url=self.storage_url,
            credential=self.credential
        )
        
    async def create_test_file(self, filename: str = "test.txt") -> str:
        """
        Create a test file with sample content.
        
        Args:
            filename: Name of the test file to create
            
        Returns:
            Path to the created test file
        """
        try:
            test_content = f"""Test file created at: {datetime.now().isoformat()}
This is a test file for Azure Storage upload using pod identity.
Storage Account: {self.storage_account_name}
Container: {self.container_name}
Pod Identity Authentication Test
"""
            
            with open(filename, 'w') as f:
                f.write(test_content)
            
            logger.info(f"Created test file: {filename}")
            return filename
            
        except Exception as e:
            logger.error(f"Failed to create test file: {e}")
            raise

    async def test_authentication(self) -> bool:
        """
        Test authentication with Azure Storage using pod identity.
        
        Returns:
            True if authentication is successful, False otherwise
        """
        try:
            logger.info("Testing authentication with Azure Storage...")
            
            # Test authentication by listing containers
            async with self.blob_service_client:
                containers = []
                async for container in self.blob_service_client.list_containers():
                    containers.append(container.name)
                
                logger.info(f"Successfully authenticated. Found {len(containers)} containers")
                if self.container_name in containers:
                    logger.info(f"Target container '{self.container_name}' exists")
                else:
                    logger.warning(f"Target container '{self.container_name}' not found in: {containers}")
                
                return True
                
        except ClientAuthenticationError as e:
            logger.error(f"Authentication failed: {e}")
            return False
        except Exception as e:
            logger.error(f"Unexpected error during authentication test: {e}")
            return False

    async def upload_file_with_retry(self, local_file_path: str, blob_name: str, max_retries: int = 3) -> bool:
        """
        Upload file to Azure Storage with retry logic and exponential backoff.
        
        Args:
            local_file_path: Path to the local file to upload
            blob_name: Name of the blob in storage
            max_retries: Maximum number of retry attempts
            
        Returns:
            True if upload is successful, False otherwise
        """
        for attempt in range(max_retries):
            try:
                logger.info(f"Upload attempt {attempt + 1}/{max_retries} for {blob_name}")
                
                # Get blob client
                blob_client = self.blob_service_client.get_blob_client(
                    container=self.container_name,
                    blob=blob_name
                )
                
                # Upload file with overwrite enabled
                with open(local_file_path, 'rb') as data:
                    await blob_client.upload_blob(
                        data,
                        overwrite=True,
                        content_type='text/plain'
                    )
                
                logger.info(f"Successfully uploaded {blob_name} to container {self.container_name}")
                
                # Verify upload by checking blob properties
                blob_properties = await blob_client.get_blob_properties()
                logger.info(f"Blob size: {blob_properties.size} bytes")
                logger.info(f"Last modified: {blob_properties.last_modified}")
                
                return True
                
            except AzureError as e:
                logger.error(f"Azure error on attempt {attempt + 1}: {e}")
                if attempt < max_retries - 1:
                    # Exponential backoff: wait 2^attempt seconds
                    wait_time = 2 ** attempt
                    logger.info(f"Retrying in {wait_time} seconds...")
                    await asyncio.sleep(wait_time)
                else:
                    logger.error(f"Upload failed after {max_retries} attempts")
                    return False
                    
            except Exception as e:
                logger.error(f"Unexpected error on attempt {attempt + 1}: {e}")
                return False
        
        return False

    async def download_and_verify(self, blob_name: str, download_path: str) -> bool:
        """
        Download the uploaded blob and verify its content.
        
        Args:
            blob_name: Name of the blob to download
            download_path: Path where to save the downloaded file
            
        Returns:
            True if download and verification are successful, False otherwise
        """
        try:
            logger.info(f"Downloading blob {blob_name} for verification...")
            
            blob_client = self.blob_service_client.get_blob_client(
                container=self.container_name,
                blob=blob_name
            )
            
            # Download blob content
            with open(download_path, 'wb') as download_file:
                download_stream = await blob_client.download_blob()
                download_file.write(await download_stream.readall())
            
            logger.info(f"Downloaded blob to {download_path}")
            
            # Verify file exists and has content
            if os.path.exists(download_path) and os.path.getsize(download_path) > 0:
                logger.info("Download verification successful")
                return True
            else:
                logger.error("Downloaded file is empty or missing")
                return False
                
        except Exception as e:
            logger.error(f"Download verification failed: {e}")
            return False

    async def cleanup_test_files(self, *file_paths):
        """
        Clean up local test files.
        
        Args:
            *file_paths: Variable number of file paths to clean up
        """
        for file_path in file_paths:
            try:
                if os.path.exists(file_path):
                    os.remove(file_path)
                    logger.info(f"Cleaned up file: {file_path}")
            except Exception as e:
                logger.warning(f"Failed to clean up {file_path}: {e}")

    async def run_upload_test(self) -> bool:
        """
        Run the complete upload test workflow.
        
        Returns:
            True if all tests pass, False otherwise
        """
        test_filename = "test.txt"
        download_filename = "downloaded_test.txt"
        blob_name = f"test-uploads/{test_filename}"
        
        try:
            # Step 1: Test authentication
            logger.info("=== Starting Azure Storage Upload Test ===")
            if not await self.test_authentication():
                logger.error("Authentication test failed")
                return False
            
            # Step 2: Create test file
            await self.create_test_file(test_filename)
            
            # Step 3: Upload file
            logger.info(f"Uploading {test_filename} to blob {blob_name}...")
            upload_success = await self.upload_file_with_retry(test_filename, blob_name)
            
            if not upload_success:
                logger.error("Upload test failed")
                return False
            
            # Step 4: Download and verify
            verify_success = await self.download_and_verify(blob_name, download_filename)
            
            if not verify_success:
                logger.error("Download verification failed")
                return False
            
            # Step 5: Cleanup
            await self.cleanup_test_files(test_filename, download_filename)
            
            logger.info("=== Upload test completed successfully! ===")
            return True
            
        except Exception as e:
            logger.error(f"Upload test failed with unexpected error: {e}")
            return False
        finally:
            # Ensure cleanup happens even if test fails
            await self.cleanup_test_files(test_filename, download_filename)
            await self.blob_service_client.close()

async def main():
    """
    Main function to run the storage upload test.
    """
    # Get configuration from environment variables
    storage_account_name = os.getenv('AZURE_STORAGE_ACCOUNT_NAME')
    container_name = os.getenv('AZURE_STORAGE_CONTAINER_NAME', 'upload-test')
    cloud_environment = os.getenv('AZURE_CLOUD_ENVIRONMENT', 'AzureGovernment')
    
    if not storage_account_name:
        logger.error("AZURE_STORAGE_ACCOUNT_NAME environment variable is required")
        sys.exit(1)
    
    logger.info(f"Configuration:")
    logger.info(f"  Storage Account: {storage_account_name}")
    logger.info(f"  Container: {container_name}")
    logger.info(f"  Cloud Environment: {cloud_environment}")
    
    # Run the upload test
    tester = StorageUploadTester(storage_account_name, container_name, cloud_environment)
    success = await tester.run_upload_test()
    
    if success:
        logger.info("All tests passed!")
        sys.exit(0)
    else:
        logger.error("Tests failed!")
        sys.exit(1)

if __name__ == "__main__":
    asyncio.run(main())