Merge pull request #1 from port-labs/PORT-3326-aws-exporter-fail-to-p…

…aginate-when-lambda-timed-out

fix close to timeout pagination bug

lambda_function/aws/resources/resource_handler.py

@@ -33,7 +33,7 @@ def handle(self):
             resources_models = self.regions_config.get(region, {}).get('resources_models', ["{}"])
             for resource_model in list(resources_models):
                 logger.info(f"List kind: {self.kind}, region: {region}, resource_model: {resource_model}")
-                self.next_token = ''
+                self.next_token = '' if self.next_token is None else self.next_token
                 while self.next_token is not None:
                     list_resources_params = {'TypeName': self.kind, 'ResourceModel': resource_model}
                     if self.next_token:
@@ -45,6 +45,7 @@ def handle(self):
                         logger.error(
                             f"Failed list kind: {self.kind}, region: {region}, resource_model: {resource_model}; {e}")
                         self.skip_delete = True
+                        self.next_token = None
                         break
 
                     self._handle_list_response(aws_cloudcontrol_client, response)
@@ -54,7 +55,7 @@ def handle(self):
                         # Lambda timeout is too close, should return checkpoint for next run
                         return self._handle_close_to_timeout(resources_models, resource_model, region)
 
-                resources_models.remove(resource_model)
+                self._cleanup_resources_models(resources_models, resource_model, region)
 
             self._cleanup_regions(region)
 
@@ -115,16 +116,32 @@ def _handle_close_to_timeout(self, resources_models, current_resource_model, reg
         if self.next_token:
             self.selector_aws['next_token'] = self.next_token
         else:
-            resources_models.remove(current_resource_model)
+            self.selector_aws.pop('next_token', None)
+            resources_models = self._cleanup_resources_models(resources_models, current_resource_model, region)
             if not resources_models:
                 self._cleanup_regions(region)
                 if not self.regions:  # Nothing left to sync
                     return {'aws_entities': self.aws_entities, 'next_resource_config': None,
                             'skip_delete': self.skip_delete}
 
+        if 'selector' not in self.resource_config:
+            self.resource_config['selector'] = {}
+        self.resource_config['selector']['aws'] = self.selector_aws
+
         return {'aws_entities': self.aws_entities, 'next_resource_config': self.resource_config,
                 'skip_delete': self.skip_delete}
 
+    def _cleanup_resources_models(self, resources_models, current_resource_model, region):
+        resources_models.remove(current_resource_model)
+        if region not in self.regions_config:
+            self.regions_config[region] = {}
+        self.regions_config[region]['resources_models'] = resources_models
+        self.selector_aws['regions_config'] = self.regions_config
+
+        return resources_models
+
     def _cleanup_regions(self, region):
         self.regions.remove(region)
         self.regions_config.pop(region, None)
+        self.selector_aws['regions'] = self.regions
+        self.selector_aws['regions_config'] = self.regions_config