Merge branch 'sea-migration' into sea-test-scripts

Author: varun-edachali-dbx

CHANGELOG.md

@@ -1,5 +1,14 @@
 # Release History
 
+# 4.0.4 (2025-06-16)
+
+- Update thrift client library after cleaning up unused fields and structs (databricks/databricks-sql-python#553 by @vikrantpuppala)
+- Refactor decimal conversion in PyArrow tables to use direct casting (databricks/databricks-sql-python#544 by @jayantsing-db)
+- Fix: `fetchall_arrow` to always return results in `arrow` format (databricks/databricks-sql-python#551 by @shivam2680)
+- Enhance cursor close handling and context manager exception management to prevent server side resource leaks (databricks/databricks-sql-python#554 by @madhav-db)
+- Added additional logging to enhance debugging (databricks/databricks-sql-python#556 by @saishreeeee)
+- Feature: Added support for complex data types such as Arrays and Map [Private Preview] (databricks/databricks-sql-python#559 by @jprakash-db)
+
 # 4.0.3 (2025-04-22)
 
 - Fix: Removed `packaging` dependency in favour of default libraries, for `urllib3` version checks (databricks/databricks-sql-python#547 by @jprakash-db)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "databricks-sql-connector"
-version = "4.0.3"
+version = "4.0.4"
 description = "Databricks SQL Connector for Python"
 authors = ["Databricks <databricks-sql-connector-maintainers@databricks.com>"]
 license = "Apache-2.0"

src/databricks/sql/__init__.py

@@ -68,7 +68,7 @@ def __repr__(self):
 DATE = DBAPITypeObject("date")
 ROWID = DBAPITypeObject()
 
-__version__ = "4.0.3"
+__version__ = "4.0.4"
 USER_AGENT_NAME = "PyDatabricksSqlConnector"
 
 # These two functions are pyhive legacy

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

@@ -4,19 +4,49 @@
 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

@@ -0,0 +1,95 @@
+"""
+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
+
+from databricks.sql.backend.types import CommandState
+
+
+@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: CommandState
+    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
+    byte_count: int = 0
+    row_count: int = 0
+    row_offset: int = 0
+    next_chunk_index: Optional[int] = None
+    next_chunk_internal_link: Optional[str] = None
+    http_headers: Optional[Dict[str, str]] = None
+
+
+@dataclass
+class ChunkInfo:
+    """Information about a chunk in the result set."""
+
+    chunk_index: int
+    byte_count: int
+    row_offset: int
+    row_count: int
+
+
+@dataclass
+class ResultData:
+    """Result data from a statement execution."""
+
+    data: Optional[List[List[Any]]] = None
+    external_links: Optional[List[ExternalLink]] = None
+    byte_count: Optional[int] = None
+    chunk_index: Optional[int] = None
+    next_chunk_index: Optional[int] = None
+    next_chunk_internal_link: Optional[str] = None
+    row_count: Optional[int] = None
+    row_offset: Optional[int] = None
+    attachment: Optional[bytes] = 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."""
+
+    format: str
+    schema: Dict[str, Any]  # Will contain column information
+    total_row_count: int
+    total_byte_count: int
+    total_chunk_count: int
+    truncated: bool = False
+    chunks: Optional[List[ChunkInfo]] = None
+    result_compression: Optional[str] = None
+    is_volume_operation: Optional[bool] = None

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

@@ -1,5 +1,99 @@
-from typing import Dict, Any, Optional
-from dataclasses import dataclass
+"""
+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:
+    """Representation of a parameter for a SQL statement."""
+
+    name: str
+    value: Optional[str] = None
+    type: Optional[str] = None
+
+
+@dataclass
+class ExecuteStatementRequest:
+    """Representation of a request to execute a SQL statement."""
+
+    session_id: str
+    statement: str
+    warehouse_id: str
+    disposition: str = "EXTERNAL_LINKS"
+    format: str = "JSON_ARRAY"
+    result_compression: Optional[str] = None
+    parameters: Optional[List[StatementParameter]] = None
+    wait_timeout: str = "10s"
+    on_wait_timeout: str = "CONTINUE"
+    row_limit: Optional[int] = 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.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:
+    """Representation of a 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:
+    """Representation of a 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:
+    """Representation of a 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}
 
 
 @dataclass