"""Tenant model with helper methods.""" import asyncio import logging from typing import Awaitable, Callable, Collection, Dict, List, Optional, cast from uuid import UUID from ddtrace.trace import tracer from pdp.connectors.ows_account import ( LookupCompanyBrandsResponse, LookupParentCompaniesResponse, LookupSubaccountsResponse, LookupVendorFetchFlags, LookupVendorsResponse, OwsAccountClient, ) from pdp.connectors.ows_participant import ( LookupParticipantsResponse, OwsParticipantClient, ) from pdp.connectors.redis_client import PydanticSchemaSerializer, RedisConnector from pdp.constants.constants import TenantType from pdp.constants.error import ERROR_MESSAGE_OWS_ACCOUNT_LOOKUP_DID_NOT_RETURN from pdp.fastapi.schemas.tenant import Tenant, TenantHierarchy from pdp.proxies.helpers import lookup_with_error_handling from pdp.utils.tenant_cache import ( TENANT_HIERARCHY_KEY_PREFIX, get_cache_key_from_tenant_uuid, get_cache_keys_from_tenant_uuid_list, ) logger = logging.getLogger(__name__) SUPPORTED_TENANT_TYPES: List[TenantType] = [ TenantType.TENANT_TYPE_ACCOUNT, TenantType.TENANT_TYPE_SUBACCOUNT, TenantType.TENANT_TYPE_LABEL_PARTICIPANT, TenantType.TENANT_TYPE_COMPANY_BRAND, TenantType.TENANT_TYPE_PARENT_COMPANY, ] class TenantHierarchyLookupError(Exception): """Tenant Hierarchy Lookup error.""" pass class MultiTenantProxy: """Handles multiple tenants to proxy requests to redis or relevant microservices.""" def __init__( self, tenants: List[Tenant], redis_client: RedisConnector, ows_account_client: OwsAccountClient, ows_participant_client: OwsParticipantClient, ): """Multi-Tenant Proxy constructor.""" self._tenants = self._get_unique_tenants(tenants) self._redis_client = redis_client self._ows_account_client = ows_account_client self._ows_participant_client = ows_participant_client self._tenants_uuids_by_type = self._init_tenants_by_type() self._gathered_hierarchies: Optional[Dict[UUID, TenantHierarchy]] = None @property def gathered_hierarchies(self) -> Dict[UUID, TenantHierarchy]: """Access gathered hierarchies, if non-None. Raises TenantHierarchyLookupError if accessed before ows-* lookups. """ if self._gathered_hierarchies is None: raise TenantHierarchyLookupError( "Attempted to access gathered_hierarchies before ows-* lookups." ) return self._gathered_hierarchies def _get_unique_tenants(self, tenants: List[Tenant]) -> List[Tenant]: """Return unique, consistently ordered list of tenants.""" return sorted(set(tenants)) def _init_tenants_by_type( self, ) -> Dict[TenantType, List[UUID]]: """ Build a data structure to organize Tenant UUIDs by Tenant Type. Example object: { TenantType.TENANT_TYPE_ACCOUNT: [ UUID('f39efd58-82f5-4257-859a-36e29fa69cda'), UUID('f3011786-2845-4498-a240-b9f22b2f5be0') ], TenantType.TENANT_TYPE_SUBACCOUNT: [ UUID('2c03bd86-ba0d-42a6-bcaf-a273861412a8') ] } """ # initialize the mapping with empty lists _tenants_by_type: Dict[TenantType, List[UUID]] = { tenant_type: [] for tenant_type in SUPPORTED_TENANT_TYPES } for tenant in self._tenants: if tenant.tenant_type in SUPPORTED_TENANT_TYPES: _tenants_by_type[tenant.tenant_type].append(tenant.tenant_uuid) return _tenants_by_type def get_tenant_hierarchy(self, uuid: UUID) -> Optional[TenantHierarchy]: """Get tenant hierarchy for a UUID. Only call this method after calling `gather_tenant_hierarchies`. >>> mtp = MultiTenantProxy(uuids, ...) >>> mtp.gather_tenant_hierarchies() >>> hierarchy = mtp.get_tenant_hierarchy(uuid) """ return self.gathered_hierarchies.get(uuid) async def _get_tenant_hierarchy_from_cache( self, tenant_uuids: List[UUID] ) -> Dict[UUID, Optional[TenantHierarchy]]: """Get cached TenantHierarchy objects.""" cache_keys = get_cache_keys_from_tenant_uuid_list( tenant_uuids=tenant_uuids, cache_prefix=TENANT_HIERARCHY_KEY_PREFIX ) cached_items: List[Optional[TenantHierarchy]] = await self._redis_client.mget( keys=cache_keys, serializer=PydanticSchemaSerializer(TenantHierarchy), ) if len(tenant_uuids) != len(cached_items): raise TenantHierarchyLookupError( "Redis MGET returned different number of items than requested." f" Requested {len(tenant_uuids)}, got {len(cached_items)}" ) return dict(zip(tenant_uuids, cached_items)) async def _save_tenant_hierarchy_to_cache( self, entries: Dict[UUID, TenantHierarchy], cache_prefix: str, ) -> Dict[UUID, bool]: """Save Tenants' TenantHierarchy objects to cache.""" if not entries: return {} cache_entries: Dict[str, TenantHierarchy] = {} for tenant_uuid, entry in entries.items(): k = get_cache_key_from_tenant_uuid( tenant_uuid=tenant_uuid, cache_prefix=cache_prefix, ) cache_entries[k] = entry mset_results = await self._redis_client.mset_with_pipeline( cache_entries, serializer=PydanticSchemaSerializer(TenantHierarchy) ) mset_response: Dict[UUID, bool] = {} for tenant_uuid in entries.keys(): mset_response[tenant_uuid] = mset_results.get( get_cache_key_from_tenant_uuid(tenant_uuid, cache_prefix), False ) return mset_response @tracer.wrap() async def _get_account_tenant_hierarchies_from_ows_account( self, tenant_uuids: Optional[List[UUID]] = None, ) -> Dict[UUID, TenantHierarchy]: """Lookup tenants with tenant_type account from ows-account.""" if not tenant_uuids: # gather_tenant_hierarchies() will pass along tenant_uuids. # For other callers, load the account type uuids here. tenant_uuids = self._tenants_uuids_by_type.get( TenantType.TENANT_TYPE_ACCOUNT, [] ) if not len(tenant_uuids): return {} response: LookupVendorsResponse = await lookup_with_error_handling( lookup_cb=self._ows_account_client.lookup_vendors_by_uuids( uuids=tenant_uuids, fetch_flags=[LookupVendorFetchFlags.TENANT_HIERARCHY], ), http_error_message_400_string="Bad request to ows-account for multiple account tenants.", # noqa: E501 http_error_message_500_string="tenant_hierarchy lookup failed for multiple account tenants", # noqa: E501 unhandled_exception_string="Server error: tenant_hierarchy lookup failed for multiple account tenants", # noqa: E501 lookup_error_cls=TenantHierarchyLookupError, ) if not response or not len(response.vendors): logger.warning(ERROR_MESSAGE_OWS_ACCOUNT_LOOKUP_DID_NOT_RETURN) return {} # Figure out what was missing to log warning results = list(zip(tenant_uuids, response.vendors)) missing_tenant_uuids = [str(uuid) for uuid, item in results if not item] if len(missing_tenant_uuids): # lookup_vendors_by_uuids does not return results for # vendor uuids not found in the DB logger.warning( "ows-account lookup did not return results for every requested tenant:" f" {','.join(missing_tenant_uuids)}" ) missing_company_brand_uuids = [ str(uuid) for uuid, vendor in results if vendor and not vendor.company_brand_uuid ] if len(missing_company_brand_uuids): raise TenantHierarchyLookupError( "tenant_hierarchy missing company_brand_uuid for:" f" '{','.join(missing_company_brand_uuids)}'" ) missing_parent_company_uuids = [ str(uuid) for uuid, vendor in results if vendor and not vendor.parent_company_uuid ] if len(missing_parent_company_uuids): raise TenantHierarchyLookupError( "tenant_hierarchy missing parent_company_uuid for:" f" '{','.join(missing_parent_company_uuids)}'" ) return { vendor.uuid: TenantHierarchy( company_brand=Tenant( tenant_type=TenantType.TENANT_TYPE_COMPANY_BRAND, tenant_uuid=vendor.company_brand_uuid, ), parent_company=Tenant( tenant_type=TenantType.TENANT_TYPE_PARENT_COMPANY, tenant_uuid=vendor.parent_company_uuid, ), ) for vendor in response.vendors if vendor } @tracer.wrap() async def _get_subaccount_tenant_hierarchies_from_ows_account( self, tenant_uuids: Optional[List[UUID]] = None, ) -> Dict[UUID, TenantHierarchy]: """Lookup tenants with tenant_type subaccount from ows-account.""" if not tenant_uuids: # gather_tenant_hierarchies() will pass along tenant_uuids. # For other callers, load the subaccount type uuids here. tenant_uuids = self._tenants_uuids_by_type.get( TenantType.TENANT_TYPE_SUBACCOUNT, [] ) if not len(tenant_uuids): return {} # Lookup subaccounts by uuid from ows-account client response: LookupSubaccountsResponse = await lookup_with_error_handling( lookup_cb=self._ows_account_client.lookup_subaccounts_by_uuids( uuids=tenant_uuids, fetch_flags=[LookupVendorFetchFlags.TENANT_HIERARCHY], ), http_error_message_400_string="Bad request to ows-account for multiple subaccount tenants.", # noqa: E501 http_error_message_500_string="tenant_hierarchy lookup failed for multiple subaccount tenants", # noqa: E501 unhandled_exception_string="Server error: tenant_hierarchy lookup failed for multiple subaccount tenants", # noqa: E501 lookup_error_cls=TenantHierarchyLookupError, ) if not response or not len(response.subaccounts): logger.warning(ERROR_MESSAGE_OWS_ACCOUNT_LOOKUP_DID_NOT_RETURN) return {} # Figure out what was missing to log warning results = list(zip(tenant_uuids, response.subaccounts)) missing_tenant_uuids = [str(uuid) for uuid, item in results if not item] if len(missing_tenant_uuids): # lookup_subaccounts_by_uuids does not return results for # subaccount uuids not found in the DB logger.warning( "ows-account lookup did not return results for every requested tenant:" f" {','.join(missing_tenant_uuids)}" ) missing_vendor_uuids = [ str(uuid) for uuid, subaccount in results if subaccount and not subaccount.vendor_uuid ] if len(missing_vendor_uuids): raise TenantHierarchyLookupError( "tenant_hierarchy missing vendor_uuid for:" f" '{','.join(missing_vendor_uuids)}'" ) missing_company_brand_uuids = [ str(uuid) for uuid, subaccount in results if subaccount and not subaccount.company_brand_uuid ] if len(missing_company_brand_uuids): raise TenantHierarchyLookupError( "tenant_hierarchy missing company_brand_uuid for:" f" '{','.join(missing_company_brand_uuids)}'" ) missing_parent_company_uuids = [ str(uuid) for uuid, subaccount in results if subaccount and not subaccount.parent_company_uuid ] if len(missing_parent_company_uuids): raise TenantHierarchyLookupError( "tenant_hierarchy missing parent_company_uuid for:" f" '{','.join(missing_parent_company_uuids)}'" ) tenant_hierarchies = { subaccount.uuid: TenantHierarchy( account=Tenant( tenant_type=TenantType.TENANT_TYPE_ACCOUNT, tenant_uuid=subaccount.vendor_uuid, ), company_brand=Tenant( tenant_type=TenantType.TENANT_TYPE_COMPANY_BRAND, tenant_uuid=subaccount.company_brand_uuid, ), parent_company=Tenant( tenant_type=TenantType.TENANT_TYPE_PARENT_COMPANY, tenant_uuid=subaccount.parent_company_uuid, ), ) for subaccount in response.subaccounts if subaccount } return tenant_hierarchies @tracer.wrap() async def _get_company_brand_tenant_hierarchies_from_ows_account( self, tenant_uuids: Optional[List[UUID]] = None, ) -> Dict[UUID, TenantHierarchy]: """Lookup hierarchies for tenants with tenant_type company_brand.""" if not tenant_uuids: # gather_tenant_hierarchies() will pass along tenant_uuids. # For other callers, load the company brand type uuids here. tenant_uuids = self._tenants_uuids_by_type.get( TenantType.TENANT_TYPE_COMPANY_BRAND, [] ) if not len(tenant_uuids): return {} response: LookupCompanyBrandsResponse = await lookup_with_error_handling( lookup_cb=self._ows_account_client.lookup_company_brands_by_uuids( uuids=tenant_uuids, ), http_error_message_400_string="Bad request to ows-account for multiple company_brand tenants.", # noqa: E501 http_error_message_500_string="tenant_hierarchy lookup failed for multiple company_brand tenants", # noqa: E501 unhandled_exception_string="Server error: tenant_hierarchy lookup failed for multiple company_brand tenants", # noqa: E501 lookup_error_cls=TenantHierarchyLookupError, ) if not response or not len(response.company_brands): logger.warning(ERROR_MESSAGE_OWS_ACCOUNT_LOOKUP_DID_NOT_RETURN) return {} # Figure out what was missing to log warning results = list(zip(tenant_uuids, response.company_brands)) missing_tenant_uuids = [str(uuid) for uuid, item in results if not item] if len(missing_tenant_uuids): # lookup_company_brands_by_uuids does not return results for # company_brand uuids not found in the DB logger.warning( "ows-account lookup did not return results for every requested tenant:" # noqa: E501 f" {','.join(missing_tenant_uuids)}" ) missing_parent_company_uuids = [ str(uuid) for uuid, company_brand in results if company_brand and not company_brand.parent_company_uuid ] if len(missing_parent_company_uuids): raise TenantHierarchyLookupError( "tenant_hierarchy missing parent_company_uuid for:" f" '{','.join(missing_parent_company_uuids)}'" ) return { company_brand.uuid: TenantHierarchy( company_brand=None, parent_company=Tenant( tenant_type=TenantType.TENANT_TYPE_PARENT_COMPANY, tenant_uuid=company_brand.parent_company_uuid, ), ) for company_brand in response.company_brands if company_brand } @tracer.wrap() async def _get_label_participant_tenant_hierarchies_from_ows_participant( self, tenant_uuids: Optional[List[UUID]] = None, ) -> Dict[UUID, TenantHierarchy]: """Lookup hierarchies for tenants with tenant_type label_participant.""" if not tenant_uuids: # gather_tenant_hierarchies() will pass along tenant_uuids. # For other callers, load the label participant type uuids here. tenant_uuids = self._tenants_uuids_by_type.get( TenantType.TENANT_TYPE_LABEL_PARTICIPANT, [] ) if not len(tenant_uuids): return {} # Lookup label participants by uuid from ows-participant client response: LookupParticipantsResponse = await lookup_with_error_handling( lookup_cb=self._ows_participant_client.lookup_participants_by_uuids( uuids=tenant_uuids, ), http_error_message_400_string="Bad request to ows-participant for multiple label_participant tenants.", # noqa: E501 http_error_message_500_string="tenant_hierarchy lookup failed for multiple label_participant tenants", # noqa: E501 unhandled_exception_string="tenant_hierarchy lookup failed for multiple label_participant tenants", # noqa: E501 lookup_error_cls=TenantHierarchyLookupError, ) if not response or not len(response.label_participants): logger.warning("ows-participant lookup did not return results.") return {} # Figure out what was missing to log warning results = list(zip(tenant_uuids, response.label_participants)) missing_tenant_uuids = [str(uuid) for uuid, item in results if not item] if len(missing_tenant_uuids): # lookup_participants_by_uuids does not return results for # label_participant uuids not found in the DB logger.warning( "ows-participant lookup did not return results for every requested tenant:" # noqa: E501 f" {','.join(missing_tenant_uuids)}" ) missing_company_brand_uuids = [ str(uuid) for uuid, label_participant in results if label_participant and not label_participant.company_brand_uuid ] if len(missing_company_brand_uuids): raise TenantHierarchyLookupError( "tenant_hierarchy missing company_brand_uuid for:" f" '{','.join(missing_company_brand_uuids)}'" ) # A label participant can be owned by either a subaccount or account. # Every label participant hierarchy should have at least an account. missing_vendor_uuids = [ str(uuid) for uuid, label_participant in results if label_participant and not label_participant.vendor_uuid ] if len(missing_vendor_uuids): raise TenantHierarchyLookupError( "tenant_hierarchy missing vendor_uuid for:" f" '{','.join(missing_vendor_uuids)}'" ) missing_parent_company_uuids = [ str(uuid) for uuid, label_participant in results if label_participant and not label_participant.parent_company_uuid ] if len(missing_parent_company_uuids): raise TenantHierarchyLookupError( "tenant_hierarchy missing parent_company_uuid for:" f" '{','.join(missing_parent_company_uuids)}'" ) return { label_participant.uuid: TenantHierarchy( subaccount=Tenant( tenant_type=TenantType.TENANT_TYPE_SUBACCOUNT, tenant_uuid=label_participant.subaccount_uuid, ) if label_participant.subaccount_uuid else None, account=Tenant( tenant_type=TenantType.TENANT_TYPE_ACCOUNT, tenant_uuid=label_participant.vendor_uuid, ), company_brand=Tenant( tenant_type=TenantType.TENANT_TYPE_COMPANY_BRAND, tenant_uuid=label_participant.company_brand_uuid, ), parent_company=Tenant( tenant_type=TenantType.TENANT_TYPE_PARENT_COMPANY, tenant_uuid=label_participant.parent_company_uuid, ), ) for label_participant in response.label_participants if label_participant } @tracer.wrap() async def _get_parent_company_tenant_hierarchies_from_ows_account( self, tenant_uuids: Optional[List[UUID]] = None, ) -> Dict[UUID, TenantHierarchy]: """Lookup hierarchies for tenants with tenant_type parent_company.""" if not tenant_uuids: # gather_tenant_hierarchies() will pass along tenant_uuids. # For other callers, load the parent company type uuids here. tenant_uuids = self._tenants_uuids_by_type.get( TenantType.TENANT_TYPE_PARENT_COMPANY, [] ) if not len(tenant_uuids): return {} response: LookupParentCompaniesResponse = await lookup_with_error_handling( lookup_cb=self._ows_account_client.lookup_parent_companies_by_uuids( uuids=tenant_uuids, ), http_error_message_400_string="Bad request to ows-account for multiple parent_company tenants.", # noqa: E501 http_error_message_500_string="tenant_hierarchy lookup failed for multiple parent_company tenants", # noqa: E501 unhandled_exception_string="Server error: tenant_hierarchy lookup failed for multiple parent_company tenants", # noqa: E501 lookup_error_cls=TenantHierarchyLookupError, ) if not response or not len(response.parent_companies): logger.warning(ERROR_MESSAGE_OWS_ACCOUNT_LOOKUP_DID_NOT_RETURN) return {} assert len(tenant_uuids) == len(response.parent_companies) # Figure out what was missing to log warning results = list(zip(tenant_uuids, response.parent_companies)) missing_tenant_uuids = [str(uuid) for uuid, item in results if not item] if len(missing_tenant_uuids): # lookup_parent_companies_by_uuids does not return results for # parent_company uuids not found in the DB logger.warning( "ows-account lookup did not return results for every requested tenant:" # noqa: E501 f" {','.join(missing_tenant_uuids)}" ) return { parent_company.uuid: TenantHierarchy( company_brand=None, parent_company=Tenant( tenant_type=TenantType.TENANT_TYPE_PARENT_COMPANY, tenant_uuid=parent_company.uuid, ), ) for parent_company in response.parent_companies if parent_company } async def _get_tenant_hierarchies( self, tenant_type: TenantType, tenant_hierarchy_lookup_cb: Callable[ [List[UUID]], Awaitable[Dict[UUID, TenantHierarchy]] ], ) -> Dict[UUID, TenantHierarchy]: """ Fetch the TenantHierarchy from cache or from the 'source of truth' service. Args: tenant_type: account, subaccount, etc.. tenant_hierarchy_lookup_cb: callback function to perform \ ows-* hierarchy lookup. Returns: Dictionary (tenant_uuid to TenantHierarchy) Example: { "ec1fd7e2-9c95-4e09-a037-e924e0244283": { "company_brand": { "tenant_type": "company_brand", "tenant_uuid": "d25a4cd1-e820-45f2-be5c-56edcfeb8298" } } } """ tenant_uuids = self._tenants_uuids_by_type.get(tenant_type, []) if not tenant_uuids: return {} # fetch entries from cache tenant_hierarchies = await self._get_tenant_hierarchy_from_cache( tenant_uuids=tenant_uuids ) tenant_uuids_to_lookup = [] for tenant_uuid in tenant_uuids: if not tenant_hierarchies.get(tenant_uuid): # Save UUIDs in `tenant_uuids` that were not found in the cache tenant_uuids_to_lookup.append(tenant_uuid) # MGET returns `None` for cache misses, remove the null entries. # So we can safely use `typing.cast(Dict[UUID, TenantHierarchy])` # in the return statements below. tenant_hierarchies.pop(tenant_uuid, None) if not len(tenant_uuids_to_lookup): # Found all entries in the cache. Return. return cast(Dict[UUID, TenantHierarchy], tenant_hierarchies) # Call the tenant hierarchy callback method to fetch # from ows-account, ows-participants, etc. lookup_response = await tenant_hierarchy_lookup_cb(tenant_uuids_to_lookup) # write ows-* response entries to cache. MSET will skip null entries await self._save_tenant_hierarchy_to_cache( lookup_response, cache_prefix=TENANT_HIERARCHY_KEY_PREFIX ) # Add lookup results to tenant_hierarchies # NOTE: This assumes the callback removed all `None` entries. tenant_hierarchies.update(lookup_response) # Return combined cached and lookup responses return cast(Dict[UUID, TenantHierarchy], tenant_hierarchies) @tracer.wrap() async def gather_tenant_hierarchies(self) -> Dict[UUID, TenantHierarchy]: """Lookup hierarchies for all tenant types and merge results.""" gathered_hierarchies: Collection[ Dict[UUID, TenantHierarchy] ] = await asyncio.gather( self._get_tenant_hierarchies( tenant_type=TenantType.TENANT_TYPE_ACCOUNT, tenant_hierarchy_lookup_cb=self._get_account_tenant_hierarchies_from_ows_account, # noqa ), self._get_tenant_hierarchies( tenant_type=TenantType.TENANT_TYPE_SUBACCOUNT, tenant_hierarchy_lookup_cb=self._get_subaccount_tenant_hierarchies_from_ows_account, # noqa ), self._get_tenant_hierarchies( tenant_type=TenantType.TENANT_TYPE_COMPANY_BRAND, tenant_hierarchy_lookup_cb=self._get_company_brand_tenant_hierarchies_from_ows_account, ), self._get_tenant_hierarchies( tenant_type=TenantType.TENANT_TYPE_LABEL_PARTICIPANT, tenant_hierarchy_lookup_cb=self._get_label_participant_tenant_hierarchies_from_ows_participant, # noqa ), self._get_tenant_hierarchies( tenant_type=TenantType.TENANT_TYPE_PARENT_COMPANY, tenant_hierarchy_lookup_cb=self._get_parent_company_tenant_hierarchies_from_ows_account, # noqa ), # PP-638: Add new lookup methods here. ) self._gathered_hierarchies = {} # Combine gathered tenant hierarchies together for hierarchy in gathered_hierarchies: self._gathered_hierarchies.update(hierarchy) # Callers can access the merged results using the # self.gathered_hierarchies property. return self._gathered_hierarchies