Merge pull request #121 from lhhyung/master

Add customer_tenant_id support to get_benefit_data

Author: lhhyung

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

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,

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)