remove excess models

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

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

@@ -4,49 +4,19 @@
 This package contains data models for SEA API requests and responses.
 """
 
-from databricks.sql.backend.sea.models.base import (
-    ServiceError,
-    StatementStatus,
-    ExternalLink,
-    ResultData,
-    ColumnInfo,
-    ResultManifest,
-)
-
 from databricks.sql.backend.sea.models.requests import (
-    StatementParameter,
-    ExecuteStatementRequest,
-    GetStatementRequest,
-    CancelStatementRequest,
-    CloseStatementRequest,
     CreateSessionRequest,
     DeleteSessionRequest,
 )
 
 from databricks.sql.backend.sea.models.responses import (
-    ExecuteStatementResponse,
-    GetStatementResponse,
     CreateSessionResponse,
 )
 
 __all__ = [
-    # Base models
-    "ServiceError",
-    "StatementStatus",
-    "ExternalLink",
-    "ResultData",
-    "ColumnInfo",
-    "ResultManifest",
     # Request models
-    "StatementParameter",
-    "ExecuteStatementRequest",
-    "GetStatementRequest",
-    "CancelStatementRequest",
-    "CloseStatementRequest",
     "CreateSessionRequest",
     "DeleteSessionRequest",
     # Response models
-    "ExecuteStatementResponse",
-    "GetStatementResponse",
     "CreateSessionResponse",
 ]

src/databricks/sql/backend/sea/models/base.py

@@ -1,111 +1,5 @@
-"""
-Request models for the SEA (Statement Execution API) backend.
-
-These models define the structures used in SEA API requests.
-"""
-
-from typing import Dict, List, Any, Optional, Union
-from dataclasses import dataclass, field
-
-
-@dataclass
-class StatementParameter:
-    """Parameter for a SQL statement."""
-
-    name: str
-    value: Optional[str] = None
-    type: Optional[str] = None
-
-
-@dataclass
-class ExecuteStatementRequest:
-    """Request to execute a SQL statement."""
-
-    warehouse_id: str
-    statement: str
-    session_id: str
-    disposition: str = "EXTERNAL_LINKS"
-    format: str = "JSON_ARRAY"
-    wait_timeout: str = "10s"
-    on_wait_timeout: str = "CONTINUE"
-    row_limit: Optional[int] = None
-    byte_limit: Optional[int] = None
-    parameters: Optional[List[StatementParameter]] = None
-    catalog: Optional[str] = None
-    schema: Optional[str] = None
-    result_compression: 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,
-            "session_id": self.session_id,
-            "statement": self.statement,
-            "disposition": self.disposition,
-            "format": self.format,
-            "wait_timeout": self.wait_timeout,
-            "on_wait_timeout": self.on_wait_timeout,
-        }
-
-        if self.row_limit is not None and self.row_limit > 0:
-            result["row_limit"] = self.row_limit
-
-        if self.byte_limit is not None and self.byte_limit > 0:
-            result["byte_limit"] = self.byte_limit
-
-        if self.catalog:
-            result["catalog"] = self.catalog
-
-        if self.schema:
-            result["schema"] = self.schema
-
-        if self.result_compression:
-            result["result_compression"] = self.result_compression
-
-        if self.parameters:
-            result["parameters"] = [
-                {
-                    "name": param.name,
-                    **({"value": param.value} if param.value is not None else {}),
-                    **({"type": param.type} if param.type is not None else {}),
-                }
-                for param in self.parameters
-            ]
-
-        return result
-
-
-@dataclass
-class GetStatementRequest:
-    """Request to get information about a statement."""
-
-    statement_id: str
-
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert the request to a dictionary for JSON serialization."""
-        return {"statement_id": self.statement_id}
-
-
-@dataclass
-class CancelStatementRequest:
-    """Request to cancel a statement."""
-
-    statement_id: str
-
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert the request to a dictionary for JSON serialization."""
-        return {"statement_id": self.statement_id}
-
-
-@dataclass
-class CloseStatementRequest:
-    """Request to close a statement."""
-
-    statement_id: str
-
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert the request to a dictionary for JSON serialization."""
-        return {"statement_id": self.statement_id}
+from typing import Dict, Any, Optional
+from dataclasses import dataclass
 
 
 @dataclass

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

@@ -1,96 +1,5 @@
-"""
-Response models for the SEA (Statement Execution API) backend.
-
-These models define the structures used in SEA API responses.
-"""
-
-from typing import Dict, List, Any, Optional, Union
-from dataclasses import dataclass, field
-
-from databricks.sql.backend.types import CommandState
-from databricks.sql.backend.sea.models.base import (
-    StatementStatus,
-    ResultManifest,
-    ResultData,
-    ServiceError,
-)
-
-
-@dataclass
-class ExecuteStatementResponse:
-    """Response from executing a SQL statement."""
-
-    statement_id: str
-    status: StatementStatus
-    manifest: Optional[ResultManifest] = None
-    result: Optional[ResultData] = None
-
-    @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "ExecuteStatementResponse":
-        """Create an ExecuteStatementResponse from a dictionary."""
-        status_data = data.get("status", {})
-        error = None
-        if "error" in status_data:
-            error_data = status_data["error"]
-            error = ServiceError(
-                message=error_data.get("message", ""),
-                error_code=error_data.get("error_code"),
-            )
-
-        state = CommandState.from_sea_state(status_data.get("state", ""))
-        if state is None:
-            raise ValueError(f"Invalid state: {status_data.get('state', '')}")
-        status = StatementStatus(
-            state=state,
-            error=error,
-            sql_state=status_data.get("sql_state"),
-        )
-
-        return cls(
-            statement_id=data.get("statement_id", ""),
-            status=status,
-            manifest=data.get("manifest"),  # We'll parse this more fully if needed
-            result=data.get("result"),  # We'll parse this more fully if needed
-        )
-
-
-@dataclass
-class GetStatementResponse:
-    """Response from getting information about a statement."""
-
-    statement_id: str
-    status: StatementStatus
-    manifest: Optional[ResultManifest] = None
-    result: Optional[ResultData] = None
-
-    @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "GetStatementResponse":
-        """Create a GetStatementResponse from a dictionary."""
-        status_data = data.get("status", {})
-        error = None
-        if "error" in status_data:
-            error_data = status_data["error"]
-            error = ServiceError(
-                message=error_data.get("message", ""),
-                error_code=error_data.get("error_code"),
-            )
-
-        state = CommandState.from_sea_state(status_data.get("state", ""))
-        if state is None:
-            raise ValueError(f"Invalid state: {status_data.get('state', '')}")
-
-        status = StatementStatus(
-            state=state,
-            error=error,
-            sql_state=status_data.get("sql_state"),
-        )
-
-        return cls(
-            statement_id=data.get("statement_id", ""),
-            status=status,
-            manifest=data.get("manifest"),  # We'll parse this more fully if needed
-            result=data.get("result"),  # We'll parse this more fully if needed
-        )
+from typing import Dict, Any
+from dataclasses import dataclass
 
 
 @dataclass