Skip to content

Commit ba8d9fd

Browse files
varun-edachali-dbxvarun-edachali-dbx
authored
Introduce models for SeaDatabricksClient (#595)
* [squash from exec-sea] bring over execution phase changes Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove excess test Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * add docstring Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remvoe exec func in sea backend Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove excess files Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove excess models Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove excess sea backend tests Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * cleanup Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * re-introduce get_schema_desc Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove SeaResultSet Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * clean imports and attributes Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * pass CommandId to ExecResp Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove changes in types Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * add back essential types (ExecResponse, from_sea_state) Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * fix fetch types Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * excess imports Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * reduce diff by maintaining logs Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * fix int test types Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * [squashed from exec-sea] init execution func Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove irrelevant changes Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove ResultSetFilter functionality Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove more irrelevant changes Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove more irrelevant changes Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * even more irrelevant changes Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove sea response as init option Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * exec test example scripts Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * formatting (black) Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * [squashed from sea-exec] merge sea stuffs Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove excess changes Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove excess removed docstring Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove excess changes in backend Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove excess imports Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove accidentally removed _get_schema_desc Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove unnecessary init with sea_response tests Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * rmeove unnecessary changes Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * formatting (black) Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * improved models and filters from cloudfetch-sea branch Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * filters stuff (align with JDBC) Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * backend from cloudfetch-sea Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove filtering, metadata ops Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * raise NotImplementedErrror for metadata ops Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove un-necessary changes Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove un-necessary backend cahnges Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove un-needed GetChunksResponse only relevant in Fetch phase Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * reduce code duplication Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * more clear docstrings Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * introduce strongly typed ChunkInfo Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove is_volume_operation from response Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * add is_volume_op and more ResultData fields Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> --------- Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 6d63df0 commit ba8d9fd

File tree

4 files changed

+363
-2
lines changed

4 files changed

+363
-2
lines changed

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,19 +4,49 @@
44
This package contains data models for SEA API requests and responses.
55
"""
66

7+
from databricks.sql.backend.sea.models.base import (
8+
ServiceError,
9+
StatementStatus,
10+
ExternalLink,
11+
ResultData,
12+
ColumnInfo,
13+
ResultManifest,
14+
)
15+
716
from databricks.sql.backend.sea.models.requests import (
17+
StatementParameter,
18+
ExecuteStatementRequest,
19+
GetStatementRequest,
20+
CancelStatementRequest,
21+
CloseStatementRequest,
822
CreateSessionRequest,
923
DeleteSessionRequest,
1024
)
1125

1226
from databricks.sql.backend.sea.models.responses import (
27+
ExecuteStatementResponse,
28+
GetStatementResponse,
1329
CreateSessionResponse,
1430
)
1531

1632
__all__ = [
33+
# Base models
34+
"ServiceError",
35+
"StatementStatus",
36+
"ExternalLink",
37+
"ResultData",
38+
"ColumnInfo",
39+
"ResultManifest",
1740
# Request models
41+
"StatementParameter",
42+
"ExecuteStatementRequest",
43+
"GetStatementRequest",
44+
"CancelStatementRequest",
45+
"CloseStatementRequest",
1846
"CreateSessionRequest",
1947
"DeleteSessionRequest",
2048
# Response models
49+
"ExecuteStatementResponse",
50+
"GetStatementResponse",
2151
"CreateSessionResponse",
2252
]

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

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
"""
2+
Base models for the SEA (Statement Execution API) backend.
3+
4+
These models define the common structures used in SEA API requests and responses.
5+
"""
6+
7+
from typing import Dict, List, Any, Optional, Union
8+
from dataclasses import dataclass, field
9+
10+
from databricks.sql.backend.types import CommandState
11+
12+
13+
@dataclass
14+
class ServiceError:
15+
"""Error information returned by the SEA API."""
16+
17+
message: str
18+
error_code: Optional[str] = None
19+
20+
21+
@dataclass
22+
class StatementStatus:
23+
"""Status information for a statement execution."""
24+
25+
state: CommandState
26+
error: Optional[ServiceError] = None
27+
sql_state: Optional[str] = None
28+
29+
30+
@dataclass
31+
class ExternalLink:
32+
"""External link information for result data."""
33+
34+
external_link: str
35+
expiration: str
36+
chunk_index: int
37+
byte_count: int = 0
38+
row_count: int = 0
39+
row_offset: int = 0
40+
next_chunk_index: Optional[int] = None
41+
next_chunk_internal_link: Optional[str] = None
42+
http_headers: Optional[Dict[str, str]] = None
43+
44+
45+
@dataclass
46+
class ChunkInfo:
47+
"""Information about a chunk in the result set."""
48+
49+
chunk_index: int
50+
byte_count: int
51+
row_offset: int
52+
row_count: int
53+
54+
55+
@dataclass
56+
class ResultData:
57+
"""Result data from a statement execution."""
58+
59+
data: Optional[List[List[Any]]] = None
60+
external_links: Optional[List[ExternalLink]] = None
61+
byte_count: Optional[int] = None
62+
chunk_index: Optional[int] = None
63+
next_chunk_index: Optional[int] = None
64+
next_chunk_internal_link: Optional[str] = None
65+
row_count: Optional[int] = None
66+
row_offset: Optional[int] = None
67+
attachment: Optional[bytes] = None
68+
69+
70+
@dataclass
71+
class ColumnInfo:
72+
"""Information about a column in the result set."""
73+
74+
name: str
75+
type_name: str
76+
type_text: str
77+
nullable: bool = True
78+
precision: Optional[int] = None
79+
scale: Optional[int] = None
80+
ordinal_position: Optional[int] = None
81+
82+
83+
@dataclass
84+
class ResultManifest:
85+
"""Manifest information for a result set."""
86+
87+
format: str
88+
schema: Dict[str, Any] # Will contain column information
89+
total_row_count: int
90+
total_byte_count: int
91+
total_chunk_count: int
92+
truncated: bool = False
93+
chunks: Optional[List[ChunkInfo]] = None
94+
result_compression: Optional[str] = None
95+
is_volume_operation: Optional[bool] = None

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

Lines changed: 96 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,99 @@
1-
from typing import Dict, Any, Optional
2-
from dataclasses import dataclass
1+
"""
2+
Request models for the SEA (Statement Execution API) backend.
3+
4+
These models define the structures used in SEA API requests.
5+
"""
6+
7+
from typing import Dict, List, Any, Optional, Union
8+
from dataclasses import dataclass, field
9+
10+
11+
@dataclass
12+
class StatementParameter:
13+
"""Representation of a parameter for a SQL statement."""
14+
15+
name: str
16+
value: Optional[str] = None
17+
type: Optional[str] = None
18+
19+
20+
@dataclass
21+
class ExecuteStatementRequest:
22+
"""Representation of a request to execute a SQL statement."""
23+
24+
session_id: str
25+
statement: str
26+
warehouse_id: str
27+
disposition: str = "EXTERNAL_LINKS"
28+
format: str = "JSON_ARRAY"
29+
result_compression: Optional[str] = None
30+
parameters: Optional[List[StatementParameter]] = None
31+
wait_timeout: str = "10s"
32+
on_wait_timeout: str = "CONTINUE"
33+
row_limit: Optional[int] = None
34+
35+
def to_dict(self) -> Dict[str, Any]:
36+
"""Convert the request to a dictionary for JSON serialization."""
37+
result: Dict[str, Any] = {
38+
"warehouse_id": self.warehouse_id,
39+
"session_id": self.session_id,
40+
"statement": self.statement,
41+
"disposition": self.disposition,
42+
"format": self.format,
43+
"wait_timeout": self.wait_timeout,
44+
"on_wait_timeout": self.on_wait_timeout,
45+
}
46+
47+
if self.row_limit is not None and self.row_limit > 0:
48+
result["row_limit"] = self.row_limit
49+
50+
if self.result_compression:
51+
result["result_compression"] = self.result_compression
52+
53+
if self.parameters:
54+
result["parameters"] = [
55+
{
56+
"name": param.name,
57+
**({"value": param.value} if param.value is not None else {}),
58+
**({"type": param.type} if param.type is not None else {}),
59+
}
60+
for param in self.parameters
61+
]
62+
63+
return result
64+
65+
66+
@dataclass
67+
class GetStatementRequest:
68+
"""Representation of a request to get information about a statement."""
69+
70+
statement_id: str
71+
72+
def to_dict(self) -> Dict[str, Any]:
73+
"""Convert the request to a dictionary for JSON serialization."""
74+
return {"statement_id": self.statement_id}
75+
76+
77+
@dataclass
78+
class CancelStatementRequest:
79+
"""Representation of a request to cancel a statement."""
80+
81+
statement_id: str
82+
83+
def to_dict(self) -> Dict[str, Any]:
84+
"""Convert the request to a dictionary for JSON serialization."""
85+
return {"statement_id": self.statement_id}
86+
87+
88+
@dataclass
89+
class CloseStatementRequest:
90+
"""Representation of a request to close a statement."""
91+
92+
statement_id: str
93+
94+
def to_dict(self) -> Dict[str, Any]:
95+
"""Convert the request to a dictionary for JSON serialization."""
96+
return {"statement_id": self.statement_id}
397

498

599
@dataclass

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

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,148 @@
1+
"""
2+
Response models for the SEA (Statement Execution API) backend.
3+
4+
These models define the structures used in SEA API responses.
5+
"""
6+
17
from typing import Dict, Any
28
from dataclasses import dataclass
39

10+
from databricks.sql.backend.types import CommandState
11+
from databricks.sql.backend.sea.models.base import (
12+
StatementStatus,
13+
ResultManifest,
14+
ResultData,
15+
ServiceError,
16+
ExternalLink,
17+
ChunkInfo,
18+
)
19+
20+
21+
def _parse_status(data: Dict[str, Any]) -> StatementStatus:
22+
"""Parse status from response data."""
23+
status_data = data.get("status", {})
24+
error = None
25+
if "error" in status_data:
26+
error_data = status_data["error"]
27+
error = ServiceError(
28+
message=error_data.get("message", ""),
29+
error_code=error_data.get("error_code"),
30+
)
31+
32+
state = CommandState.from_sea_state(status_data.get("state", ""))
33+
if state is None:
34+
raise ValueError(f"Invalid state: {status_data.get('state', '')}")
35+
36+
return StatementStatus(
37+
state=state,
38+
error=error,
39+
sql_state=status_data.get("sql_state"),
40+
)
41+
42+
43+
def _parse_manifest(data: Dict[str, Any]) -> ResultManifest:
44+
"""Parse manifest from response data."""
45+
46+
manifest_data = data.get("manifest", {})
47+
chunks = None
48+
if "chunks" in manifest_data:
49+
chunks = [
50+
ChunkInfo(
51+
chunk_index=chunk.get("chunk_index", 0),
52+
byte_count=chunk.get("byte_count", 0),
53+
row_offset=chunk.get("row_offset", 0),
54+
row_count=chunk.get("row_count", 0),
55+
)
56+
for chunk in manifest_data.get("chunks", [])
57+
]
58+
59+
return ResultManifest(
60+
format=manifest_data.get("format", ""),
61+
schema=manifest_data.get("schema", {}),
62+
total_row_count=manifest_data.get("total_row_count", 0),
63+
total_byte_count=manifest_data.get("total_byte_count", 0),
64+
total_chunk_count=manifest_data.get("total_chunk_count", 0),
65+
truncated=manifest_data.get("truncated", False),
66+
chunks=chunks,
67+
result_compression=manifest_data.get("result_compression"),
68+
is_volume_operation=manifest_data.get("is_volume_operation"),
69+
)
70+
71+
72+
def _parse_result(data: Dict[str, Any]) -> ResultData:
73+
"""Parse result data from response data."""
74+
result_data = data.get("result", {})
75+
external_links = None
76+
77+
if "external_links" in result_data:
78+
external_links = []
79+
for link_data in result_data["external_links"]:
80+
external_links.append(
81+
ExternalLink(
82+
external_link=link_data.get("external_link", ""),
83+
expiration=link_data.get("expiration", ""),
84+
chunk_index=link_data.get("chunk_index", 0),
85+
byte_count=link_data.get("byte_count", 0),
86+
row_count=link_data.get("row_count", 0),
87+
row_offset=link_data.get("row_offset", 0),
88+
next_chunk_index=link_data.get("next_chunk_index"),
89+
next_chunk_internal_link=link_data.get("next_chunk_internal_link"),
90+
http_headers=link_data.get("http_headers"),
91+
)
92+
)
93+
94+
return ResultData(
95+
data=result_data.get("data_array"),
96+
external_links=external_links,
97+
byte_count=result_data.get("byte_count"),
98+
chunk_index=result_data.get("chunk_index"),
99+
next_chunk_index=result_data.get("next_chunk_index"),
100+
next_chunk_internal_link=result_data.get("next_chunk_internal_link"),
101+
row_count=result_data.get("row_count"),
102+
row_offset=result_data.get("row_offset"),
103+
attachment=result_data.get("attachment"),
104+
)
105+
106+
107+
@dataclass
108+
class ExecuteStatementResponse:
109+
"""Representation of the response from executing a SQL statement."""
110+
111+
statement_id: str
112+
status: StatementStatus
113+
manifest: ResultManifest
114+
result: ResultData
115+
116+
@classmethod
117+
def from_dict(cls, data: Dict[str, Any]) -> "ExecuteStatementResponse":
118+
"""Create an ExecuteStatementResponse from a dictionary."""
119+
return cls(
120+
statement_id=data.get("statement_id", ""),
121+
status=_parse_status(data),
122+
manifest=_parse_manifest(data),
123+
result=_parse_result(data),
124+
)
125+
126+
127+
@dataclass
128+
class GetStatementResponse:
129+
"""Representation of the response from getting information about a statement."""
130+
131+
statement_id: str
132+
status: StatementStatus
133+
manifest: ResultManifest
134+
result: ResultData
135+
136+
@classmethod
137+
def from_dict(cls, data: Dict[str, Any]) -> "GetStatementResponse":
138+
"""Create a GetStatementResponse from a dictionary."""
139+
return cls(
140+
statement_id=data.get("statement_id", ""),
141+
status=_parse_status(data),
142+
manifest=_parse_manifest(data),
143+
result=_parse_result(data),
144+
)
145+
4146

5147
@dataclass
6148
class CreateSessionResponse:

0 commit comments

Comments
 (0)