Merge branch 'release/8.0'

Author: kevanescence

dataikuapi/dss/future.py

@@ -72,6 +72,8 @@ def wait_for_result(self):
         """
         Wait and get the future result
         """
+        if self.state.get('hasResult', False):
+            return self.result_wrapper(self.state.get('result', None))
         if self.state is None or not self.state.get('hasResult', False) or self.state_is_peek:
             self.get_state()
         while not self.state.get('hasResult', False):

dataikuapi/dss/recipe.py

@@ -76,15 +76,25 @@ def run(self, job_type="NON_RECURSIVE_FORCED_BUILD", partitions=None, wait=True,
         :return: the :class:`dataikuapi.dss.job.DSSJob` job handle corresponding to the built job
         :rtype: :class:`dataikuapi.dss.job.DSSJob`
         """
+        project = self.client.get_project(self.project_key)
+        outputs = project.get_flow().get_graph().get_successor_computables(self)
 
-        settings = self.get_settings()
-        output_refs = settings.get_flat_output_refs()
-
-        if len(output_refs) == 0:
+        if len(outputs) == 0:
             raise Exception("recipe has no outputs, can't run it")
 
-        jd = self.client.get_project(self.project_key).new_job(job_type)
-        jd.with_output(output_refs[0], partition=partitions)
+        first_output = outputs[0]
+
+        object_type_map = {
+            "COMPUTABLE_DATASET": "DATASET",
+            "COMPUTABLE_FOLDER": "MANAGED_FOLDER",
+            "COMPUTABLE_SAVED_MODEL": "SAVED_MODEL",
+            "COMPUTABLE_STREAMING_ENDPOINT": "STREAMING_ENDPOINT",
+        }
+        if first_output["type"] in object_type_map:
+            jd = project.new_job(job_type)
+            jd.with_output(first_output["ref"], object_type=object_type_map[first_output["type"]], partition=partitions)
+        else:
+            raise Exception("Recipe has unsupported output type {}, can't run it".format(first_output["type"]))
 
         if wait:
             return jd.start_and_wait()
@@ -1121,7 +1131,7 @@ def set_code_env(self, code_env=None, inherit=False, use_builtin=False):
             raise ValueError("This recipe kind does not seem to take a code env selection")
 
         if code_env is not None:
-            rp["envSelection"] = {"envMode": "EXPLICIT_ENV", "envName": "code_env"}
+            rp["envSelection"] = {"envMode": "EXPLICIT_ENV", "envName": code_env}
         elif inherit:
             rp["envSelection"] = {"envMode": "INHERIT"}
         elif use_builtin:

dataikuapi/dss/sqlquery.py

@@ -3,11 +3,13 @@
 from ..utils import DataikuStreamedHttpUTF8CSVReader
 import json
 
+
 class DSSSQLQuery(object):
     """
-    A connection to a database or database-like on which queries can be run through DSS
+    A connection to a database or database-like on which queries can be run through DSS.
+    Do not create this class directly, instead use :meth:`dataikuapi.DSSClient.sql_query`
     """
-    def __init__(self, client, query, connection, database, dataset_full_name, pre_queries, post_queries, type, extra_conf, script_steps, script_input_schema, script_output_schema, script_report_location, read_timestamp_without_timezone_as_string, read_date_as_string):
+    def __init__(self, client, query, connection, database, dataset_full_name, pre_queries, post_queries, type, extra_conf, script_steps, script_input_schema, script_output_schema, script_report_location, read_timestamp_without_timezone_as_string, read_date_as_string, project_key):
         self.client = client
 
         self.streaming_session = self.client._perform_json(
@@ -19,6 +21,7 @@ def __init__(self, client, query, connection, database, dataset_full_name, pre_q
                     "connection" : connection,
                     "database" : database,
                     "datasetFullName" : dataset_full_name,
+                    "projectKey" : project_key,
                     "type" : type,
                     "extraConf" : extra_conf,
                     "scriptSteps" : script_steps,

dataikuapi/dssclient.py

@@ -302,7 +302,7 @@ def get_plugin(self, plugin_id):
     # SQL queries
     ########################################################
 
-    def sql_query(self, query, connection=None, database=None, dataset_full_name=None, pre_queries=None, post_queries=None, type='sql', extra_conf=None, script_steps=None, script_input_schema=None, script_output_schema=None, script_report_location=None, read_timestamp_without_timezone_as_string=True, read_date_as_string=False):
+    def sql_query(self, query, connection=None, database=None, dataset_full_name=None, pre_queries=None, post_queries=None, type='sql', extra_conf=None, script_steps=None, script_input_schema=None, script_output_schema=None, script_report_location=None, read_timestamp_without_timezone_as_string=True, read_date_as_string=False, project_key=None):
         """
         Initiate a SQL, Hive or Impala query and get a handle to retrieve the results of the query.
         Internally, the query is run by DSS. The  database to run the query on is specified either by 
@@ -316,12 +316,13 @@ def sql_query(self, query, connection=None, database=None, dataset_full_name=Non
         :param list pre_queries: (optional) array of queries to run before the query
         :param list post_queries: (optional) array of queries to run after the query
         :param str type: the type of query : either 'sql', 'hive' or 'impala'
-        
+        :param str project_key: The project_key on which the query should be run (especially useful for user isolation/impersonation scenario)
+
         :returns: A :class:`dataikuapi.dss.sqlquery.DSSSQLQuery` query handle
         """
         if extra_conf is None:
             extra_conf = {}
-        return DSSSQLQuery(self, query, connection, database, dataset_full_name, pre_queries, post_queries, type, extra_conf, script_steps, script_input_schema, script_output_schema, script_report_location, read_timestamp_without_timezone_as_string, read_date_as_string)
+        return DSSSQLQuery(self, query, connection, database, dataset_full_name, pre_queries, post_queries, type, extra_conf, script_steps, script_input_schema, script_output_schema, script_report_location, read_timestamp_without_timezone_as_string, read_date_as_string, project_key)
 
     ########################################################
     # Users