Skip to content

Commit bb5cf17

Browse files
minor fixes
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 329c9af commit bb5cf17

File tree

2 files changed

+21
-27
lines changed

2 files changed

+21
-27
lines changed

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

Lines changed: 20 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111

1212
import databricks
1313
from databricks.sql.auth.retry import CommandType
14-
from databricks.sql.backend.sea.models.base import ExternalLink
14+
from databricks.sql.backend.sea.models.base import ExternalLink, ResultManifest
1515
from databricks.sql.backend.sea.utils.constants import (
1616
ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
1717
ResultFormat,
@@ -320,6 +320,7 @@ def _handle_request_error(self, error_info, attempt, elapsed):
320320
user_friendly_error_message = error_info.user_friendly_error_message(
321321
no_retry_reason, attempt, elapsed
322322
)
323+
logger.info(f"User friendly error message: {user_friendly_error_message}")
323324
network_request_error = RequestError(
324325
user_friendly_error_message, full_error_info_context, error_info.error
325326
)
@@ -594,14 +595,6 @@ def open_session(
594595

595596
session_response = CreateSessionResponse.from_dict(response)
596597
session_id = session_response.session_id
597-
if not session_id:
598-
raise ServerOperationError(
599-
"Failed to create session: No session ID returned",
600-
{
601-
"operation-id": None,
602-
"diagnostic-info": None,
603-
},
604-
)
605598

606599
return SessionId.from_sea_session_id(session_id)
607600
except Exception as e:
@@ -733,7 +726,9 @@ def get_chunk_link(self, statement_id: str, chunk_index: int) -> ExternalLink:
733726

734727
return link
735728

736-
def _results_message_to_execute_response(self, sea_response, command_id):
729+
def _results_message_to_execute_response(
730+
self, response: GetStatementResponse, command_id: CommandId
731+
) -> ExecuteResponse:
737732
"""
738733
Convert a SEA response to an ExecuteResponse and extract result data.
739734
@@ -811,7 +806,7 @@ def execute_command(
811806
lz4_compression: bool,
812807
cursor: "Cursor",
813808
use_cloud_fetch: bool,
814-
parameters: List[Dict[str, Any]],
809+
parameters: List[ttypes.TSparkParameter],
815810
async_op: bool,
816811
enforce_embedded_schema_correctness: bool,
817812
) -> Union["ResultSet", None]:
@@ -845,9 +840,9 @@ def execute_command(
845840
for param in parameters:
846841
sea_parameters.append(
847842
StatementParameter(
848-
name=param["name"],
849-
value=param["value"],
850-
type=param["type"] if "type" in param else None,
843+
name=param.name,
844+
value=param.value.stringValue,
845+
type=param.type,
851846
)
852847
)
853848

@@ -1057,25 +1052,24 @@ def get_execution_result(
10571052
params=None,
10581053
headers=None,
10591054
)
1055+
response = GetStatementResponse.from_dict(response_data)
10601056

10611057
# Create and return a SeaResultSet
10621058
from databricks.sql.result_set import SeaResultSet
10631059

10641060
# Convert the response to an ExecuteResponse and extract result data
1065-
(
1066-
execute_response,
1067-
result_data,
1068-
manifest,
1069-
) = self._results_message_to_execute_response(response_data, command_id)
1061+
execute_response = self._results_message_to_execute_response(
1062+
response, command_id
1063+
)
10701064

10711065
return SeaResultSet(
10721066
connection=cursor.connection,
10731067
execute_response=execute_response,
10741068
sea_client=self,
10751069
buffer_size_bytes=cursor.buffer_size_bytes,
10761070
arraysize=cursor.arraysize,
1077-
result_data=result_data,
1078-
manifest=manifest,
1071+
result_data=response.result,
1072+
manifest=response.manifest,
10791073
)
10801074
except Exception as e:
10811075
logger.error("SeaDatabricksClient.get_execution_result: Exception: %s", e)
@@ -1091,9 +1085,12 @@ def get_chunk_link(self, statement_id: str, chunk_index: int) -> ExternalLink:
10911085
ExternalLink: External link for the chunk
10921086
"""
10931087

1094-
response_data = self.http_client._make_request(
1095-
method="GET",
1088+
response_data = self.make_request(
1089+
method_name="GET",
10961090
path=self.CHUNK_PATH_WITH_ID_AND_INDEX.format(statement_id, chunk_index),
1091+
data=None,
1092+
params=None,
1093+
headers=None,
10971094
)
10981095
response = GetChunksResponse.from_dict(response_data)
10991096

@@ -1180,9 +1177,6 @@ def get_tables(
11801177
table_types: Optional[List[str]] = None,
11811178
) -> "ResultSet":
11821179
"""Get tables by executing 'SHOW TABLES IN catalog [SCHEMA LIKE pattern] [LIKE pattern]'."""
1183-
if not catalog_name:
1184-
raise ValueError("Catalog name is required for get_tables")
1185-
11861180
operation = (
11871181
MetadataCommands.SHOW_TABLES_ALL_CATALOGS.value
11881182
if catalog_name in [None, "*", "%"]

tests/e2e/common/large_queries_mixin.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ def test_long_running_query(self):
8787
and asserts that the query completes successfully.
8888
"""
8989
minutes = 60
90-
min_duration = 5 * minutes
90+
min_duration = 3 * minutes
9191

9292
duration = -1
9393
scale0 = 10000

0 commit comments

Comments
 (0)