[incubator-superset] 01/02: Return full info when doing query cost estimation

[be...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

beto pushed a commit to branch VIZ-1202
in repository https://gitbox.apache.org/repos/asf/incubator-superset.git

commit 18f249e474a43aae2b4dcee5b1d0c6b69bfbf669
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Mon Nov 25 11:36:54 2019 -0800

    Return full info when doing query cost estimation
---
 superset/db_engine_specs/presto.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py
index 661fe35..f3564cf 100644
--- a/superset/db_engine_specs/presto.py
+++ b/superset/db_engine_specs/presto.py
@@ -445,7 +445,7 @@ class PrestoEngineSpec(BaseEngineSpec):
     @classmethod
     def estimate_statement_cost(  # pylint: disable=too-many-locals
         cls, statement: str, database, cursor, user_name: str
-    ) -> Dict[str, float]:
+    ) -> Dict[str, Any]:
         """
         Run a SQL query that estimates the cost of a given statement.
 
@@ -453,7 +453,7 @@ class PrestoEngineSpec(BaseEngineSpec):
         :param database: Database instance
         :param cursor: Cursor instance
         :param username: Effective username
-        :return: JSON estimate from Presto
+        :return: JSON response from Presto
         """
         parsed_query = ParsedQuery(statement)
         sql = parsed_query.stripped()
@@ -479,11 +479,11 @@ class PrestoEngineSpec(BaseEngineSpec):
         #     }
         #   }
         result = json.loads(cursor.fetchone()[0])
-        return result["estimate"]
+        return result
 
     @classmethod
     def query_cost_formatter(
-        cls, raw_cost: List[Dict[str, float]]
+        cls, raw_cost: List[Dict[str, Any]]
     ) -> List[Dict[str, str]]:
         """
         Format cost estimate.
@@ -516,10 +516,11 @@ class PrestoEngineSpec(BaseEngineSpec):
             ("networkCost", "Network cost", ""),
         ]
         for row in raw_cost:
+            estimate: Dict[str, float] = row.get("estimate", {})
             statement_cost = {}
             for key, label, suffix in columns:
-                if key in row:
-                    statement_cost[label] = humanize(row[key], suffix).strip()
+                if key in estimate:
+                    statement_cost[label] = humanize(estimate[key], suffix).strip()
             cost.append(statement_cost)
 
         return cost