introduce unit tests for sea backend

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

Author: varun-edachali-dbx

src/databricks/sql/backend/models/requests.py

@@ -139,6 +139,14 @@ class DeleteSessionRequest:
     warehouse_id: str
     session_id: str
 
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert the request to a dictionary for JSON serialization."""
-        return {"warehouse_id": self.warehouse_id, "session_id": self.session_id}
+    def to_query_params(self) -> Dict[str, str]:
+        """
+        Convert the request to query parameters.
+
+        In the SEA API, only the warehouse_id is sent as a query parameter.
+        The session_id is included in the URL path.
+
+        Returns:
+            A dictionary containing the warehouse_id as a query parameter
+        """
+        return {"warehouse_id": self.warehouse_id}

src/databricks/sql/backend/sea_backend.py

@@ -222,7 +222,7 @@ def close_session(self, session_id: SessionId) -> None:
         self.http_client._make_request(
             method="DELETE",
             path=self.SESSION_PATH_WITH_ID.format(sea_session_id),
-            data=request.to_dict(),
+            data=request.to_query_params(),  # Use to_query_params to get only the warehouse_id
         )
 
     def execute_command(
@@ -338,7 +338,7 @@ def execute_command(
 
             # Check for errors
             if state == "FAILED" and status.error:
-                error_message = status.error.message
+                error_message = status.error["message"]
                 raise Error(f"Statement execution failed: {error_message}")
 
             # Check for cancellation

src/databricks/sql/backend/sea_result_set.py

@@ -58,7 +58,7 @@ def __init__(
         if manifest_data:
             schema_data = manifest_data.get("schema", [])
             columns = []
-            for col_data in schema_data.get("columns", []):
+            for col_data in schema_data:
                 columns.append(
                     ColumnInfo(
                         name=col_data.get("name", ""),
@@ -142,28 +142,35 @@ def _extract_description_from_manifest(
 
     def _fill_results_buffer(self) -> None:
         """Fill the results buffer from the backend."""
-        raise NotImplementedError
+        raise NotImplementedError("Not implemented yet")
 
     def fetchone(self) -> Optional[Row]:
         """Fetch the next row of a query result set."""
-        raise NotImplementedError
+        raise NotImplementedError("Not implemented yet")
 
     def fetchmany(self, size: int) -> List[Row]:
         """Fetch the next set of rows of a query result."""
-        raise NotImplementedError
+        raise NotImplementedError("Not implemented yet")
 
     def fetchall(self) -> List[Row]:
         """Fetch all remaining rows of a query result."""
-        raise NotImplementedError
+        raise NotImplementedError("Not implemented yet")
 
     def fetchmany_arrow(self, size: int) -> Any:
         """Fetch the next set of rows as an Arrow table."""
-        raise NotImplementedError
+        raise NotImplementedError("Not implemented yet")
 
     def fetchall_arrow(self) -> Any:
         """Fetch all remaining rows as an Arrow table."""
-        raise NotImplementedError
+        raise NotImplementedError("Not implemented yet")
 
     def close(self) -> None:
         """Close the result set and release any resources."""
-        logger.info("Closing SeaResultSet (but not really)")
+        # Basic implementation to close the statement
+        if self.connection.open and self.statement_id:
+            try:
+                self.backend.close_command(
+                    CommandId.from_sea_statement_id(self.statement_id)
+                )
+            except Exception as e:
+                logger.warning(f"Error closing SEA statement: {e}")