Skip to content

Commit ab2e43d

Browse files
varun-edachali-dbxvarun-edachali-dbx
committed
Revert "Complete Fetch Phase (EXTERNAL_LINKS disposition and ARROW format) (#598)"
This reverts commit 1a0575a.
1 parent 0d6b53c commit ab2e43d

17 files changed

+331
-1294
lines changed

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

Lines changed: 16 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import re
66
from typing import Any, Dict, Tuple, List, Optional, Union, TYPE_CHECKING, Set
77

8-
from databricks.sql.backend.sea.models.base import ExternalLink, ResultManifest
8+
from databricks.sql.backend.sea.models.base import ResultManifest
99
from databricks.sql.backend.sea.utils.constants import (
1010
ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
1111
ResultFormat,
@@ -29,7 +29,7 @@
2929
BackendType,
3030
ExecuteResponse,
3131
)
32-
from databricks.sql.exc import DatabaseError, ServerOperationError
32+
from databricks.sql.exc import DatabaseError, ProgrammingError, ServerOperationError
3333
from databricks.sql.backend.sea.utils.http_client import SeaHttpClient
3434
from databricks.sql.types import SSLOptions
3535

@@ -45,7 +45,6 @@
4545
GetStatementResponse,
4646
CreateSessionResponse,
4747
)
48-
from databricks.sql.backend.sea.models.responses import GetChunksResponse
4948

5049
logger = logging.getLogger(__name__)
5150

@@ -90,7 +89,6 @@ class SeaDatabricksClient(DatabricksClient):
9089
STATEMENT_PATH = BASE_PATH + "statements"
9190
STATEMENT_PATH_WITH_ID = STATEMENT_PATH + "/{}"
9291
CANCEL_STATEMENT_PATH_WITH_ID = STATEMENT_PATH + "/{}/cancel"
93-
CHUNK_PATH_WITH_ID_AND_INDEX = STATEMENT_PATH + "/{}/result/chunks/{}"
9492

9593
# SEA constants
9694
POLL_INTERVAL_SECONDS = 0.2
@@ -137,13 +135,13 @@ def __init__(
137135
self.warehouse_id = self._extract_warehouse_id(http_path)
138136

139137
# Initialize HTTP client
140-
self._http_client = SeaHttpClient(
138+
self.http_client = SeaHttpClient(
141139
server_hostname=server_hostname,
142140
port=port,
143141
http_path=http_path,
144142
http_headers=http_headers,
145143
auth_provider=auth_provider,
146-
ssl_options=self._ssl_options,
144+
ssl_options=ssl_options,
147145
**kwargs,
148146
)
149147

@@ -182,7 +180,7 @@ def _extract_warehouse_id(self, http_path: str) -> str:
182180
f"Note: SEA only works for warehouses."
183181
)
184182
logger.error(error_message)
185-
raise ValueError(error_message)
183+
raise ProgrammingError(error_message)
186184

187185
@property
188186
def max_download_threads(self) -> int:
@@ -229,7 +227,7 @@ def open_session(
229227
schema=schema,
230228
)
231229

232-
response = self._http_client._make_request(
230+
response = self.http_client._make_request(
233231
method="POST", path=self.SESSION_PATH, data=request_data.to_dict()
234232
)
235233

@@ -254,7 +252,7 @@ def close_session(self, session_id: SessionId) -> None:
254252
session_id: The session identifier returned by open_session()
255253
256254
Raises:
257-
ValueError: If the session ID is invalid
255+
ProgrammingError: If the session ID is invalid
258256
OperationalError: If there's an error closing the session
259257
"""
260258

@@ -269,7 +267,7 @@ def close_session(self, session_id: SessionId) -> None:
269267
session_id=sea_session_id,
270268
)
271269

272-
self._http_client._make_request(
270+
self.http_client._make_request(
273271
method="DELETE",
274272
path=self.SESSION_PATH_WITH_ID.format(sea_session_id),
275273
data=request_data.to_dict(),
@@ -351,7 +349,7 @@ def _results_message_to_execute_response(
351349

352350
# Check for compression
353351
lz4_compressed = (
354-
response.manifest.result_compression == ResultCompression.LZ4_FRAME.value
352+
response.manifest.result_compression == ResultCompression.LZ4_FRAME
355353
)
356354

357355
execute_response = ExecuteResponse(
@@ -457,7 +455,7 @@ def execute_command(
457455
enforce_embedded_schema_correctness: Whether to enforce schema correctness
458456
459457
Returns:
460-
SeaResultSet: A SeaResultSet instance for the executed command
458+
ResultSet: A SeaResultSet instance for the executed command
461459
"""
462460

463461
if session_id.backend_type != BackendType.SEA:
@@ -508,7 +506,7 @@ def execute_command(
508506
result_compression=result_compression,
509507
)
510508

511-
response_data = self._http_client._make_request(
509+
response_data = self.http_client._make_request(
512510
method="POST", path=self.STATEMENT_PATH, data=request.to_dict()
513511
)
514512
response = ExecuteStatementResponse.from_dict(response_data)
@@ -545,7 +543,7 @@ def cancel_command(self, command_id: CommandId) -> None:
545543
command_id: Command identifier to cancel
546544
547545
Raises:
548-
ValueError: If the command ID is invalid
546+
ProgrammingError: If the command ID is invalid
549547
"""
550548

551549
if command_id.backend_type != BackendType.SEA:
@@ -556,7 +554,7 @@ def cancel_command(self, command_id: CommandId) -> None:
556554
raise ValueError("Not a valid SEA command ID")
557555

558556
request = CancelStatementRequest(statement_id=sea_statement_id)
559-
self._http_client._make_request(
557+
self.http_client._make_request(
560558
method="POST",
561559
path=self.CANCEL_STATEMENT_PATH_WITH_ID.format(sea_statement_id),
562560
data=request.to_dict(),
@@ -570,7 +568,7 @@ def close_command(self, command_id: CommandId) -> None:
570568
command_id: Command identifier to close
571569
572570
Raises:
573-
ValueError: If the command ID is invalid
571+
ProgrammingError: If the command ID is invalid
574572
"""
575573

576574
if command_id.backend_type != BackendType.SEA:
@@ -581,7 +579,7 @@ def close_command(self, command_id: CommandId) -> None:
581579
raise ValueError("Not a valid SEA command ID")
582580

583581
request = CloseStatementRequest(statement_id=sea_statement_id)
584-
self._http_client._make_request(
582+
self.http_client._make_request(
585583
method="DELETE",
586584
path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
587585
data=request.to_dict(),
@@ -600,7 +598,7 @@ def _poll_query(self, command_id: CommandId) -> GetStatementResponse:
600598
raise ValueError("Not a valid SEA command ID")
601599

602600
request = GetStatementRequest(statement_id=sea_statement_id)
603-
response_data = self._http_client._make_request(
601+
response_data = self.http_client._make_request(
604602
method="GET",
605603
path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
606604
data=request.to_dict(),

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
ExecuteStatementResponse,
2828
GetStatementResponse,
2929
CreateSessionResponse,
30-
GetChunksResponse,
3130
)
3231

3332
__all__ = [
@@ -50,5 +49,4 @@
5049
"ExecuteStatementResponse",
5150
"GetStatementResponse",
5251
"CreateSessionResponse",
53-
"GetChunksResponse",
5452
]

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

Lines changed: 0 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -160,37 +160,3 @@ class CreateSessionResponse:
160160
def from_dict(cls, data: Dict[str, Any]) -> "CreateSessionResponse":
161161
"""Create a CreateSessionResponse from a dictionary."""
162162
return cls(session_id=data.get("session_id", ""))
163-
164-
165-
@dataclass
166-
class GetChunksResponse:
167-
"""
168-
Response from getting chunks for a statement.
169-
170-
The response model can be found in the docs, here:
171-
https://docs.databricks.com/api/workspace/statementexecution/getstatementresultchunkn
172-
"""
173-
174-
data: Optional[List[List[Any]]] = None
175-
external_links: Optional[List[ExternalLink]] = None
176-
byte_count: Optional[int] = None
177-
chunk_index: Optional[int] = None
178-
next_chunk_index: Optional[int] = None
179-
next_chunk_internal_link: Optional[str] = None
180-
row_count: Optional[int] = None
181-
row_offset: Optional[int] = None
182-
183-
@classmethod
184-
def from_dict(cls, data: Dict[str, Any]) -> "GetChunksResponse":
185-
"""Create a GetChunksResponse from a dictionary."""
186-
result = _parse_result({"result": data})
187-
return cls(
188-
data=result.data,
189-
external_links=result.external_links,
190-
byte_count=result.byte_count,
191-
chunk_index=result.chunk_index,
192-
next_chunk_index=result.next_chunk_index,
193-
next_chunk_internal_link=result.next_chunk_internal_link,
194-
row_count=result.row_count,
195-
row_offset=result.row_offset,
196-
)

0 commit comments

Comments
 (0)