bakery-ia/services/auth/app/utils/subscription_fetcher.py

"""Fetches subscription data for JWT enrichment at login time"""

from typing import Dict, Any, Optional, List
import logging
from fastapi import HTTPException, status

from shared.clients.tenant_client import TenantServiceClient
from shared.config.base import BaseServiceSettings

logger = logging.getLogger(__name__)


class SubscriptionFetcher:
    def __init__(self, config: BaseServiceSettings):
        """
        Initialize SubscriptionFetcher with service configuration
        
        Args:
            config: BaseServiceSettings containing service configuration
        """
        self.tenant_client = TenantServiceClient(config)
        logger.info("SubscriptionFetcher initialized with TenantServiceClient")

    async def get_user_subscription_context(
        self,
        user_id: str
    ) -> Dict[str, Any]:
        """
        Fetch user's tenant memberships and subscription data using shared tenant client.
        Called ONCE at login, not per-request.
        
        This method uses the shared TenantServiceClient instead of direct HTTP calls,
        providing better error handling, circuit breaking, and consistency.

        Returns:
            {
                "tenant_id": "primary-tenant-uuid",
                "tenant_role": "owner",
                "subscription": {
                    "tier": "professional",
                    "status": "active",
                    "valid_until": "2025-02-15T00:00:00Z"
                },
                "tenant_access": [
                    {"id": "uuid", "role": "admin", "tier": "starter"}
                ]
            }
        """
        try:
            logger.debug("Fetching subscription data for user: %s", user_id)
            
            # Get user's tenant memberships using shared tenant client
            memberships = await self.tenant_client.get_user_memberships(user_id)
            
            if not memberships:
                logger.info(f"User {user_id} has no tenant memberships - returning default subscription context")
                return {
                    "tenant_id": None,
                    "tenant_role": None,
                    "subscription": {
                        "tier": "starter",
                        "status": "active",
                        "valid_until": None
                    },
                    "tenant_access": []
                }

            # Get primary tenant (first one, or the one with highest role)
            primary_membership = memberships[0]
            for membership in memberships:
                if membership.get("role") == "owner":
                    primary_membership = membership
                    break

            primary_tenant_id = primary_membership["tenant_id"]
            primary_role = primary_membership["role"]

            # Get subscription for primary tenant using shared tenant client
            subscription_data = await self.tenant_client.get_subscription_details(primary_tenant_id)

            if not subscription_data:
                logger.warning(f"No subscription data found for primary tenant {primary_tenant_id}")
                # Return with basic info but no subscription
                return {
                    "tenant_id": primary_tenant_id,
                    "tenant_role": primary_role,
                    "subscription": None,
                    "tenant_access": memberships
                }

            # Build tenant access list with subscription info
            tenant_access = []
            for membership in memberships:
                tenant_id = membership["tenant_id"]
                role = membership["role"]
                
                # Get subscription for each tenant using shared tenant client
                tenant_sub = await self.tenant_client.get_subscription_details(tenant_id)
                
                tier = "starter"  # default
                if tenant_sub:
                    tier = tenant_sub.get("plan", "starter")

                tenant_access.append({
                    "id": tenant_id,
                    "role": role,
                    "tier": tier
                })

            return {
                "tenant_id": primary_tenant_id,
                "tenant_role": primary_role,
                "subscription": {
                    "tier": subscription_data.get("plan", "starter"),
                    "status": subscription_data.get("status", "active"),
                    "valid_until": subscription_data.get("valid_until", None)
                },
                "tenant_access": tenant_access
            }

        except Exception as e:
            logger.error(f"Error fetching subscription data: {str(e)}", exc_info=True)
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"Error fetching subscription data: {str(e)}"
            )