introduce models

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

Author: varun-edachali-dbx

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

@@ -0,0 +1,54 @@
+"""
+Models for the SEA (Statement Execution API) backend.
+
+This package contains data models for SEA API requests and responses.
+"""
+
+from databricks.sql.backend.models.base import (
+    ServiceError,
+    StatementStatus,
+    ExternalLink,
+    ResultData,
+    ColumnInfo,
+    ResultManifest
+)
+
+from databricks.sql.backend.models.requests import (
+    StatementParameter,
+    ExecuteStatementRequest,
+    GetStatementRequest,
+    CancelStatementRequest,
+    CloseStatementRequest,
+    CreateSessionRequest,
+    DeleteSessionRequest
+)
+
+from databricks.sql.backend.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/models/base.py

@@ -0,0 +1,60 @@
+"""
+Base models for the SEA (Statement Execution API) backend.
+
+These models define the common structures used in SEA API requests and responses.
+"""
+
+from typing import Dict, List, Any, Optional, Union
+from dataclasses import dataclass, field
+
+
+@dataclass
+class ServiceError:
+    """Error information returned by the SEA API."""
+    message: str
+    error_code: Optional[str] = None
+
+
+@dataclass
+class StatementStatus:
+    """Status information for a statement execution."""
+    state: str
+    error: Optional[ServiceError] = None
+    sql_state: Optional[str] = None
+
+
+@dataclass
+class ExternalLink:
+    """External link information for result data."""
+    external_link: str
+    expiration: str
+    chunk_index: int
+
+
+@dataclass
+class ResultData:
+    """Result data from a statement execution."""
+    data: Optional[List[List[Any]]] = None
+    external_links: Optional[List[ExternalLink]] = None
+
+
+@dataclass
+class ColumnInfo:
+    """Information about a column in the result set."""
+    name: str
+    type_name: str
+    type_text: str
+    nullable: bool = True
+    precision: Optional[int] = None
+    scale: Optional[int] = None
+    ordinal_position: Optional[int] = None
+
+
+@dataclass
+class ResultManifest:
+    """Manifest information for a result set."""
+    schema: List[ColumnInfo]
+    total_row_count: int
+    total_byte_count: int
+    truncated: bool = False
+    chunk_count: Optional[int] = None

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

@@ -0,0 +1,149 @@
+"""
+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
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the request to a dictionary for JSON serialization."""
+        result = {
+            "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.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."""
+    warehouse_id: str
+    statement_id: str
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the request to a dictionary for JSON serialization."""
+        return {
+            "warehouse_id": self.warehouse_id,
+            "statement_id": self.statement_id
+        }
+
+
+@dataclass
+class CancelStatementRequest:
+    """Request to cancel a statement."""
+    warehouse_id: str
+    statement_id: str
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the request to a dictionary for JSON serialization."""
+        return {
+            "warehouse_id": self.warehouse_id,
+            "statement_id": self.statement_id
+        }
+
+
+@dataclass
+class CloseStatementRequest:
+    """Request to close a statement."""
+    warehouse_id: str
+    statement_id: str
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the request to a dictionary for JSON serialization."""
+        return {
+            "warehouse_id": self.warehouse_id,
+            "statement_id": self.statement_id
+        }
+
+
+@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 = {"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, Any]:
+        """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,81 @@
+"""
+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.models.base import (
+    StatementStatus, ResultManifest, ResultData
+)
+
+
+@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", {})
+        status = StatementStatus(
+            state=status_data.get("state", ""),
+            error=None if "error" not in status_data else {
+                "message": status_data["error"].get("message", ""),
+                "error_code": status_data["error"].get("error_code")
+            },
+            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", {})
+        status = StatementStatus(
+            state=status_data.get("state", ""),
+            error=None if "error" not in status_data else {
+                "message": status_data["error"].get("message", ""),
+                "error_code": status_data["error"].get("error_code")
+            },
+            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 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", ""))