Merge pull request #11 from ImMin5/master

Add task_type option for job tasks

Author: ImMin5

src/plugin/main.py

@@ -7,9 +7,9 @@
 app = DataSourcePluginServer()
 
 
-@app.route('DataSource.init')
+@app.route("DataSource.init")
 def data_source_init(params: dict) -> dict:
-    """ init plugin by options
+    """init plugin by options
 
     Args:
         params (DataSourceInitRequest): {
@@ -22,15 +22,15 @@ def data_source_init(params: dict) -> dict:
             'metadata': 'dict'
         }
     """
-    options = params['options']
+    options = params["options"]
 
     data_source_mgr = DataSourceManager()
     return data_source_mgr.init_response(options)
 
 
-@app.route('DataSource.verify')
+@app.route("DataSource.verify")
 def data_source_verify(params: dict) -> None:
-    """ Verifying data source plugin
+    """Verifying data source plugin
 
     Args:
         params (CollectorVerifyRequest): {
@@ -44,18 +44,18 @@ def data_source_verify(params: dict) -> None:
         None
     """
 
-    options = params['options']
-    secret_data = params['secret_data']
-    domain_id = params.get('domain_id')
-    schema = params.get('schema')
+    options = params["options"]
+    secret_data = params["secret_data"]
+    domain_id = params.get("domain_id")
+    schema = params.get("schema")
 
     data_source_mgr = DataSourceManager()
     data_source_mgr.verify_plugin(options, secret_data, domain_id, schema)
 
 
-@app.route('Job.get_tasks')
+@app.route("Job.get_tasks")
 def job_get_tasks(params: dict) -> dict:
-    """ Get job tasks
+    """Get job tasks
 
     Args:
         params (JobGetTaskRequest): {
@@ -75,20 +75,29 @@ def job_get_tasks(params: dict) -> dict:
 
     """
 
-    domain_id = params['domain_id']
-    options = params['options']
-    secret_data = params['secret_data']
-    schema = params.get('schema')
-    start = params.get('start')
-    last_synchronized_at = params.get('last_synchronized_at')
+    domain_id = params["domain_id"]
+    options = params["options"]
+    secret_data = params["secret_data"]
+    schema = params.get("schema")
+    start = params.get("start")
+    last_synchronized_at = params.get("last_synchronized_at")
 
     job_mgr = JobManager()
-    return job_mgr.get_tasks(domain_id, options, secret_data, schema, start, last_synchronized_at)
+    task_type = options.get("task_type", "identity")
 
+    if task_type == "identity":
+        return job_mgr.get_tasks(
+            domain_id, options, secret_data, schema, start, last_synchronized_at
+        )
+    else:
+        return job_mgr.get_tasks_directory_type(
+            domain_id, options, secret_data, schema, start, last_synchronized_at
+        )
 
-@app.route('Cost.get_data')
+
+@app.route("Cost.get_data")
 def cost_get_data(params: dict) -> Generator[dict, None, None]:
-    """ Get external cost data
+    """Get external cost data
 
     Args:
         params (CostGetDataRequest): {
@@ -117,11 +126,11 @@ def cost_get_data(params: dict) -> Generator[dict, None, None]:
         }
     """
 
-    options = params['options']
-    secret_data = params['secret_data']
+    options = params["options"]
+    secret_data = params["secret_data"]
 
-    task_options = params.get('task_options', {})
-    schema = params.get('schema')
+    task_options = params.get("task_options", {})
+    schema = params.get("schema")
 
     cost_mgr = CostManager()
     return cost_mgr.get_data(options, secret_data, task_options, schema)

src/plugin/manager/cost_manager.py

@@ -12,38 +12,38 @@
 _LOGGER = logging.getLogger(__name__)
 
 _REGION_MAP = {
-    'APE1': 'ap-east-1',
-    'APN1': 'ap-northeast-1',
-    'APN2': 'ap-northeast-2',
-    'APN3': 'ap-northeast-3',
-    'APS1': 'ap-southeast-1',
-    'APS2': 'ap-southeast-2',
-    'APS3': 'ap-south-1',
-    'CAN1': 'ca-central-1',
-    'CPT': 'af-south-1',
-    'EUN1': 'eu-north-1',
-    'EUC1': 'eu-central-1',
-    'EU': 'eu-west-1',
-    'EUW2': 'eu-west-2',
-    'EUW3': 'eu-west-3',
-    'MES1': 'me-south-1',
-    'SAE1': 'sa-east-1',
-    'UGW1': 'AWS GovCloud (US-West)',
-    'UGE1': 'AWS GovCloud (US-East)',
-    'USE1': 'us-east-1',
-    'USE2': 'us-east-2',
-    'USW1': 'us-west-1',
-    'USW2': 'us-west-2',
-    'AP': 'Asia Pacific',
-    'AU': 'Australia',
-    'CA': 'Canada',
+    "APE1": "ap-east-1",
+    "APN1": "ap-northeast-1",
+    "APN2": "ap-northeast-2",
+    "APN3": "ap-northeast-3",
+    "APS1": "ap-southeast-1",
+    "APS2": "ap-southeast-2",
+    "APS3": "ap-south-1",
+    "CAN1": "ca-central-1",
+    "CPT": "af-south-1",
+    "EUN1": "eu-north-1",
+    "EUC1": "eu-central-1",
+    "EU": "eu-west-1",
+    "EUW2": "eu-west-2",
+    "EUW3": "eu-west-3",
+    "MES1": "me-south-1",
+    "SAE1": "sa-east-1",
+    "UGW1": "AWS GovCloud (US-West)",
+    "UGE1": "AWS GovCloud (US-East)",
+    "USE1": "us-east-1",
+    "USE2": "us-east-2",
+    "USW1": "us-west-1",
+    "USW2": "us-west-2",
+    "AP": "Asia Pacific",
+    "AU": "Australia",
+    "CA": "Canada",
     # 'EU': 'Europe and Israel',
-    'IN': 'India',
-    'JP': 'Japan',
-    'ME': 'Middle East',
-    'SA': 'South America',
-    'US': 'United States',
-    'ZA': 'South Africa',
+    "IN": "India",
+    "JP": "Japan",
+    "ME": "Middle East",
+    "SA": "South America",
+    "US": "United States",
+    "ZA": "South Africa",
 }
 
 
@@ -54,41 +54,47 @@ def __init__(self, *args, **kwargs):
         self.aws_s3_connector = AWSS3Connector()
         self.space_connector = SpaceONEConnector()
 
-    def get_data(self, options: dict, secret_data: dict, task_options: dict, schema: str = None) \
-            -> Generator[dict, None, None]:
+    def get_data(
+        self, options: dict, secret_data: dict, task_options: dict, schema: str = None
+    ) -> Generator[dict, None, None]:
+        task_type = task_options.get("task_type", "identity")
+
         self.aws_s3_connector.create_session(options, secret_data, schema)
         self._check_task_options(task_options)
 
-        start = task_options['start']
-        account_id = task_options['account_id']
-        service_account_id = task_options['service_account_id']
-        database = task_options['database']
-        is_sync = task_options['is_sync']
+        start = task_options["start"]
+        account_id = task_options["account_id"]
+        database = task_options["database"]
 
-        if is_sync == 'false':
-            self._update_sync_state(options, secret_data, schema, service_account_id)
+        if task_type == "identity":
+            service_account_id = task_options["service_account_id"]
+            is_sync = task_options["is_sync"]
+            if is_sync == "false":
+                self._update_sync_state(
+                    options, secret_data, schema, service_account_id
+                )
 
         date_ranges = self._get_date_range(start)
 
         for date in date_ranges:
-            year, month = date.split('-')
-            path = f'SPACE_ONE/billing/database={database}/account_id={account_id}/year={year}/month={month}'
+            year, month = date.split("-")
+            path = f"SPACE_ONE/billing/database={database}/account_id={account_id}/year={year}/month={month}"
             response = self.aws_s3_connector.list_objects(path)
-            contents = response.get('Contents', [])
+            contents = response.get("Contents", [])
             for content in contents:
-                response_stream = self.aws_s3_connector.get_cost_data(content['Key'])
+                response_stream = self.aws_s3_connector.get_cost_data(content["Key"])
                 for results in response_stream:
                     yield self._make_cost_data(results, account_id)
 
-        yield {
-            'results': []
-        }
+        yield {"results": []}
 
     def _update_sync_state(self, options, secret_data, schema, service_account_id):
         self.space_connector.init_client(options, secret_data, schema)
-        service_account_info = self.space_connector.get_service_account(service_account_id)
-        tags = service_account_info.get('tags', {})
-        tags['is_sync'] = 'true'
+        service_account_info = self.space_connector.get_service_account(
+            service_account_id
+        )
+        tags = service_account_info.get("tags", {})
+        tags["is_sync"] = "true"
         self.space_connector.update_service_account(service_account_id, tags)
 
     def _make_cost_data(self, results, account_id):
@@ -113,95 +119,97 @@ class CostSummaryItem(BaseModel):
 
         for result in results:
             try:
-                region = result['region'] or 'USE1'
-                service_code = result['service_code']
-                usage_type = result['usage_type']
+                region = result["region"] or "USE1"
+                service_code = result["service_code"]
+                usage_type = result["usage_type"]
                 data = {
-                    'cost': result.get('usage_cost', 0.0) or 0.0,
-                    'usage_quantity': result['usage_quantity'],
-                    'usage_unit': None,
-                    'provider': 'aws',
-                    'region_code': _REGION_MAP.get(region, region),
-                    'product': service_code,
-                    'usage_type': usage_type,
-                    'billed_date': result['usage_date'],
-                    'additional_info': {
-                        'Instance Type': result['instance_type'],
-                        'Account ID': account_id
+                    "cost": result.get("usage_cost", 0.0) or 0.0,
+                    "usage_quantity": result["usage_quantity"],
+                    "usage_unit": None,
+                    "provider": "aws",
+                    "region_code": _REGION_MAP.get(region, region),
+                    "product": service_code,
+                    "usage_type": usage_type,
+                    "billed_date": result["usage_date"],
+                    "additional_info": {
+                        "Instance Type": result["instance_type"],
+                        "Account ID": account_id,
                     },
-                    'tags': self._get_tags_from_cost_data(result)
+                    "tags": self._get_tags_from_cost_data(result),
                 }
 
-                if service_code == 'AWSDataTransfer':
-                    data['usage_unit'] = 'Bytes'
-                    if usage_type.find('-In-Bytes') > 0:
-                        data['additional_info']['Usage Type Details'] = 'Transfer In'
-                    elif usage_type.find('-Out-Bytes') > 0:
-                        data['additional_info']['Usage Type Details'] = 'Transfer Out'
+                if service_code == "AWSDataTransfer":
+                    data["usage_unit"] = "Bytes"
+                    if usage_type.find("-In-Bytes") > 0:
+                        data["additional_info"]["Usage Type Details"] = "Transfer In"
+                    elif usage_type.find("-Out-Bytes") > 0:
+                        data["additional_info"]["Usage Type Details"] = "Transfer Out"
                     else:
-                        data['additional_info']['Usage Type Details'] = 'Transfer Etc'
-                elif service_code == 'AmazonCloudFront':
-                    if usage_type.find('-HTTPS') > 0:
-                        data['usage_unit'] = 'Count'
-                        data['additional_info']['Usage Type Details'] = 'HTTPS Requests'
-                    elif usage_type.find('-Out-Bytes') > 0:
-                        data['usage_unit'] = 'GB'
-                        data['additional_info']['Usage Type Details'] = 'Transfer Out'
+                        data["additional_info"]["Usage Type Details"] = "Transfer Etc"
+                elif service_code == "AmazonCloudFront":
+                    if usage_type.find("-HTTPS") > 0:
+                        data["usage_unit"] = "Count"
+                        data["additional_info"]["Usage Type Details"] = "HTTPS Requests"
+                    elif usage_type.find("-Out-Bytes") > 0:
+                        data["usage_unit"] = "GB"
+                        data["additional_info"]["Usage Type Details"] = "Transfer Out"
                     else:
-                        data['usage_unit'] = 'Count'
-                        data['additional_info']['Usage Type Details'] = 'HTTP Requests'
+                        data["usage_unit"] = "Count"
+                        data["additional_info"]["Usage Type Details"] = "HTTP Requests"
                 else:
-                    data['additional_info']['Usage Type Details'] = None
+                    data["additional_info"]["Usage Type Details"] = None
 
             except Exception as e:
-                _LOGGER.error(f'[_make_cost_data] make data error: {e}', exc_info=True)
+                _LOGGER.error(f"[_make_cost_data] make data error: {e}", exc_info=True)
                 raise e
 
             costs_data.append(data)
 
-        return {
-            'results': costs_data
-        }
+        return {"results": costs_data}
 
     @staticmethod
     def _get_tags_from_cost_data(cost_data: dict) -> dict:
         tags = {}
 
-        if tags_str := cost_data.get('tags'):
+        if tags_str := cost_data.get("tags"):
             try:
                 tags_dict: dict = utils.load_json(tags_str)
                 for key, value in tags_dict.items():
-                    key = key.replace('user:', '')
+                    key = key.replace("user:", "")
                     tags[key] = value
             except Exception as e:
                 _LOGGER.debug(e)
 
         return tags
 
     @staticmethod
-    def _check_task_options(task_options):
-        if 'start' not in task_options:
-            raise ERROR_REQUIRED_PARAMETER(key='task_options.start')
+    def _check_task_options(task_options: dict):
+        task_type = task_options.get("task_type", "identity")
+
+        if "start" not in task_options:
+            raise ERROR_REQUIRED_PARAMETER(key="task_options.start")
+
+        if "database" not in task_options:
+            raise ERROR_REQUIRED_PARAMETER(key="task_options.database")
 
-        if 'account_id' not in task_options:
-            raise ERROR_REQUIRED_PARAMETER(key='task_options.account_id')
+        if "is_sync" not in task_options:
+            raise ERROR_REQUIRED_PARAMETER(key="task_options.is_sync")
 
-        if 'service_account_id' not in task_options:
-            raise ERROR_REQUIRED_PARAMETER(key='task_options.service_account_id')
+        if task_type == "identity":
 
-        if 'database' not in task_options:
-            raise ERROR_REQUIRED_PARAMETER(key='task_options.database')
+            if "account_id" not in task_options:
+                raise ERROR_REQUIRED_PARAMETER(key="task_options.account_id")
 
-        if 'is_sync' not in task_options:
-            raise ERROR_REQUIRED_PARAMETER(key='task_options.is_sync')
+            if "service_account_id" not in task_options:
+                raise ERROR_REQUIRED_PARAMETER(key="task_options.service_account_id")
 
     @staticmethod
     def _get_date_range(start):
         date_ranges = []
-        start_time = datetime.strptime(start, '%Y-%m')
+        start_time = datetime.strptime(start, "%Y-%m")
         now = datetime.utcnow()
         for dt in rrule.rrule(rrule.MONTHLY, dtstart=start_time, until=now):
-            billed_month = dt.strftime('%Y-%m')
+            billed_month = dt.strftime("%Y-%m")
             date_ranges.append(billed_month)
 
         return date_ranges