Reorganized query function, got rid of tu[le return values

superset/models.py

@@ -1238,7 +1238,8 @@ def values_for_column(self,
         )
 
     def get_query_str(  # sqla
-            self, groupby, metrics,
+            self, engine, qry_start_dttm,
+            groupby, metrics,
             granularity,
             from_dttm, to_dttm,
             filter=None,  # noqa
@@ -1261,7 +1262,6 @@ def get_query_str(  # sqla
 
         cols = {col.column_name: col for col in self.columns}
         metrics_dict = {m.metric_name: m for m in self.metrics}
-        qry_start_dttm = datetime.now()
 
         if not granularity and is_timeseries:
             raise Exception(_(
@@ -1421,17 +1421,18 @@ def visit_column(element, compiler, **kw):
 
         qry = qry.select_from(tbl)
 
-        engine = self.database.get_sqla_engine()
         sql = "{}".format(
             qry.compile(
                 engine, compile_kwargs={"literal_binds": True},),
         )
         logging.info(sql)
         sql = sqlparse.format(sql, reindent=True)
-        return sql, engine, qry_start_dttm
+        return sql
 
-    def query(self, sql, engine, qry_start_dttm):
-        start_dttm = datetime.now()
+    def query(self, query_obj):
+        qry_start_dttm = datetime.now()
+        engine = self.database.get_sqla_engine()
+        sql = self.get_query_str(engine, qry_start_dttm, **query_obj)
         status = QueryStatus.SUCCESS
         error_message = None
         df = None
@@ -2271,7 +2272,8 @@ def values_for_column(self,
         return df
 
     def get_query_str(  # druid
-            self, groupby, metrics,
+            self, client, qry_start_dttm,
+            groupby, metrics,
             granularity,
             from_dttm, to_dttm,
             filter=None,  # noqa
@@ -2283,13 +2285,12 @@ def get_query_str(  # druid
             orderby=None,
             extras=None,  # noqa
             select=None,  # noqa
-            columns=None, ):
+            columns=None, phase=2):
         """Runs a query against Druid and returns a dataframe.
 
         This query interface is common to SqlAlchemy and Druid
         """
         # TODO refactor into using a TBD Query object
-        qry_start_dttm = datetime.now()
         if not is_timeseries:
             granularity = 'all'
         inner_from_dttm = inner_from_dttm or from_dttm
@@ -2385,7 +2386,6 @@ def recursive_get_fields(_conf):
         if having_filters:
             qry['having'] = having_filters
 
-        client = self.cluster.get_pydruid_client()
         orig_filters = filters
         if len(groupby) == 0:
             del qry['dimensions']
@@ -2424,6 +2424,8 @@ def recursive_get_fields(_conf):
                 query_str += json.dumps(
                     client.query_builder.last_query.query_dict, indent=2)
                 query_str += "\n"
+                if phase == 1:
+                    return query_str
                 query_str += (
                     "//\nPhase 2 (built based on phase one's results)\n")
                 df = client.export_pandas()
@@ -2463,30 +2465,35 @@ def recursive_get_fields(_conf):
             client.groupby(**qry)
         query_str += json.dumps(
             client.query_builder.last_query.query_dict, indent=2)
-        return query_str, client, qry_start_dttm
+        return query_str
 
-    def query(query_str, client, qry_start_dttm):
+    def query(self, query_obj):
+        qry_start_dttm = datetime.now()
+        client = self.cluster.get_pydruid_client()
+        query_str = self.get_query_str(client, qry_start_dttm, **query_obj)
         df = client.export_pandas()
         if df is None or df.size == 0:
             raise Exception(_("No data was returned."))
         df.columns = [
             DTTM_ALIAS if c == 'timestamp' else c for c in df.columns]
 
+        is_timeseries = query_obj['is_timeseries'] \
+            if 'is_timeseries' in query_obj else True
         if (
                 not is_timeseries and
-                granularity == "all" and
+                query_obj['granularity'] == "all" and
                 DTTM_ALIAS in df.columns):
             del df[DTTM_ALIAS]
 
         # Reordering columns
         cols = []
         if DTTM_ALIAS in df.columns:
             cols += [DTTM_ALIAS]
-        cols += [col for col in groupby if col in df.columns]
-        cols += [col for col in metrics if col in df.columns]
+        cols += [col for col in query_obj['groupby'] if col in df.columns]
+        cols += [col for col in query_obj['metrics'] if col in df.columns]
         df = df[cols]
 
-        time_offset = DruidDatasource.time_offset(granularity)
+        time_offset = DruidDatasource.time_offset(query_obj['granularity'])
 
         def increment_timestamp(ts):
             dt = utils.parse_human_datetime(ts).replace(

superset/views.py

@@ -1488,12 +1488,17 @@ def explore_json(self, datasource_type, datasource_id):
 
         if request.args.get("query") == "true":
             try:
-                query = viz_obj.get_query()
+                query_obj = viz_obj.query_obj()
+                engine = viz_obj.datasource.database.get_sqla_engine() \
+                    if datasource_type == 'table' \
+                    else viz_obj.datasource.cluster.get_pydruid_client()
+                if datasource_type == 'druid':
+                    # only retrive first phase query for druid
+                    query_obj['phase'] = 1
+                query = viz_obj.datasource.get_query_str(
+                    engine, datetime.now(), **query_obj)
             except Exception as e:
-                return Response(
-                    json.dumps({'error': e}),
-                    status=404,
-                    mimetype="application/json")
+                return json_error_response(e)
             return Response(
                 json.dumps({'query': query}),
                 status=200,

superset/viz.py

@@ -106,9 +106,8 @@ def get_df(self, query_obj=None):
                 timestamp_format = dttm_col.python_date_format
 
         # The datasource here can be different backend but the interface is common
-        query_str, engine, start_dttm = self.datasource.get_query_str(**query_obj)
-        self.query = query_str
-        self.results = self.datasource.query(query_str, engine, start_dttm)
+        self.results = self.datasource.query(query_obj)
+        self.query = self.results.query
         self.status = self.results.status
         self.error_message = self.results.error_message
 
@@ -312,11 +311,6 @@ def get_csv(self):
         include_index = not isinstance(df.index, pd.RangeIndex)
         return df.to_csv(index=include_index, encoding="utf-8")
 
-    def get_query(self):
-        query_obj = self.query_obj()
-        query_str, engine, start_dttm = self.datasource.get_query_str(**query_obj)
-        return query_str
-
     def get_values_for_column(self, column):
         """
         Retrieves values for a column to be used by the filter dropdown.