From f6fd7a7956a4dbc78ad36b5e079fe8d74176a0f1 Mon Sep 17 00:00:00 2001 From: Levko Kravets Date: Mon, 4 Mar 2024 20:20:14 +0200 Subject: [PATCH] [PECO-1440] Expose current query id on cursor object (#364) * [PECO-1440] Expose current query id on cursor object Signed-off-by: Levko Kravets * Clear `active_op_handle` when closing the cursor Signed-off-by: Levko Kravets --------- Signed-off-by: Levko Kravets --- src/databricks/sql/client.py | 14 ++++++++++++++ tests/unit/test_client.py | 21 +++++++++++++++++++++ 2 files changed, 35 insertions(+) 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__])