Skip to content

Add raw query response to cursor class #103

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jul 13, 2024
Merged

Add raw query response to cursor class #103

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
bc2d6e6
Add raw query response to cursor class
piby180 Jul 4, 2024
4e81b23
Add raw response to cursor class
piby180 Jul 4, 2024
8c7c76a
Rename attribute
piby180 Jul 4, 2024
8e02c1a
Added tests for cursor.raw_query_response
piby180 Jul 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions pinotdb/db.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -395,7 +396,15 @@ def finalize_query_payload(
def normalize_query_response(self, input_query, query_response):
try:
payload = query_response.json()
self.raw_query_response = {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you please add a test using normal query response and failure results.

As well as v1/v2 query engine.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@xiangfu0 I have added the tests. Please review.

"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}"
Expand Down
114 changes: 114 additions & 0 deletions tests/unit/test_db.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
Loading