feat: modify change cost data logic (#199)

Author: ImMin5

src/spaceone/cost_analysis/manager/cost_manager.py

@@ -46,11 +46,9 @@ def _rollback(vo: Cost):
         params["billed_year"] = billed_at.strftime("%Y")
         params["billed_month"] = billed_at.strftime("%Y-%m")
 
-        params, use_account_routing = self.data_source_account_mgr.connect_cost_data(
-            params
-        )
-        if not use_account_routing:
-            params = self.data_source_rule_mgr.change_cost_data(params)
+        params, ds_account_vo = self.data_source_account_mgr.connect_cost_data(params)
+
+        params = self.data_source_rule_mgr.change_cost_data(params, ds_account_vo)
 
         cost_vo: Cost = self.cost_model.create(params)
 

src/spaceone/cost_analysis/manager/data_source_account_manager.py

@@ -83,7 +83,9 @@ def filter_data_source_accounts(self, **conditions) -> QuerySet:
     def stat_data_source_accounts(self, query: dict) -> dict:
         return self.data_source_account_model.stat(**query)
 
-    def connect_cost_data(self, cost_data: dict) -> Tuple[dict, bool]:
+    def connect_cost_data(
+        self, cost_data: dict
+    ) -> Tuple[dict, Union[DataSourceAccount, None]]:
         data_source_id = cost_data["data_source_id"]
         domain_id = cost_data["domain_id"]
 
@@ -92,6 +94,7 @@ def connect_cost_data(self, cost_data: dict) -> Tuple[dict, bool]:
 
         use_account_routing = plugin_info_metadata.get("user_account_routing", False)
 
+        ds_account_vo = None
         if use_account_routing:
             account_connect_polices: list = plugin_info_metadata.get(
                 "account_connect_polices"
@@ -107,25 +110,23 @@ def connect_cost_data(self, cost_data: dict) -> Tuple[dict, bool]:
                     operator = polices["connect_account_to_workspace"].get("operator")
 
                     if target_value:
-                        ds_account_info = self._get_data_source_account(
+                        ds_account_vo = self._get_data_source_account_vo(
                             target_key,
                             target_value,
                             data_source_id,
                             domain_id,
                             operator,
                         )
-                        if ds_account_info:
-                            cost_data.update(
-                                {"account_id": ds_account_info.get("account_id")}
-                            )
+                        if ds_account_vo:
+                            cost_data.update({"account_id": ds_account_vo.account_id})
 
-        return cost_data, use_account_routing
+        return cost_data, ds_account_vo
 
     def connect_account_by_data_source_vo(
         self,
         data_source_account_vo: DataSourceAccount,
         data_source_vo: DataSource,
-    ) -> None:
+    ) -> DataSourceAccount:
         domain_id = data_source_vo.domain_id
 
         plugin_info_metadata = data_source_vo.plugin_info.metadata
@@ -153,10 +154,11 @@ def connect_account_by_data_source_vo(
                     )
 
                     if workspace_info:
-                        self.update_data_source_account_by_vo(
+                        data_source_account_vo = self.update_data_source_account_by_vo(
                             {"workspace_id": workspace_info.get("workspace_id")},
                             data_source_account_vo,
                         )
+        return data_source_account_vo
 
     def _get_workspace(
         self, target_key: str, target_value: str, domain_id: str, operator: str = "eq"
@@ -203,14 +205,14 @@ def _get_data_source(self, data_source_id: str, domain_id: str) -> DataSource:
 
         return data_source_vo
 
-    def _get_data_source_account(
+    def _get_data_source_account_vo(
         self,
         target_key: str,
         target_value: str,
         data_source_id: str,
         domain_id: str,
         operator: str = "eq",
-    ) -> Union[dict, None]:
+    ) -> Union[DataSourceAccount, None]:
         query = {
             "filter": [
                 {"k": "domain_id", "v": domain_id, "o": "eq"},
@@ -220,8 +222,8 @@ def _get_data_source_account(
         }
 
         data_source_account_vos, total_count = self.list_data_source_accounts(query)
-        data_source_account_info = None
+        data_source_account_vo = None
         if total_count > 0:
-            data_source_account_info = data_source_account_vos[0].to_dict()
+            data_source_account_vo = data_source_account_vos[0]
 
-        return data_source_account_info
+        return data_source_account_vo

src/spaceone/cost_analysis/manager/data_source_rule_manager.py

@@ -1,9 +1,11 @@
 import logging
 import functools
 
+from mongoengine import QuerySet
 from spaceone.core import utils
 from spaceone.core.manager import BaseManager
 from spaceone.cost_analysis.manager.identity_manager import IdentityManager
+from spaceone.cost_analysis.model import DataSourceAccount
 from spaceone.cost_analysis.model.data_source_rule_model import (
     DataSourceRule,
     DataSourceRuleCondition,
@@ -86,7 +88,9 @@ def list_data_source_rules(self, query: dict):
     def stat_data_source_rules(self, query):
         return self.data_source_rule_model.stat(**query)
 
-    def change_cost_data(self, cost_data):
+    def change_cost_data(
+        self, cost_data: dict, data_source_account_vo: DataSourceAccount = None
+    ) -> dict:
         data_source_id = cost_data["data_source_id"]
         domain_id = cost_data["domain_id"]
         (
@@ -95,31 +99,49 @@ def change_cost_data(self, cost_data):
         ) = self._get_data_source_rules(data_source_id, domain_id)
 
         cost_data = self._apply_data_source_rule_to_cost_data(
-            cost_data, managed_data_source_rule_vos, domain_id
+            cost_data, managed_data_source_rule_vos, domain_id, data_source_account_vo
         )
 
         cost_data = self._apply_data_source_rule_to_cost_data(
-            cost_data, custom_data_source_rule_vos, domain_id
+            cost_data, custom_data_source_rule_vos, domain_id, data_source_account_vo
         )
 
         return cost_data
 
     def _apply_data_source_rule_to_cost_data(
-        self, cost_data, data_source_rule_vos, domain_id
+        self,
+        cost_data: dict,
+        data_source_rule_vos: QuerySet,
+        domain_id: str,
+        data_source_account_vo: DataSourceAccount = None,
     ):
         for data_source_rule_vo in data_source_rule_vos:
             is_match = self._change_cost_data_by_rule(cost_data, data_source_rule_vo)
             if is_match:
                 cost_data = self._change_cost_data_with_actions(
-                    cost_data, data_source_rule_vo.actions, domain_id
+                    cost_data,
+                    data_source_rule_vo.actions,
+                    domain_id,
+                    data_source_account_vo,
                 )
 
             if is_match and data_source_rule_vo.options.stop_processing:
                 break
 
         return cost_data
 
-    def _change_cost_data_with_actions(self, cost_data, actions, domain_id):
+    def _change_cost_data_with_actions(
+        self,
+        cost_data: dict,
+        actions: dict,
+        domain_id: str,
+        data_source_account_vo: DataSourceAccount = None,
+    ):
+        if data_source_account_vo:
+            workspace_id = data_source_account_vo.workspace_id
+        else:
+            workspace_id = None
+
         for action, value in actions.items():
             if action == "change_project" and value:
                 cost_data["project_id"] = value
@@ -130,7 +152,7 @@ def _change_cost_data_with_actions(self, cost_data, actions, domain_id):
                 target_value = utils.get_dict_value(cost_data, source)
                 if target_value:
                     project_info = self._get_project(
-                        target_key, target_value, domain_id
+                        target_key, target_value, domain_id, workspace_id
                     )
                     if project_info:
                         cost_data["workspace_id"] = project_info.get("workspace_id")
@@ -142,7 +164,7 @@ def _change_cost_data_with_actions(self, cost_data, actions, domain_id):
                 target_value = utils.get_dict_value(cost_data, source)
                 if target_value:
                     service_account_info = self._get_service_account(
-                        target_key, target_value, domain_id
+                        target_key, target_value, domain_id, workspace_id
                     )
                     if service_account_info:
                         cost_data["service_account_id"] = service_account_info[
@@ -159,7 +181,9 @@ def _change_cost_data_with_actions(self, cost_data, actions, domain_id):
 
         return cost_data
 
-    def _get_service_account(self, target_key, target_value, domain_id):
+    def _get_service_account(
+        self, target_key, target_value, domain_id: str, workspace_id: str = None
+    ):
         if (
             f"service-account:{domain_id}:{target_key}:{target_value}"
             in self._service_account_info
@@ -176,6 +200,9 @@ def _get_service_account(self, target_key, target_value, domain_id):
             "only": ["service_account_id", "project_id", "workspace_id"],
         }
 
+        if workspace_id:
+            query["filter"].append({"k": "workspace_id", "v": workspace_id, "o": "eq"})
+
         identity_mgr: IdentityManager = self.locator.get_manager("IdentityManager")
         response = identity_mgr.list_service_accounts(query, domain_id)
         results = response.get("results", [])
@@ -190,7 +217,9 @@ def _get_service_account(self, target_key, target_value, domain_id):
         ] = service_account_info
         return service_account_info
 
-    def _get_project(self, target_key, target_value, domain_id):
+    def _get_project(
+        self, target_key, target_value, domain_id: str, workspace_id: str = None
+    ):
         if f"project:{domain_id}:{target_key}:{target_value}" in self._project_info:
             return self._project_info[
                 f"project:{domain_id}:{target_key}:{target_value}"
@@ -200,6 +229,8 @@ def _get_project(self, target_key, target_value, domain_id):
             "filter": [{"k": target_key, "v": target_value, "o": "eq"}],
             "only": ["project_id"],
         }
+        if workspace_id:
+            query["filter"].append({"k": "workspace_id", "v": workspace_id, "o": "eq"})
 
         identity_mgr: IdentityManager = self.locator.get_manager("IdentityManager")
         response = identity_mgr.list_projects({"query": query}, domain_id)