Skip to content

Commit d552695

Browse files
introduce unit tests for sea backend
Signed-off-by: varun-edachali-dbx <[email protected]>
1 parent ef09dbe commit d552695

File tree

4 files changed

+542
-14
lines changed

4 files changed

+542
-14
lines changed

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

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,14 @@ class DeleteSessionRequest:
139139
warehouse_id: str
140140
session_id: str
141141

142-
def to_dict(self) -> Dict[str, Any]:
143-
"""Convert the request to a dictionary for JSON serialization."""
144-
return {"warehouse_id": self.warehouse_id, "session_id": self.session_id}
142+
def to_query_params(self) -> Dict[str, str]:
143+
"""
144+
Convert the request to query parameters.
145+
146+
In the SEA API, only the warehouse_id is sent as a query parameter.
147+
The session_id is included in the URL path.
148+
149+
Returns:
150+
A dictionary containing the warehouse_id as a query parameter
151+
"""
152+
return {"warehouse_id": self.warehouse_id}

src/databricks/sql/backend/sea_backend.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -222,7 +222,7 @@ def close_session(self, session_id: SessionId) -> None:
222222
self.http_client._make_request(
223223
method="DELETE",
224224
path=self.SESSION_PATH_WITH_ID.format(sea_session_id),
225-
data=request.to_dict(),
225+
data=request.to_query_params(), # Use to_query_params to get only the warehouse_id
226226
)
227227

228228
def execute_command(
@@ -338,7 +338,7 @@ def execute_command(
338338

339339
# Check for errors
340340
if state == "FAILED" and status.error:
341-
error_message = status.error.message
341+
error_message = status.error["message"]
342342
raise Error(f"Statement execution failed: {error_message}")
343343

344344
# Check for cancellation

src/databricks/sql/backend/sea_result_set.py

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ def __init__(
5858
if manifest_data:
5959
schema_data = manifest_data.get("schema", [])
6060
columns = []
61-
for col_data in schema_data.get("columns", []):
61+
for col_data in schema_data:
6262
columns.append(
6363
ColumnInfo(
6464
name=col_data.get("name", ""),
@@ -142,28 +142,35 @@ def _extract_description_from_manifest(
142142

143143
def _fill_results_buffer(self) -> None:
144144
"""Fill the results buffer from the backend."""
145-
raise NotImplementedError
145+
raise NotImplementedError("Not implemented yet")
146146

147147
def fetchone(self) -> Optional[Row]:
148148
"""Fetch the next row of a query result set."""
149-
raise NotImplementedError
149+
raise NotImplementedError("Not implemented yet")
150150

151151
def fetchmany(self, size: int) -> List[Row]:
152152
"""Fetch the next set of rows of a query result."""
153-
raise NotImplementedError
153+
raise NotImplementedError("Not implemented yet")
154154

155155
def fetchall(self) -> List[Row]:
156156
"""Fetch all remaining rows of a query result."""
157-
raise NotImplementedError
157+
raise NotImplementedError("Not implemented yet")
158158

159159
def fetchmany_arrow(self, size: int) -> Any:
160160
"""Fetch the next set of rows as an Arrow table."""
161-
raise NotImplementedError
161+
raise NotImplementedError("Not implemented yet")
162162

163163
def fetchall_arrow(self) -> Any:
164164
"""Fetch all remaining rows as an Arrow table."""
165-
raise NotImplementedError
165+
raise NotImplementedError("Not implemented yet")
166166

167167
def close(self) -> None:
168168
"""Close the result set and release any resources."""
169-
logger.info("Closing SeaResultSet (but not really)")
169+
# Basic implementation to close the statement
170+
if self.connection.open and self.statement_id:
171+
try:
172+
self.backend.close_command(
173+
CommandId.from_sea_statement_id(self.statement_id)
174+
)
175+
except Exception as e:
176+
logger.warning(f"Error closing SEA statement: {e}")

0 commit comments

Comments
 (0)