refactor: modify JobTasks for Google Cloud Storage

Author: seolmin

src/cloudforet/cost_analysis/connector/google_storage_collector.py

@@ -32,10 +32,9 @@ def __init__(self, *args, **kwargs):
             project=self.secret_data["project_id"], credentials=self.credentials
         )
 
-    def get_cost_data(self, base_url):
-        _LOGGER.debug(f"[get_cost_data] base url: {base_url}")
+    def get_cost_data(self, bucket_name: str):
+        _LOGGER.debug(f"[get_cost_data] bucket name: {bucket_name}")
 
-        bucket_name = self.secret_data.get("bucket_name")
         bucket = self.client.get_bucket(bucket_name)
         blob_names = [blob.name for blob in bucket.list_blobs()]
 
@@ -47,7 +46,9 @@ def get_cost_data(self, base_url):
                 csv_file_path = os.path.join(tmpdir, blob_name)
                 blob.download_to_filename(csv_file_path)
                 costs_data = self._get_csv(csv_file_path)
-                _LOGGER.debug(f"[get_cost_data] costs count: {len(costs_data)}")
+                _LOGGER.debug(
+                    f"[get_cost_data] costs count of {blob_name} : {len(costs_data)}"
+                )
 
                 # Paginate
                 page_count = int(len(costs_data) / _PAGE_SIZE) + 1

src/cloudforet/cost_analysis/connector/http_file_connector.py

@@ -3,9 +3,6 @@
 import numpy as np
 import chardet
 import requests
-import google.oauth2.service_account
-
-from spaceone.core.transaction import Transaction
 from spaceone.core.connector import BaseConnector
 from typing import List
 

src/cloudforet/cost_analysis/manager/cost_manager.py

@@ -9,34 +9,37 @@
 
 _LOGGER = logging.getLogger(__name__)
 
-_REQUIRED_FIELDS = [
-    "cost",
-    "currency",
-    "year",
-    "month",
-]
+_REQUIRED_FIELDS = ["cost", "currency", "billed_date"]
 
 
 class CostManager(BaseManager):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.http_file_connector: HTTPFileConnector = self.locator.get_connector(
-            HTTPFileConnector
-        )
+        self.default_vars = None
+        self.field_mapper = None
 
     def get_data(self, options, secret_data, schema, task_options):
-        self.http_file_connector.create_session(options, secret_data, schema)
         self._check_task_options(task_options)
 
-        base_url = task_options["base_url"]
+        if "default_vars" in options:
+            self.default_vars = options["default_vars"]
+
+        if "field_mapper" in options:
+            self.field_mapper = options["field_mapper"]
+
         if not secret_data:
-            response_stream = self.http_file_connector.get_cost_data(base_url)
+            base_url = task_options["base_url"]
+            http_file_connector = self.locator.get_connector(HTTPFileConnector)
+            http_file_connector.create_session(options, secret_data, schema)
+            response_stream = http_file_connector.get_cost_data(base_url)
         else:
             # just for Google Cloud Storage
+            bucket_name = task_options["bucket_name"]
             storage_connector = self.locator.get_connector(
                 GoogleStorageConnector, secret_data=secret_data
             )
-            response_stream = storage_connector.get_cost_data(base_url)
+            response_stream = storage_connector.get_cost_data(bucket_name)
+
         for results in response_stream:
             yield self._make_cost_data(results)
 
@@ -46,10 +49,10 @@ def _make_cost_data(self, results):
             result = self._apply_strip_to_dict_keys(result)
             result = self._apply_strip_to_dict_values(result)
 
-            if self.http_file_connector.field_mapper:
+            if self.default_vars:
                 result = self._change_result_by_field_mapper(result)
 
-            if self.http_file_connector.default_vars:
+            if self.field_mapper:
                 self._set_default_vars(result)
 
             self._create_billed_date(result)
@@ -86,8 +89,8 @@ def _make_cost_data(self, results):
 
     @staticmethod
     def _check_task_options(task_options):
-        if "base_url" not in task_options:
-            raise ERROR_REQUIRED_PARAMETER(key="task_options.base_url")
+        if "base_url" not in task_options or "bucket_name" not in task_options:
+            raise ERROR_REQUIRED_PARAMETER(key="task_options")
 
     @staticmethod
     def _apply_strip_to_dict_keys(result):
@@ -106,7 +109,7 @@ def _apply_strip_to_dict_values(result):
         return result
 
     def _change_result_by_field_mapper(self, result):
-        for origin_field, actual_field in self.http_file_connector.field_mapper.items():
+        for origin_field, actual_field in self.field_mapper.items():
             if isinstance(actual_field, str):
                 if actual_field in result:
                     result[origin_field] = result[actual_field]
@@ -124,7 +127,7 @@ def _change_result_by_field_mapper(self, result):
                     del result[actual_additional_field]
                 result[origin_field] = additional_info
 
-            return result
+        return result
 
     def _create_billed_date(self, result):
         if self._exist_billed_date(result):
@@ -170,7 +173,7 @@ def _apply_parse_date(date):
             raise e
 
     def _set_default_vars(self, result):
-        for key, value in self.http_file_connector.default_vars.items():
+        for key, value in self.default_vars.items():
             result[key] = value
 
     @staticmethod

src/cloudforet/cost_analysis/manager/job_manager.py

@@ -1,6 +1,6 @@
 import logging
-from datetime import datetime, timedelta
 
+from spaceone.core.error import ERROR_REQUIRED_PARAMETER
 from spaceone.core.manager import BaseManager
 from cloudforet.cost_analysis.model.job_model import Tasks
 from cloudforet.cost_analysis.connector.http_file_connector import HTTPFileConnector
@@ -11,21 +11,28 @@
 class JobManager(BaseManager):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.http_file_connector: HTTPFileConnector = self.locator.get_connector(
-            HTTPFileConnector
-        )
 
     def get_tasks(self, options, secret_data, schema, start, last_synchronized_at):
-        self.http_file_connector.create_session(options, secret_data, schema)
-
         tasks = []
         changed = []
+        if "base_url" in options:
+            http_file_connector = self.locator.get_connector(HTTPFileConnector)
+            http_file_connector.create_session(options, secret_data, schema)
+
+            for base_url in http_file_connector.base_url:
+                task_options = {"base_url": base_url}
+
+                tasks.append({"task_options": task_options})
+                changed.append({"start": "1900-01"})
 
-        for base_url in self.http_file_connector.base_url:
-            task_options = {"base_url": base_url}
+        elif "bucket_name" in options:
+            for bucket_name in options["bucket_name"]:
+                task_options = {"bucket_name": bucket_name}
 
-            tasks.append({"task_options": task_options})
-            changed.append({"start": "1900-01"})
+                tasks.append({"task_options": task_options})
+                changed.append({"start": "1900-01"})
+        else:
+            raise ERROR_REQUIRED_PARAMETER(key="options")
 
         _LOGGER.debug(f"[get_tasks] tasks: {tasks}")
         _LOGGER.debug(f"[get_tasks] changed: {changed}")

src/cloudforet/cost_analysis/model/job_model.py

@@ -6,7 +6,8 @@
 
 
 class TaskOptions(Model):
-    base_url = StringType(required=True)
+    base_url = StringType(default=None)
+    bucket_name = StringType(default=None)
 
 
 class Task(Model):