Merge pull request #18 from dataiku/feature/sc-76763-add-follow-authorization-header-option

Add "Follow authorization header" option

Author: alexbourret

custom-recipes/api-connect/recipe.json

@@ -227,6 +227,12 @@
             "type": "BOOLEAN",
             "defaultValue": false
         },
+        {
+            "name": "redirect_auth_header",
+            "label": "Redirect authorization header",
+            "type": "BOOLEAN",
+            "defaultValue": false
+        },
         {
             "name": "display_metadata",
             "label": "Display metadata",

python-connectors/api-connect_dataset/connector.json

@@ -186,6 +186,12 @@
             "type": "BOOLEAN",
             "defaultValue": false
         },
+        {
+            "name": "redirect_auth_header",
+            "label": "Redirect authorization header",
+            "type": "BOOLEAN",
+            "defaultValue": false
+        },
         {
             "name": "display_metadata",
             "label": "Display metadata",

python-lib/dku_utils.py

@@ -19,6 +19,7 @@ def get_endpoint_parameters(configuration):
         "extraction_key",
         "raw_output",
         "ignore_ssl_check",
+        "redirect_auth_header",
         "timeout",
         "requests_per_minute",
         "pagination_type",

python-lib/rest_api_client.py

@@ -1,5 +1,6 @@
 import requests
 import time
+import copy
 from pagination import Pagination
 from safe_logger import SafeLogger
 from loop_detector import LoopDetector
@@ -74,6 +75,7 @@ def __init__(self, credential, endpoint, custom_key_values={}):
             self.requests_kwargs.update({"verify": False})
         else:
             self.requests_kwargs.update({"verify": True})
+        self.redirect_auth_header = endpoint.get("redirect_auth_header", False)
         self.timeout = endpoint.get("timeout", -1)
         if self.timeout > 0:
             self.requests_kwargs.update({"timeout": self.timeout})
@@ -141,7 +143,7 @@ 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.")
         try:
             request_start_time = time.time()
-            response = requests.request(method, url, **kwargs)
+            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
@@ -168,6 +170,17 @@ def request(self, method, url, can_raise_exeption=True, **kwargs):
         self.pagination.update_next_page(json_response, response.links)
         return json_response
 
+    def request_with_redirect_retry(self, method, url, **kwargs):
+        # In case of redirection to another domain, the authorization header is not kept
+        # If redirect_auth_header is true, another attempt is made with initial headers to the redirected url
+        response = requests.request(method, url, **kwargs)
+        if self.redirect_auth_header and not response.url.startswith(url):
+            redirection_kwargs = copy.deepcopy(kwargs)
+            redirection_kwargs.pop("params", None)  # params are contained in the redirected url
+            logger.warning("Redirection ! Accessing endpoint {} with initial authorization headers".format(response.url))
+            response = requests.request(method, response.url, **redirection_kwargs)
+        return response
+
     def paginated_api_call(self, can_raise_exeption=True):
         pagination_params = self.pagination.get_params()
         params = self.requests_kwargs.get("params")

tests/python/integration/test_scenario.py

@@ -27,5 +27,9 @@ def test_run_api_connect_search_path(user_dss_clients):
     dss_scenario.run(user_dss_clients, project_key=TEST_PROJECT_KEY, scenario_id="SearchPath")
 
 
+def test_run_api_connect_redirection(user_dss_clients):
+    dss_scenario.run(user_dss_clients, project_key=TEST_PROJECT_KEY, scenario_id="REDIRECTION")
+
+
 def test_run_api_connect_check_sc_84465(user_dss_clients):
     dss_scenario.run(user_dss_clients, project_key=TEST_PROJECT_KEY, scenario_id="CHECKSC84465")