[PECO-1440] Expose current query id on cursor object (#364)

* [PECO-1440] Expose current query id on cursor object

Signed-off-by: Levko Kravets <[email protected]>

* Clear `active_op_handle` when closing the cursor

Signed-off-by: Levko Kravets <[email protected]>

---------

Signed-off-by: Levko Kravets <[email protected]>

Author: kravets-levko

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]]:
         """

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__])