From e3c8ee93e730f68536c78553183512739c0d1f8a Mon Sep 17 00:00:00 2001
From: seolmin <seolmin@megazone.com>
Date: Mon, 2 Sep 2024 17:27:17 +0900
Subject: [PATCH] fix: add type mapper in cost manager

---
 .../cost_analysis/manager/cost_manager.py     | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/src/cloudforet/cost_analysis/manager/cost_manager.py b/src/cloudforet/cost_analysis/manager/cost_manager.py
index 82dc120..515b6a2 100644
--- a/src/cloudforet/cost_analysis/manager/cost_manager.py
+++ b/src/cloudforet/cost_analysis/manager/cost_manager.py
@@ -17,6 +17,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.default_vars = None
         self.field_mapper = None
+        self.type_mapper = None
 
     def get_data(self, options, secret_data, schema, task_options):
         if "default_vars" in options:
@@ -25,6 +26,9 @@ def get_data(self, options, secret_data, schema, task_options):
         if "field_mapper" in options:
             self.field_mapper = options["field_mapper"]
 
+        if "type_mapper" in options:
+            self.type_mapper = options["type_mapper"]
+
         if "base_url" in task_options:
             base_url = task_options["base_url"]
             http_file_connector = self.locator.get_connector(HTTPFileConnector)
@@ -53,6 +57,9 @@ def _make_cost_data(self, results):
             if self.field_mapper:
                 self._set_default_vars(result)
 
+            if self.type_mapper:
+                self._set_type_mapper(result)
+
             self._create_billed_date(result)
 
             if not self._convert_cost_and_usage_quantity_types(result):
@@ -199,3 +206,16 @@ def _check_required_fields(result):
         for field in _REQUIRED_FIELDS:
             if field not in result:
                 raise ERROR_REQUIRED_PARAMETER(key=field)
+
+    def _set_type_mapper(self, result):
+        # Not Implemented
+        if "additional_info" in self.type_mapper:
+            if (
+                "Account ID" in self.type_mapper["additional_info"]
+                and "additional_info" in result
+                and "Account ID" in result.get("additional_info", {})
+            ):
+                account_id = result["additional_info"]["Account ID"]
+                if isinstance(account_id, int) or isinstance(account_id, float):
+                    result["additional_info"]["Account ID"] = str(account_id).zfill(12)
+        return result