Merge pull request #12 from dataiku/bug/sc-72912-output-json-compatible-with-unnest

sc 72912 output json compatible with unnest

Author: alexbourret

custom-recipes/api-connect/recipe.json

@@ -226,6 +226,13 @@
             "type": "BOOLEAN",
             "defaultValue": false
         },
+        {
+            "name": "display_metadata",
+            "label": "Display metadata",
+            "description": "Status code, request time...",
+            "type": "BOOLEAN",
+            "defaultValue": false
+        },
         {
             "name": "timeout",
             "label": "Timeout (s)",
@@ -239,6 +246,13 @@
             "description": "-1 for no limit",
             "type": "INT",
             "defaultValue": -1
+        },
+        {
+            "name": "maximum_number_rows",
+            "label": "Maximum number of rows",
+            "description": "-1 for no limit",
+            "type": "INT",
+            "defaultValue": -1
         }
     ],
     "resourceKeys": []

custom-recipes/api-connect/recipe.py

@@ -40,6 +40,8 @@ def get_partitioning_keys(id_list, dku_flow_variables):
     raise ValueError("There is no parameter column selected.")
 parameter_renamings = get_dku_key_values(config.get("parameter_renamings", {}))
 custom_key_values = get_dku_key_values(config.get("custom_key_values", {}))
+display_metadata = config.get("display_metadata", False)
+maximum_number_rows = config.get("maximum_number_rows", -1)
 input_parameters_dataset = dataiku.Dataset(input_A_names[0])
 partitioning_keys = get_partitioning_keys(input_parameters_dataset, dku_flow_variables)
 custom_key_values.update(partitioning_keys)
@@ -51,7 +53,9 @@ def get_partitioning_keys(id_list, dku_flow_variables):
     endpoint_parameters,
     extraction_key,
     parameter_columns,
-    parameter_renamings
+    parameter_renamings,
+    display_metadata,
+    maximum_number_rows=maximum_number_rows
 )
 results = recipe_session.process_dataframe(input_parameters_dataframe, is_raw_output)
 

python-connectors/api-connect_dataset/connector.json

@@ -110,8 +110,8 @@
         },
         {
             "name": "extraction_key",
-            "label": "Key to data array (optional)",
-            "description": "",
+            "label": "Path to data array (optional)",
+            "description": "Dot separated key path",
             "defaultValue": null,
             "type": "STRING"
         },
@@ -186,6 +186,13 @@
             "type": "BOOLEAN",
             "defaultValue": false
         },
+        {
+            "name": "display_metadata",
+            "label": "Display metadata",
+            "description": "Status code, request time...",
+            "type": "BOOLEAN",
+            "defaultValue": false
+        },
         {
             "name": "timeout",
             "label": "Timeout (s)",

python-connectors/api-connect_dataset/connector.py

@@ -2,7 +2,10 @@
 from dataikuapi.utils import DataikuException
 from safe_logger import SafeLogger
 from rest_api_client import RestAPIClient
-from dku_utils import get_dku_key_values, get_endpoint_parameters
+from dku_utils import get_dku_key_values, get_endpoint_parameters, parse_keys_for_json, get_value_from_path
+from dku_constants import DKUConstants
+import json
+
 
 logger = SafeLogger("api-connect plugin", forbiden_keys=["token", "password"])
 
@@ -18,11 +21,11 @@ def __init__(self, config, plugin_config):
         custom_key_values = get_dku_key_values(config.get("custom_key_values", {}))
         self.client = RestAPIClient(credential, endpoint_parameters, custom_key_values)
         extraction_key = endpoint_parameters.get("extraction_key", None)
-        if extraction_key == '':
-            extraction_key = None
-        self.extraction_key = extraction_key
+        self.extraction_key = extraction_key or ''
+        self.extraction_path = self.extraction_key.split('.')
         self.raw_output = endpoint_parameters.get("raw_output", None)
         self.maximum_number_rows = config.get("maximum_number_rows", -1)
+        self.display_metadata = config.get("display_metadata", False)
 
     def get_read_schema(self):
         # In this example, we don't specify a schema here, so DSS will infer the schema
@@ -37,25 +40,31 @@ def generate_rows(self, dataset_schema=None, dataset_partitioning=None,
         record_count = 0
         while self.client.has_more_data():
             json_response = self.client.paginated_api_call()
-            if self.extraction_key is None:
-                # Todo: check api_response key is free and add something overwise
-                if isinstance(json_response, list):
-                    record_count += len(json_response)
-                    for row in json_response:
-                        yield {"api_response": row}
-                else:
-                    record_count += 1
-                    yield {"api_response": json_response}
+            metadata = self.client.get_metadata() if self.display_metadata else None
+            if self.extraction_key:
+                data = get_value_from_path(json_response, self.extraction_path)
             else:
-                data = json_response.get(self.extraction_key, None)
-                if data is None:
-                    raise DataikuException("Extraction key '{}' was not found in the incoming data".format(self.extraction_key))
+                data = json_response
+            if isinstance(data, list):
                 record_count += len(data)
-                for result in data:
-                    yield {"api_response": result} if self.raw_output else result
+                for row in data:
+                    yield self.format_output(row, metadata)
+            else:
+                record_count += 1
+                yield self.format_output(data, metadata)
             if is_records_limit and record_count >= records_limit:
                 break
 
+    def format_output(self, item, metadata=None):
+        output = metadata or {}
+        if self.raw_output:
+            output.update({
+                DKUConstants.API_RESPONSE_KEY: json.dumps(item)
+            })
+        else:
+            output.update(parse_keys_for_json(item))
+        return output
+
     def get_writer(self, dataset_schema=None, dataset_partitioning=None,
                    partition_id=None):
         """

python-lib/dku_constants.py

@@ -1,3 +1,4 @@
 class DKUConstants(object):
+    API_RESPONSE_KEY = "api_response"
     RAW_BODY_FORMAT = "RAW"
     FORM_DATA_BODY_FORMAT = "FORM_DATA"

python-lib/dku_utils.py

@@ -1,3 +1,7 @@
+import json
+import copy
+
+
 def get_dku_key_values(endpoint_query_string):
     return {key_value.get("from"): key_value.get("to") for key_value in endpoint_query_string if key_value.get("from")}
 
@@ -19,7 +23,30 @@ def get_endpoint_parameters(configuration):
         "requests_per_minute",
         "pagination_type",
         "next_page_url_key",
-        "top_key", "skip_key"
+        "top_key", "skip_key", "maximum_number_rows"
     ]
     parameters = {endpoint_parameter: configuration.get(endpoint_parameter) for endpoint_parameter in endpoint_parameters if configuration.get(endpoint_parameter) is not None}
     return parameters
+
+
+def parse_keys_for_json(items):
+    ret = {}
+    for key in items:
+        value = items.get(key)
+        if isinstance(value, dict) or isinstance(value, list):
+            ret.update({key: json.dumps(value)})
+        elif value is None:
+            continue
+        else:
+            ret.update({key: value})
+    return ret
+
+
+def get_value_from_path(dictionary, path):
+    ret = copy.deepcopy(dictionary)
+    for key in path:
+        if key in ret and isinstance(ret, dict):
+            ret = ret.get(key)
+        else:
+            raise ValueError("The extraction path {} was not found in the incoming data".format(path))
+    return ret

python-lib/rest_api_client.py

@@ -107,6 +107,7 @@ def __init__(self, credential, endpoint, custom_key_values={}):
         elif body_format in [DKUConstants.FORM_DATA_BODY_FORMAT]:
             key_value_body = endpoint.get("key_value_body", {})
             self.requests_kwargs.update({"json": get_dku_key_values(key_value_body)})
+        self.metadata = {}
 
     def set_login(self, credential):
         login_type = credential.get("login_type", "no_auth")
@@ -140,22 +141,30 @@ def request(self, method, url, can_raise_exeption=True, **kwargs):
         if self.loop_detector.is_stuck_in_loop(url, kwargs.get("params", {}), kwargs.get("headers", {})):
             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)
+            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}
+        self.set_metadata("request_duration", request_finish_time - request_start_time)
         self.time_last_request = time.time()
+        self.set_metadata("status_code", response.status_code)
+        self.set_metadata("response_headers", "{}".format(response.headers))
         if response.status_code >= 400:
             error_message = "Error {}: {}".format(response.status_code, response.content)
             self.pagination.is_last_batch_empty = True
             if can_raise_exeption:
                 raise RestAPIClientError(error_message)
             else:
                 return {"error": error_message}
+        if response.status_code in [204]:
+            self.pagination.update_next_page({}, response.links)
+            return self.empty_json_response()
         json_response = response.json()
         self.pagination.update_next_page(json_response, response.links)
         return json_response
@@ -167,6 +176,12 @@ def paginated_api_call(self, can_raise_exeption=True):
         self.requests_kwargs.update({"params": params})
         return self.request(self.http_method, self.pagination.get_next_page_url(), can_raise_exeption, **self.requests_kwargs)
 
+    def empty_json_response(self):
+        return {self.extraction_key: {}} if self.extraction_key else {}
+
+    def set_metadata(self, metadata_name, value):
+        self.metadata["dku_{}".format(metadata_name)] = value
+
     @staticmethod
     def get_params(endpoint_query_string, keywords):
         templated_query_string = get_dku_key_values(endpoint_query_string)
@@ -191,3 +206,6 @@ def enforce_throttling(self):
             if time_since_last_resquests < self.time_between_requests:
                 logger.info("Enforcing {}s throttling".format(self.time_between_requests - time_since_last_resquests))
                 time.sleep(self.time_between_requests - time_since_last_resquests)
+
+    def get_metadata(self):
+        return self.metadata

python-lib/rest_api_recipe_session.py

@@ -1,20 +1,28 @@
 from dataikuapi.utils import DataikuException
 from rest_api_client import RestAPIClient
 from safe_logger import SafeLogger
+from dku_utils import parse_keys_for_json
+from dku_constants import DKUConstants
 import copy
+import json
 
 logger = SafeLogger("api-connect plugin", forbiden_keys=["token", "password"])
 
 
 class RestApiRecipeSession:
-    def __init__(self, custom_key_values, credential_parameters, endpoint_parameters, extraction_key, parameter_columns, parameter_renamings):
+    def __init__(self, custom_key_values, credential_parameters, endpoint_parameters, extraction_key, parameter_columns, parameter_renamings,
+                 display_metadata=False,
+                 maximum_number_rows=-1):
         self.custom_key_values = custom_key_values
         self.credential_parameters = credential_parameters
         self.endpoint_parameters = endpoint_parameters
         self.extraction_key = extraction_key
         self.client = None
         self.initial_parameter_columns = None
         self.column_to_parameter_dict = self.get_column_to_parameter_dict(parameter_columns, parameter_renamings)
+        self.display_metadata = display_metadata
+        self.maximum_number_rows = maximum_number_rows
+        self.is_row_limit = (self.maximum_number_rows > 0)
 
     @staticmethod
     def get_column_to_parameter_dict(parameter_columns, parameter_renamings):
@@ -30,6 +38,7 @@ def process_dataframe(self, input_parameters_dataframe, is_raw_output):
         results = []
         time_last_request = None
         for index, input_parameters_row in input_parameters_dataframe.iterrows():
+            rows_count = 0
             self.initial_parameter_columns = {}
             for column_name in self.column_to_parameter_dict:
                 parameter_name = self.column_to_parameter_dict[column_name]
@@ -46,42 +55,64 @@ def process_dataframe(self, input_parameters_dataframe, is_raw_output):
             while self.client.has_more_data():
                 page_results = self.retrieve_next_page(is_raw_output)
                 results.extend(page_results)
+                rows_count += len(page_results)
+                if self.is_row_limit and rows_count >= self.maximum_number_rows:
+                    break
             time_last_request = self.client.time_last_request
         return results
 
     def retrieve_next_page(self, is_raw_output):
         page_rows = []
-        base_row = copy.deepcopy(self.initial_parameter_columns)
         logger.info("retrieve_next_page: Calling next page")
         json_response = self.client.paginated_api_call(can_raise_exeption=False)
+        metadata = self.client.get_metadata() if self.display_metadata else {}
+        is_api_returning_dict = True
         if self.extraction_key:
             data_rows = json_response.get(self.extraction_key, [json_response])
             if data_rows is None:
                 raise DataikuException("Extraction key '{}' was not found in the incoming data".format(self.extraction_key))
-            page_rows.extend(self.format_page_rows(data_rows, is_raw_output))
+            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
+            base_row = copy.deepcopy(metadata)
             if is_raw_output:
                 if is_error_message(json_response):
-                    base_row.update(json_response)
+                    base_row.update(parse_keys_for_json(json_response))
                 else:
-                    base_row.update({"api_response": json_response})
+                    base_row.update({
+                        DKUConstants.API_RESPONSE_KEY: json.dumps(json_response)
+                    })
             else:
-                base_row.update(json_response)
-            page_rows.append(base_row)
+                if isinstance(json_response, dict):
+                    base_row.update(parse_keys_for_json(json_response))
+                elif isinstance(json_response, list):
+                    is_api_returning_dict = False
+                    for row in json_response:
+                        base_row = copy.deepcopy(metadata)
+                        base_row.update(parse_keys_for_json(row))
+                        base_row.update(self.initial_parameter_columns)
+                        page_rows.append(base_row)
+
+            if is_api_returning_dict:
+                base_row.update(self.initial_parameter_columns)
+                page_rows.append(base_row)
         return page_rows
 
-    def format_page_rows(self, data_rows, is_raw_output):
+    def format_page_rows(self, data_rows, is_raw_output, metadata=None):
         page_rows = []
+        metadata = metadata or {}
         for data_row in data_rows:
             base_row = copy.deepcopy(self.initial_parameter_columns)
+            base_row.update(metadata)
             if is_raw_output:
                 if is_error_message(data_row):
-                    base_row.update(data_row)
+                    base_row.update(parse_keys_for_json(data_row))
                 else:
-                    base_row.update({"api_response": data_row})
+                    base_row.update({
+                        DKUConstants.API_RESPONSE_KEY: json.dumps(data_row)
+                    })
             else:
-                base_row.update(data_row)
+                base_row.update(parse_keys_for_json(data_row))
             page_rows.append(base_row)
         return page_rows
 

tests/python/integration/test_scenario.py

@@ -9,3 +9,19 @@ def test_run_api_connect_authentication_modes(user_dss_clients):
 
 def test_run_api_connect_pagination_modes(user_dss_clients):
     dss_scenario.run(user_dss_clients, project_key=TEST_PROJECT_KEY, scenario_id="PAGINATION")
+
+
+def test_run_api_connect_recipes(user_dss_clients):
+    dss_scenario.run(user_dss_clients, project_key=TEST_PROJECT_KEY, scenario_id="Recipes")
+
+
+def test_run_api_connect_using_global_variable(user_dss_clients):
+    dss_scenario.run(user_dss_clients, project_key=TEST_PROJECT_KEY, scenario_id="UsingGlobalVariable")
+
+
+def test_run_api_connect_array_api(user_dss_clients):
+    dss_scenario.run(user_dss_clients, project_key=TEST_PROJECT_KEY, scenario_id="ArrayAPI")
+
+
+def test_run_api_connect_search_path(user_dss_clients):
+    dss_scenario.run(user_dss_clients, project_key=TEST_PROJECT_KEY, scenario_id="SearchPath")