Pass full response to query_cost_formatter (#8652)

* Return full info when doing query cost estimation

* Add unit test

* Fix isort

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