Merge pull request #34 from dataiku/feature/shaker2018

Feature/shaker2018

Author: FChataigner

dataikuapi/dss/sqlquery.py

@@ -1,33 +1,30 @@
 from ..utils import DataikuException
 from ..utils import DataikuUTF8CSVReader
 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
     """
-    def __init__(self, client, query, connection, database, dataset_full_name, pre_queries, post_queries, type, extra_conf):
+    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):
         self.client = client
-        self.query = query
-        self.connection = connection
-        self.database = database
-        self.dataset_full_name = dataset_full_name
-        self.pre_queries = pre_queries
-        self.post_queries = post_queries
-        self.extra_conf = extra_conf
-        self.type = type
 
         self.streaming_session = self.client._perform_json(
                 "POST", "/sql/queries/",
                 body = {
-                    "query" : self.query,
-                    "preQueries" : self.pre_queries,
-                    "postQueries" : self.post_queries,
-                    "connection" : self.connection,
-                    "database" : self.database,
-                    "datasetFullName" : self.dataset_full_name,
-                    "type" : self.type,
-                    "extraConf" : extra_conf
+                    "query" : query,
+                    "preQueries" : pre_queries,
+                    "postQueries" : post_queries,
+                    "connection" : connection,
+                    "database" : database,
+                    "datasetFullName" : dataset_full_name,
+                    "type" : type,
+                    "extraConf" : extra_conf,
+                    "scriptSteps" : script_steps,
+                    "scriptInputSchema" : script_input_schema,
+                    "scriptOutputSchema" : script_output_schema,
+                    "scriptReportLocation" : script_report_location
                 })
         self.queryId = self.streaming_session['queryId']
 

dataikuapi/dssclient.py

@@ -182,7 +182,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={}):
+    def sql_query(self, query, connection=None, database=None, dataset_full_name=None, pre_queries=None, post_queries=None, type='sql', extra_conf={}, script_steps=None, script_input_schema=None, script_output_schema=None, script_report_location=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 
@@ -199,7 +199,7 @@ def sql_query(self, query, connection=None, database=None, dataset_full_name=Non
         
         :returns: A :class:`dataikuapi.dss.sqlquery.DSSSQLQuery` query handle
         """
-        return DSSSQLQuery(self, query, connection, database, dataset_full_name, pre_queries, post_queries, type, 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)
 
     ########################################################
     # Users