From e3f4822aed8e280600bdbe3c0f7f44e6f754edc3 Mon Sep 17 00:00:00 2001
From: piby180 <piby180@gmail.com>
Date: Sat, 13 Jul 2024 18:44:17 +0200
Subject: [PATCH] Add raw query response to cursor class (#103)

---
 pinotdb/db.py         |   9 ++++
 tests/unit/test_db.py | 114 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 123 insertions(+)

diff --git a/pinotdb/db.py b/pinotdb/db.py
index e3d73b7..31cbb3f 100644
--- a/pinotdb/db.py
+++ b/pinotdb/db.py
@@ -315,6 +315,7 @@ def __init__(
         self.schema = None
         self.rowcount = -1
         self._results = None
+        self.raw_query_response = None
         self.timeUsedMs = -1
         self._debug = debug
         self._preserve_types = preserve_types
@@ -396,7 +397,15 @@ def finalize_query_payload(
     def normalize_query_response(self, input_query, query_response):
         try:
             payload = query_response.json()
+            self.raw_query_response = {
+                    "response" : payload, 
+                    "status_code" : query_response.status_code
+                }
         except Exception as e:
+            self.raw_query_response = {
+                    "response" : query_response.text, 
+                    "status_code" : query_response.status_code
+                }
             raise exceptions.DatabaseError(
                 f"Error when querying {input_query} from {self.url}, "
                 f"raw response is:\n{query_response.text}"
diff --git a/tests/unit/test_db.py b/tests/unit/test_db.py
index 3bd9ddd..a05b610 100644
--- a/tests/unit/test_db.py
+++ b/tests/unit/test_db.py
@@ -678,6 +678,120 @@ def test_does_nothing_for_setoutputsizes(self):
         cursor.setoutputsizes(123)
 
         # All good, nothing happened
+    
+    def test_checks_raw_query_response_with_single_stage_and_status_code_200(self):
+        cursor = self.create_cursor(
+                    {
+                        'dataSchema': {
+                            'columnNames': ['age'],
+                            'columnDataTypes': ['INT'],
+                        },
+                        'rows': [[1], [2], [3]],
+                    },
+                    status_code=200,
+                    extra_payload={'numGroupsLimitReached': True},
+                    use_multistage_engine=False
+                )
+        cursor.execute('some statement')
+        raw_query_response = {
+            'response': {
+                'numServersResponded': 1,
+                'numServersQueried': 1,
+                'resultTable': {
+                    'dataSchema': {'columnNames': ['age'], 'columnDataTypes': ['INT']},
+                    'rows': [[1], [2], [3]],
+                },
+                'numGroupsLimitReached': True,
+            },
+            'status_code': 200,
+        }
+        self.assertEqual(cursor.raw_query_response, raw_query_response)
+
+    def test_checks_raw_query_response_with_multi_stage_and_status_code_200(self):
+        cursor = self.create_cursor(
+                    {
+                        'dataSchema': {
+                            'columnNames': ['age'],
+                            'columnDataTypes': ['INT'],
+                        },
+                        'rows': [[1], [2], [3]],
+                    },
+                    status_code=200,
+                    extra_payload={'numGroupsLimitReached': True},
+                    use_multistage_engine=True
+                )
+        cursor.execute('some statement')
+        raw_query_response = {
+            'response': {
+                'numServersResponded': 1,
+                'numServersQueried': 1,
+                'resultTable': {
+                    'dataSchema': {'columnNames': ['age'], 'columnDataTypes': ['INT']},
+                    'rows': [[1], [2], [3]],
+                },
+                'numGroupsLimitReached': True,
+            },
+            'status_code': 200,
+        }
+        self.assertEqual(cursor.raw_query_response, raw_query_response)
+
+    def test_checks_raw_query_response_with_single_stage_and_status_code_400(self):
+        cursor = self.create_cursor(
+            {
+                'dataSchema': {
+                    'columnNames': ['age'],
+                    'columnDataTypes': ['INT'],
+                },
+                'rows': [[1], [2], [3]],
+            },
+            status_code=400,
+            extra_payload={'exceptions': ['something', 'wrong']},
+            use_multistage_engine=False
+        )
+        with self.assertRaises(exceptions.ProgrammingError):
+            cursor.execute('some statement')
+        raw_query_response = {
+            'response': {
+                'numServersResponded': 1,
+                'numServersQueried': 1,
+                'resultTable': {
+                    'dataSchema': {'columnNames': ['age'], 'columnDataTypes': ['INT']},
+                    'rows': [[1], [2], [3]],
+                },
+                'exceptions': ['something', 'wrong'],
+            },
+            'status_code': 400,
+        }
+        self.assertEqual(cursor.raw_query_response, raw_query_response)
+
+    def test_checks_raw_query_response_with_multi_stage_and_status_code_400(self):
+        cursor = self.create_cursor(
+            {
+                'dataSchema': {
+                    'columnNames': ['age'],
+                    'columnDataTypes': ['INT'],
+                },
+                'rows': [[1], [2], [3]],
+            },
+            status_code=400,
+            extra_payload={'exceptions': ['something', 'wrong']},
+            use_multistage_engine=True
+        )
+        with self.assertRaises(exceptions.ProgrammingError):
+            cursor.execute('some statement')
+        raw_query_response = {
+            'response': {
+                'numServersResponded': 1,
+                'numServersQueried': 1,
+                'resultTable': {
+                    'dataSchema': {'columnNames': ['age'], 'columnDataTypes': ['INT']},
+                    'rows': [[1], [2], [3]],
+                },
+                'exceptions': ['something', 'wrong'],
+            },
+            'status_code': 400,
+        }
+        self.assertEqual(cursor.raw_query_response, raw_query_response)
 
 
 class AsyncCursorTest(IsolatedAsyncioTestCase):