introduce models for requests and responses

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

Author: varun-edachali-dbx

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

@@ -0,0 +1,22 @@
+"""
+Models for the SEA (Statement Execution API) backend.
+
+This package contains data models for SEA API requests and responses.
+"""
+
+from databricks.sql.backend.models.requests import (
+    CreateSessionRequest,
+    DeleteSessionRequest,
+)
+
+from databricks.sql.backend.models.responses import (
+    CreateSessionResponse,
+)
+
+__all__ = [
+    # Request models
+    "CreateSessionRequest",
+    "DeleteSessionRequest",
+    # Response models
+    "CreateSessionResponse",
+]

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

@@ -0,0 +1,39 @@
+from typing import Dict, Any, Optional
+from dataclasses import dataclass
+
+
+@dataclass
+class CreateSessionRequest:
+    """Request to create a new session."""
+
+    warehouse_id: str
+    session_confs: Optional[Dict[str, str]] = None
+    catalog: Optional[str] = None
+    schema: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the request to a dictionary for JSON serialization."""
+        result: Dict[str, Any] = {"warehouse_id": self.warehouse_id}
+
+        if self.session_confs:
+            result["session_confs"] = self.session_confs
+
+        if self.catalog:
+            result["catalog"] = self.catalog
+
+        if self.schema:
+            result["schema"] = self.schema
+
+        return result
+
+
+@dataclass
+class DeleteSessionRequest:
+    """Request to delete a session."""
+
+    warehouse_id: str
+    session_id: str
+
+    def to_dict(self) -> Dict[str, str]:
+        """Convert the request to a dictionary for JSON serialization."""
+        return {"warehouse_id": self.warehouse_id, "session_id": self.session_id}

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

@@ -0,0 +1,14 @@
+from typing import Dict, Any
+from dataclasses import dataclass
+
+
+@dataclass
+class CreateSessionResponse:
+    """Response from creating a new session."""
+
+    session_id: str
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "CreateSessionResponse":
+        """Create a CreateSessionResponse from a dictionary."""
+        return cls(session_id=data.get("session_id", ""))

src/databricks/sql/backend/sea_backend.py

@@ -7,11 +7,17 @@
 
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.backend.types import SessionId, CommandId, CommandState, BackendType
-from databricks.sql.exc import Error, NotSupportedError
+from databricks.sql.exc import Error, NotSupportedError, ServerOperationError
 from databricks.sql.backend.utils.http_client import CustomHttpClient
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.types import SSLOptions
 
+from databricks.sql.backend.models import (
+    CreateSessionRequest,
+    DeleteSessionRequest,
+    CreateSessionResponse,
+)
+
 logger = logging.getLogger(__name__)
 
 
@@ -163,21 +169,27 @@ def open_session(
             schema,
         )
 
-        request_data: Dict[str, Any] = {"warehouse_id": self.warehouse_id}
-        if session_configuration:
-            request_data["session_confs"] = session_configuration
-        if catalog:
-            request_data["catalog"] = catalog
-        if schema:
-            request_data["schema"] = schema
+        request_data = CreateSessionRequest(
+            warehouse_id=self.warehouse_id,
+            session_confs=session_configuration,
+            catalog=catalog,
+            schema=schema,
+        )
 
         response = self.http_client._make_request(
-            method="POST", path=self.SESSION_PATH, data=request_data
+            method="POST", path=self.SESSION_PATH, data=request_data.to_dict()
         )
 
-        session_id = response.get("session_id")
+        session_response = CreateSessionResponse.from_dict(response)
+        session_id = session_response.session_id
         if not session_id:
-            raise Error("Failed to create session: No session ID returned")
+            raise ServerOperationError(
+                "Failed to create session: No session ID returned",
+                {
+                    "operation-id": None,
+                    "diagnostic-info": None,
+                },
+            )
 
         return SessionId.from_sea_session_id(session_id)
 
@@ -199,12 +211,15 @@ def close_session(self, session_id: SessionId) -> None:
             raise ValueError("Not a valid SEA session ID")
         sea_session_id = session_id.to_sea_session_id()
 
-        request_data = {"warehouse_id": self.warehouse_id}
+        request_data = DeleteSessionRequest(
+            warehouse_id=self.warehouse_id,
+            session_id=sea_session_id,
+        )
 
         self.http_client._make_request(
             method="DELETE",
             path=self.SESSION_PATH_WITH_ID.format(sea_session_id),
-            data=request_data,
+            data=request_data.to_dict(),
         )
 
     # == Not Implemented Operations ==

tests/unit/test_sea_backend.py

@@ -150,7 +150,7 @@ def test_close_session_valid_id(self, sea_client, mock_http_client):
         mock_http_client._make_request.assert_called_once_with(
             method="DELETE",
             path=sea_client.SESSION_PATH_WITH_ID.format("test-session-789"),
-            data={"warehouse_id": "abc123"},
+            data={"session_id": "test-session-789", "warehouse_id": "abc123"},
         )
 
     def test_close_session_invalid_id_type(self, sea_client):