Skip to content

Commit

Permalink
[PECO-1440] Expose current query id on cursor object (#364)
Browse files Browse the repository at this point in the history
* [PECO-1440] Expose current query id on cursor object

Signed-off-by: Levko Kravets <levko.ne@gmail.com>

* Clear `active_op_handle` when closing the cursor

Signed-off-by: Levko Kravets <levko.ne@gmail.com>

---------

Signed-off-by: Levko Kravets <levko.ne@gmail.com>
  • Loading branch information
kravets-levko authored Mar 4, 2024
1 parent 5048934 commit f6fd7a7
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 0 deletions.
14 changes: 14 additions & 0 deletions src/databricks/sql/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import json
import os
import decimal
from uuid import UUID

from databricks.sql import __version__
from databricks.sql import *
Expand Down Expand Up @@ -1004,9 +1005,22 @@ def cancel(self) -> None:
def close(self) -> None:
"""Close cursor"""
self.open = False
self.active_op_handle = None
if self.active_result_set:
self._close_and_clear_active_result_set()

@property
def query_id(self) -> Optional[str]:
"""
This attribute is an identifier of last executed query.
This attribute will be ``None`` if the cursor has not had an operation
invoked via the execute method yet, or if cursor was closed.
"""
if self.active_op_handle is not None:
return str(UUID(bytes=self.active_op_handle.operationId.guid))
return None

@property
def description(self) -> Optional[List[Tuple]]:
"""
Expand Down
21 changes: 21 additions & 0 deletions tests/unit/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,14 @@
import itertools
from decimal import Decimal
from datetime import datetime, date
from uuid import UUID

from databricks.sql.thrift_api.TCLIService.ttypes import (
TOpenSessionResp,
TExecuteStatementResp,
TOperationHandle,
THandleIdentifier,
TOperationType
)
from databricks.sql.thrift_backend import ThriftBackend

Expand Down Expand Up @@ -610,6 +614,23 @@ def test_staging_operation_response_is_handled(self, mock_client_class, mock_han

mock_handle_staging_operation.call_count == 1

@patch("%s.client.ThriftBackend" % PACKAGE_NAME, ThriftBackendMockFactory.new())
def test_access_current_query_id(self):
operation_id = 'EE6A8778-21FC-438B-92D8-96AC51EE3821'

connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
cursor = connection.cursor()

self.assertIsNone(cursor.query_id)

cursor.active_op_handle = TOperationHandle(
operationId=THandleIdentifier(guid=UUID(operation_id).bytes, secret=0x00),
operationType=TOperationType.EXECUTE_STATEMENT)
self.assertEqual(cursor.query_id.upper(), operation_id.upper())

cursor.close()
self.assertIsNone(cursor.query_id)


if __name__ == '__main__':
suite = unittest.TestLoader().loadTestsFromModule(sys.modules[__name__])
Expand Down

0 comments on commit f6fd7a7

Please sign in to comment.