fix: active op handle -> active command id in Cursor

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -578,6 +578,8 @@ def open_session(self, session_configuration, catalog, schema) -> SessionId:
             response = self.make_request(self._client.OpenSession, open_session_req)
             self._check_initial_namespace(catalog, schema, response)
             self._check_protocol_version(response)
+            if response.sessionHandle is None:
+                return None
             return SessionId.from_thrift_handle(response.sessionHandle)
         except:
             self._transport.close()
@@ -1072,7 +1074,9 @@ def get_columns(
         return self._handle_execute_response(resp, cursor)
 
     def _handle_execute_response(self, resp, cursor):
-        cursor.active_op_handle = resp.operationHandle
+        command_id = CommandId.from_thrift_handle(resp.operationHandle)
+
+        cursor.active_command_id = command_id
         self._check_direct_results_for_error(resp.directResults)
 
         final_operation_state = self._wait_until_command_done(
@@ -1083,12 +1087,12 @@ def _handle_execute_response(self, resp, cursor):
         execute_response = self._results_message_to_execute_response(
             resp, final_operation_state
         )
-        command_id = CommandId.from_thrift_handle(resp.operationHandle)
         execute_response = execute_response._replace(command_id=command_id)
         return execute_response
 
     def _handle_execute_response_async(self, resp, cursor):
-        cursor.active_op_handle = resp.operationHandle
+        command_id = CommandId.from_thrift_handle(resp.operationHandle)
+        cursor.active_command_id = command_id
         self._check_direct_results_for_error(resp.directResults)
 
     def fetch_results(

src/databricks/sql/client.py

@@ -403,7 +403,7 @@ def __init__(
         self.open = True
         self.executing_command_id = None
         self.backend = backend
-        self.active_op_handle = None
+        self.active_command_id = None
         self.escaper = ParamEscaper()
         self.lastrowid = None
 
@@ -862,7 +862,7 @@ def get_query_state(self) -> "TOperationState":
         :return:
         """
         self._check_not_closed()
-        return self.backend.get_query_state(self.active_op_handle)
+        return self.backend.get_query_state(self.active_command_id)
 
     def is_query_pending(self):
         """
@@ -893,7 +893,7 @@ def get_async_execution_result(self):
         operation_state = self.get_query_state()
         if operation_state == ttypes.TOperationState.FINISHED_STATE:
             execute_response = self.backend.get_execution_result(
-                self.active_op_handle, self
+                self.active_command_id, self
             )
             self.active_result_set = ResultSet(
                 self.connection,
@@ -1124,8 +1124,8 @@ def cancel(self) -> None:
         The command should be closed to free resources from the server.
         This method can be called from another thread.
         """
-        if self.active_op_handle is not None:
-            self.backend.cancel_command(self.active_op_handle)
+        if self.active_command_id is not None:
+            self.backend.cancel_command(self.active_command_id)
         else:
             logger.warning(
                 "Attempting to cancel a command, but there is no "
@@ -1137,9 +1137,9 @@ def close(self) -> None:
         self.open = False
 
         # Close active operation handle if it exists
-        if self.active_op_handle:
+        if self.active_command_id:
             try:
-                self.backend.close_command(self.active_op_handle)
+                self.backend.close_command(self.active_command_id)
             except RequestError as e:
                 if isinstance(e.args[1], CursorAlreadyClosedError):
                     logger.info("Operation was canceled by a prior request")
@@ -1148,7 +1148,7 @@ def close(self) -> None:
             except Exception as e:
                 logging.warning(f"Error closing operation handle: {e}")
             finally:
-                self.active_op_handle = None
+                self.active_command_id = None
 
         if self.active_result_set:
             self._close_and_clear_active_result_set()
@@ -1161,8 +1161,8 @@ def query_id(self) -> Optional[str]:
         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 self.active_op_handle.to_hex_id()
+        if self.active_command_id is not None:
+            return self.active_command_id.to_hex_id()
         return None
 
     @property

src/databricks/sql/ids.py

@@ -67,6 +67,9 @@ def from_thrift_handle(cls, session_handle):
         Returns:
             A SessionId instance
         """
+        if session_handle is None or session_handle.sessionId is None:
+            return None
+
         guid_bytes = session_handle.sessionId.guid
         secret_bytes = session_handle.sessionId.secret
 
@@ -171,6 +174,9 @@ def from_thrift_handle(cls, operation_handle):
         Returns:
             A CommandId instance
         """
+        if operation_handle is None or operation_handle.operationId is None:
+            return None
+
         guid_bytes = operation_handle.operationId.guid
         secret_bytes = operation_handle.operationId.secret