app/utils/http_client.py

"""
HTTP client utilities with correlation ID propagation.
"""
import httpx
from typing import Optional, Dict, Any
from app.middleware.correlation_id import get_correlation_id, CORRELATION_ID_HEADER


class CorrelationHttpClient:
    """
    HTTP client that automatically propagates correlation IDs to downstream services.
    """
    
    @staticmethod
    def get_headers_with_correlation_id(additional_headers: Optional[Dict[str, str]] = None) -> Dict[str, str]:
        """
        Get headers with correlation ID included.
        
        Args:
            additional_headers: Optional additional headers to include
            
        Returns:
            Dict[str, str]: Headers with correlation ID
        """
        headers = {}
        
        # Add correlation ID if available
        correlation_id = get_correlation_id()
        if correlation_id:
            headers[CORRELATION_ID_HEADER] = correlation_id
        
        # Add additional headers
        if additional_headers:
            headers.update(additional_headers)
        
        return headers
    
    @staticmethod
    async def get(
        url: str,
        headers: Optional[Dict[str, str]] = None,
        **kwargs
    ) -> httpx.Response:
        """
        Make GET request with correlation ID.
        
        Args:
            url: Request URL
            headers: Optional headers
            **kwargs: Additional arguments for httpx.get
            
        Returns:
            httpx.Response: Response object
        """
        headers_with_correlation = CorrelationHttpClient.get_headers_with_correlation_id(headers)
        
        async with httpx.AsyncClient() as client:
            return await client.get(url, headers=headers_with_correlation, **kwargs)
    
    @staticmethod
    async def post(
        url: str,
        headers: Optional[Dict[str, str]] = None,
        json: Optional[Dict[str, Any]] = None,
        **kwargs
    ) -> httpx.Response:
        """
        Make POST request with correlation ID.
        
        Args:
            url: Request URL
            headers: Optional headers
            json: Optional JSON body
            **kwargs: Additional arguments for httpx.post
            
        Returns:
            httpx.Response: Response object
        """
        headers_with_correlation = CorrelationHttpClient.get_headers_with_correlation_id(headers)
        
        async with httpx.AsyncClient() as client:
            return await client.post(url, headers=headers_with_correlation, json=json, **kwargs)
    
    @staticmethod
    async def put(
        url: str,
        headers: Optional[Dict[str, str]] = None,
        json: Optional[Dict[str, Any]] = None,
        **kwargs
    ) -> httpx.Response:
        """
        Make PUT request with correlation ID.
        
        Args:
            url: Request URL
            headers: Optional headers
            json: Optional JSON body
            **kwargs: Additional arguments for httpx.put
            
        Returns:
            httpx.Response: Response object
        """
        headers_with_correlation = CorrelationHttpClient.get_headers_with_correlation_id(headers)
        
        async with httpx.AsyncClient() as client:
            return await client.put(url, headers=headers_with_correlation, json=json, **kwargs)
    
    @staticmethod
    async def delete(
        url: str,
        headers: Optional[Dict[str, str]] = None,
        **kwargs
    ) -> httpx.Response:
        """
        Make DELETE request with correlation ID.
        
        Args:
            url: Request URL
            headers: Optional headers
            **kwargs: Additional arguments for httpx.delete
            
        Returns:
            httpx.Response: Response object
        """
        headers_with_correlation = CorrelationHttpClient.get_headers_with_correlation_id(headers)
        
        async with httpx.AsyncClient() as client:
            return await client.delete(url, headers=headers_with_correlation, **kwargs)