Merge branch 'feature/sc-124815-use-requests-session' into release/1.1.3

Author: alexbourret

python-lib/rest_api_client.py

@@ -17,7 +17,7 @@ class RestAPIClientError(ValueError):
 
 class RestAPIClient(object):
 
-    def __init__(self, credential, endpoint, custom_key_values={}):
+    def __init__(self, credential, endpoint, custom_key_values={}, session=None):
         logger.info("Initialising RestAPIClient, credential={}, endpoint={}".format(logger.filter_secrets(credential), endpoint))
 
         #  presets_variables contains all variables available in templates using the {{variable_name}} notation
@@ -91,7 +91,7 @@ def __init__(self, credential, endpoint, custom_key_values={}):
             self.requests_kwargs.update({"json": get_dku_key_values(key_value_body)})
         self.metadata = {}
         self.call_number = 0
-        self.session = requests.Session()
+        self.session = session or requests.Session()
 
     def set_login(self, credential):
         login_type = credential.get("login_type", "no_auth")
@@ -132,19 +132,23 @@ def request(self, method, url, can_raise_exeption=True, **kwargs):
             raise RestAPIClientError("The api-connect plugin is stuck in a loop. Please check the pagination parameters.")
         request_start_time = time.time()
         self.time_last_request = request_start_time
+        error_message = None
         try:
             response = self.request_with_redirect_retry(method, url, **kwargs)
-            request_finish_time = time.time()
         except Exception as err:
             self.pagination.is_last_batch_empty = True
             error_message = "Error: {}".format(err)
             if can_raise_exeption:
                 raise RestAPIClientError(error_message)
-            else:
-                return {"error": error_message}
+
+        request_finish_time = time.time()
         self.set_metadata("request_duration", request_finish_time - request_start_time)
         self.set_metadata("status_code", response.status_code)
         self.set_metadata("response_headers", "{}".format(response.headers))
+
+        if error_message:
+            return {"error": error_message}
+
         if response.status_code >= 400:
             error_message = "Error {}: {}".format(response.status_code, response.content)
             self.pagination.is_last_batch_empty = True

python-lib/rest_api_recipe_session.py

@@ -5,6 +5,8 @@
 from dku_constants import DKUConstants
 import copy
 import json
+import requests
+
 
 logger = SafeLogger("api-connect plugin", forbidden_keys=DKUConstants.FORBIDDEN_KEYS)
 
@@ -38,6 +40,7 @@ def get_column_to_parameter_dict(parameter_columns, parameter_renamings):
     def process_dataframe(self, input_parameters_dataframe, is_raw_output):
         results = []
         time_last_request = None
+        session = requests.Session()
         for index, input_parameters_row in input_parameters_dataframe.iterrows():
             rows_count = 0
             self.initial_parameter_columns = {}
@@ -52,7 +55,7 @@ def process_dataframe(self, input_parameters_dataframe, is_raw_output):
                 updated_endpoint_parameters,
                 self.custom_key_values
             ))
-            self.client = RestAPIClient(self.credential_parameters, updated_endpoint_parameters, custom_key_values=self.custom_key_values)
+            self.client = RestAPIClient(self.credential_parameters, updated_endpoint_parameters, custom_key_values=self.custom_key_values, session=session)
             self.client.time_last_request = time_last_request
             while self.client.has_more_data():
                 page_results = self.retrieve_next_page(is_raw_output)
@@ -76,7 +79,7 @@ def retrieve_next_page(self, is_raw_output):
                 if self.can_raise:
                     raise DataikuException(error_message)
                 else:
-                    return [{"error": error_message}]
+                    return self.format_page_rows([{"error": error_message}], is_raw_output, metadata)
             page_rows.extend(self.format_page_rows(data_rows, is_raw_output, metadata))
         else:
             # Todo: check api_response key is free and add something overwise