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

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

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

logger = logging.getLogger(__name__)


class SubscriptionFetcher:
    def __init__(self, tenant_service_url: str):
        self.tenant_service_url = tenant_service_url.rstrip('/')
        logger.info("SubscriptionFetcher initialized with URL: %s", self.tenant_service_url)

    async def get_user_subscription_context(
        self,
        user_id: str,
        service_token: str
    ) -> Dict[str, Any]:
        """
        Fetch user's tenant memberships and subscription data.
        Called ONCE at login, not per-request.

        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)
            
            async with httpx.AsyncClient(timeout=10.0) as client:
                # Get user's tenant memberships - corrected URL
                memberships_url = f"{self.tenant_service_url}/api/v1/tenants/members/user/{user_id}"
                headers = {
                    "Authorization": f"Bearer {service_token}",
                    "Content-Type": "application/json"
                }

                logger.debug("Fetching user memberships from URL: %s", memberships_url)
                response = await client.get(memberships_url, headers=headers)
                
                if response.status_code != 200:
                    logger.error(f"Failed to fetch user memberships: {response.status_code}")
                    raise HTTPException(
                        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                        detail="Failed to fetch user memberships"
                    )

                memberships = response.json()
                
                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 - FIXED: Use correct endpoint
                subscription_url = f"{self.tenant_service_url}/api/v1/tenants/subscriptions/{primary_tenant_id}/active"
                subscription_response = await client.get(subscription_url, headers=headers)

                if subscription_response.status_code != 200:
                    logger.error(f"Failed to fetch subscription for tenant {primary_tenant_id}: {subscription_response.status_code}")
                    # Return with basic info but no subscription
                    return {
                        "tenant_id": primary_tenant_id,
                        "tenant_role": primary_role,
                        "subscription": None,
                        "tenant_access": memberships
                    }

                subscription_data = subscription_response.json()

                # 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 - FIXED: Use correct endpoint
                    tenant_sub_url = f"{self.tenant_service_url}/api/v1/tenants/subscriptions/{tenant_id}/active"
                    tenant_sub_response = await client.get(tenant_sub_url, headers=headers)
                    
                    tier = "starter"  # default
                    if tenant_sub_response.status_code == 200:
                        tenant_sub = tenant_sub_response.json()
                        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 httpx.HTTPError as e:
            logger.error(f"HTTP error fetching subscription data: {str(e)}", exc_info=True)
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"HTTP error fetching subscription data: {str(e)}"
            )
        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)}"
            )