wrap error resoponse into exception

Signed-off-by: Zhongnan Su <szhongna@amazon.com>
opensearch-project · Nov 30, 2021 · e61b69f · e61b69f
1 parent 7f5ed9b
commit e61b69f
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/sql-cli/src/opensearch_sql_cli/formatter.py b/sql-cli/src/opensearch_sql_cli/formatter.py
@@ -55,10 +55,15 @@ def format_output(self, data):
         formatter = TabularOutputFormatter(format_name=self.table_format)
 
         # parse response data
-        datarows = data["datarows"]
-        schema = data["schema"]
-        total_hits = data["total"]
-        cur_size = data["size"]
+        try:
+            datarows = data["datarows"]
+            schema = data["schema"]
+            total_hits = data["total"]
+            cur_size = data["size"]
+        except KeyError as e:
+            # in case some errors in query engine don't throw exception, leading to a KeyError
+            raise Exception(data)
+
         # unused data for now,
         fields = []
         types = []