From bc2d6e6e671bf52a636e29ef0dd0bcbd0dc28435 Mon Sep 17 00:00:00 2001 From: piby180 Date: Thu, 4 Jul 2024 15:22:16 +0000 Subject: [PATCH 1/4] Add raw query response to cursor class --- pinotdb/db.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pinotdb/db.py b/pinotdb/db.py index 4580895..c055bc9 100644 --- a/pinotdb/db.py +++ b/pinotdb/db.py @@ -314,6 +314,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 @@ -395,6 +396,7 @@ def finalize_query_payload( def normalize_query_response(self, input_query, query_response): try: payload = query_response.json() + self.raw_query_response = payload except Exception as e: raise exceptions.DatabaseError( f"Error when querying {input_query} from {self.url}, " From 4e81b234dcab0e3c7063b2683ac2aab4960a8e8b Mon Sep 17 00:00:00 2001 From: piby180 Date: Thu, 4 Jul 2024 15:32:02 +0000 Subject: [PATCH 2/4] Add raw response to cursor class --- pinotdb/db.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/pinotdb/db.py b/pinotdb/db.py index c055bc9..277b9b6 100644 --- a/pinotdb/db.py +++ b/pinotdb/db.py @@ -396,8 +396,15 @@ def finalize_query_payload( def normalize_query_response(self, input_query, query_response): try: payload = query_response.json() - self.raw_query_response = payload + self.raw_query_response = { + "response" : payload, + "response_status" : query_response.status_code + } except Exception as e: + self.raw_query_response = { + "response" : query_response.text, + "response_status" : query_response.status_code + } raise exceptions.DatabaseError( f"Error when querying {input_query} from {self.url}, " f"raw response is:\n{query_response.text}" From 8c7c76a59ea605dd57ee8dd7649fb1d5926f8252 Mon Sep 17 00:00:00 2001 From: piby180 Date: Thu, 4 Jul 2024 16:35:01 +0000 Subject: [PATCH 3/4] Rename attribute --- pinotdb/db.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pinotdb/db.py b/pinotdb/db.py index 277b9b6..eb0bcbe 100644 --- a/pinotdb/db.py +++ b/pinotdb/db.py @@ -397,13 +397,13 @@ def normalize_query_response(self, input_query, query_response): try: payload = query_response.json() self.raw_query_response = { - "response" : payload, - "response_status" : query_response.status_code + "response" : payload, + "status_code" : query_response.status_code } except Exception as e: self.raw_query_response = { - "response" : query_response.text, - "response_status" : query_response.status_code + "response" : query_response.text, + "status_code" : query_response.status_code } raise exceptions.DatabaseError( f"Error when querying {input_query} from {self.url}, " From 8e02c1a8f4041b5775c1bebfb0b10146916f9a4e Mon Sep 17 00:00:00 2001 From: piby180 Date: Fri, 5 Jul 2024 19:55:10 +0000 Subject: [PATCH 4/4] Added tests for cursor.raw_query_response --- tests/unit/test_db.py | 114 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 114 insertions(+) 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):