From 52ff3e0f0bc75a160e2644174d25d6f63b93566d Mon Sep 17 00:00:00 2001 From: lhhyung Date: Tue, 2 Sep 2025 13:37:50 +0900 Subject: [PATCH] feat: Add customer_tenant_id support to get_benefit_data --- .../connector/azure_cost_mgmt_connector.py | 6 ++- .../cost_analysis/manager/cost_manager.py | 37 +++++++++++++--- .../cost_analysis/manager/job_manager.py | 44 ++++++++++++++----- 3 files changed, 69 insertions(+), 18 deletions(-) diff --git a/src/cloudforet/cost_analysis/connector/azure_cost_mgmt_connector.py b/src/cloudforet/cost_analysis/connector/azure_cost_mgmt_connector.py index 3750ee2..8d40e74 100644 --- a/src/cloudforet/cost_analysis/connector/azure_cost_mgmt_connector.py +++ b/src/cloudforet/cost_analysis/connector/azure_cost_mgmt_connector.py @@ -151,8 +151,9 @@ def query_usage_http( secret_data: dict, start: datetime, end: datetime, + collect_scope: str, account_agreement_type: str, - options=None, + tenant_id: str, ): try: billing_account_id = secret_data["billing_account_id"] @@ -173,6 +174,9 @@ def query_usage_http( parameters["dataset"]["grouping"] = ( BENEFIT_GROUPING + BENEFIT_GROUPING_MPA ) + if collect_scope == "customer_tenant_id": + self.next_link = f"https://management.azure.com/providers/Microsoft.Billing/billingAccounts/{billing_account_id}/customers/{tenant_id}/providers/Microsoft.CostManagement/query?api-version={api_version}" + elif account_agreement_type == "EnterpriseAgreement": parameters["dataset"]["grouping"] = ( BENEFIT_GROUPING + BENEFIT_GROUPING_EA diff --git a/src/cloudforet/cost_analysis/manager/cost_manager.py b/src/cloudforet/cost_analysis/manager/cost_manager.py index 4c48ec9..2c0f643 100644 --- a/src/cloudforet/cost_analysis/manager/cost_manager.py +++ b/src/cloudforet/cost_analysis/manager/cost_manager.py @@ -381,6 +381,9 @@ def get_benefit_data( domain_id: str, ): self.azure_cm_connector.create_session(options, secret_data, schema) + + collect_scope: str = task_options["collect_scope"] + tenant_ids: list = self._get_tenant_ids(task_options, collect_scope) account_agreement_type = task_options.get("account_agreement_type") start: datetime = self._get_first_date_of_month(task_options["start"]) end: datetime = datetime.utcnow() @@ -390,17 +393,37 @@ def get_benefit_data( for time_period in monthly_time_period: _start = time_period["start"] _end = time_period["end"] - response_stream = self.azure_cm_connector.query_usage_http( - secret_data, _start, _end, account_agreement_type + + start_time = time.time() + _LOGGER.info( + f"[get_benefit_data] {tenant_ids} start to collect data from {_start} to {_end}" ) + for idx, tenant_id in enumerate(tenant_ids): + _LOGGER.info( + f"[get_benefit_data] #{idx + 1} {tenant_id} tenant start to collect data from {_start} to {_end}, domain_id: {domain_id}" + ) - for results in response_stream: - yield self._make_benefit_cost_data( - results=results, - end=_end, - options=options, + response_stream = self.azure_cm_connector.query_usage_http( + secret_data, + _start, + _end, + collect_scope, + account_agreement_type, + tenant_id, ) + for results in response_stream: + yield self._make_benefit_cost_data( + results=results, + end=_end, + options=options, + ) + + end_time = time.time() + _LOGGER.info( + f"[get_benefit_data] all collect is done in {int(end_time - start_time)} seconds" + ) + def _make_benefit_cost_data( self, results: dict, diff --git a/src/cloudforet/cost_analysis/manager/job_manager.py b/src/cloudforet/cost_analysis/manager/job_manager.py index 6c968b7..8e79ac5 100644 --- a/src/cloudforet/cost_analysis/manager/job_manager.py +++ b/src/cloudforet/cost_analysis/manager/job_manager.py @@ -171,17 +171,41 @@ def get_tasks( # Benefit Job Task if options.get("cost_metric") == "AmortizedCost": - tasks.append( - { - "task_options": { - "start": start_month, - "account_agreement_type": billing_account_agreement_type, - "collect_scope": "billing_account_id", - "billing_tenant_id": secret_data["tenant_id"], - "is_benefit_job": True, + if ( + billing_account_agreement_type == "MicrosoftPartnerAgreement" + and options.get("collect_scope") == "customer_tenant_id" + ): + customer_tenants, first_sync_tenants = self._get_customer_tenants( + secret_data, linked_accounts + ) + divided_customer_tenants = self._get_divided_customer_tenants( + customer_tenants + ) + for divided_customer_tenant_info in divided_customer_tenants: + tasks.append( + { + "task_options": { + "start": start_month, + "account_agreement_type": billing_account_agreement_type, + "collect_scope": "customer_tenant_id", + "customer_tenants": divided_customer_tenant_info, + "billing_tenant_id": secret_data["tenant_id"], + "is_benefit_job": True, + } + } + ) + else: + tasks.append( + { + "task_options": { + "start": start_month, + "account_agreement_type": billing_account_agreement_type, + "collect_scope": "billing_account_id", + "billing_tenant_id": secret_data["tenant_id"], + "is_benefit_job": True, + } } - } - ) + ) elif secret_type == "USE_SERVICE_ACCOUNT_SECRET": task = self._get_task_scope_subscription(secret_data, options, start_month)