introduce __str__ methods for CommandId and SessionId

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

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -1153,7 +1153,7 @@ def cancel_command(self, command_id: CommandId) -> None:
         if not thrift_handle:
             raise ValueError("Not a valid Thrift command ID")
 
-        logger.debug("Cancelling command {}".format(command_id.guid))
+        logger.debug("Cancelling command {}".format(guid_to_hex_id(command_id.guid)))
         req = ttypes.TCancelOperationReq(thrift_handle)
         self.make_request(self._client.CancelOperation, req)
 

src/databricks/sql/backend/types.py

@@ -59,6 +59,22 @@ def __init__(
         self.secret = secret
         self.info = info or {}
 
+    def __str__(self) -> str:
+        """
+        Return a string representation of the SessionId.
+
+        For SEA backend, returns the guid.
+        For Thrift backend, returns a format like "guid|secret".
+
+        Returns:
+            A string representation of the session ID
+        """
+        if self.backend_type == BackendType.SEA:
+            return str(self.guid)
+        elif self.backend_type == BackendType.THRIFT:
+            return f"{self.to_hex_id()}|{guid_to_hex_id(self.secret) if isinstance(self.secret, bytes) else str(self.secret)}"
+        return str(self.guid)
+
     @classmethod
     def from_thrift_handle(cls, session_handle, info: Optional[Dict[str, Any]] = None):
         """
@@ -185,6 +201,22 @@ def __init__(
         self.has_result_set = has_result_set
         self.modified_row_count = modified_row_count
 
+    def __str__(self) -> str:
+        """
+        Return a string representation of the CommandId.
+
+        For SEA backend, returns the guid.
+        For Thrift backend, returns a format like "guid|secret".
+
+        Returns:
+            A string representation of the command ID
+        """
+        if self.backend_type == BackendType.SEA:
+            return str(self.guid)
+        elif self.backend_type == BackendType.THRIFT:
+            return f"{self.to_hex_id()}|{guid_to_hex_id(self.secret) if isinstance(self.secret, bytes) else str(self.secret)}"
+        return str(self.guid)
+
     @classmethod
     def from_thrift_handle(cls, operation_handle):
         """

src/databricks/sql/client.py

@@ -1256,7 +1256,7 @@ def __iter__(self):
 
     def _fill_results_buffer(self):
         if not isinstance(self.backend, ThriftDatabricksClient):
-            # This specific logic is for Thrift. SEA will have its own way.
+            # currently, we are assuming only the Thrift backend exists
             raise NotImplementedError(
                 "Fetching further result batches is currently only implemented for the Thrift backend."
             )