diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 17369b0a..ad5146d1 100644 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -6,6 +6,7 @@ import json import os import decimal +from uuid import UUID from databricks.sql import __version__ from databricks.sql import * @@ -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]]: """ diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 9e1a66c7..72f8e7d5 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -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 @@ -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__])