Skip to content

Commit 8ede414

Browse files
use spec-aligned Exceptions in SEA backend
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent e9b1314 commit 8ede414

File tree

2 files changed

+22
-21
lines changed

2 files changed

+22
-21
lines changed

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

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
BackendType,
2828
ExecuteResponse,
2929
)
30-
from databricks.sql.exc import DatabaseError, ServerOperationError
30+
from databricks.sql.exc import DatabaseError, ProgrammingError, ServerOperationError
3131
from databricks.sql.backend.sea.utils.http_client import SeaHttpClient
3232
from databricks.sql.types import SSLOptions
3333

@@ -172,7 +172,7 @@ def _extract_warehouse_id(self, http_path: str) -> str:
172172
f"Note: SEA only works for warehouses."
173173
)
174174
logger.error(error_message)
175-
raise ValueError(error_message)
175+
raise ProgrammingError(error_message)
176176

177177
@property
178178
def max_download_threads(self) -> int:
@@ -244,14 +244,14 @@ def close_session(self, session_id: SessionId) -> None:
244244
session_id: The session identifier returned by open_session()
245245
246246
Raises:
247-
ValueError: If the session ID is invalid
247+
ProgrammingError: If the session ID is invalid
248248
OperationalError: If there's an error closing the session
249249
"""
250250

251251
logger.debug("SeaDatabricksClient.close_session(session_id=%s)", session_id)
252252

253253
if session_id.backend_type != BackendType.SEA:
254-
raise ValueError("Not a valid SEA session ID")
254+
raise ProgrammingError("Not a valid SEA session ID")
255255
sea_session_id = session_id.to_sea_session_id()
256256

257257
request_data = DeleteSessionRequest(
@@ -429,7 +429,7 @@ def execute_command(
429429
"""
430430

431431
if session_id.backend_type != BackendType.SEA:
432-
raise ValueError("Not a valid SEA session ID")
432+
raise ProgrammingError("Not a valid SEA session ID")
433433

434434
sea_session_id = session_id.to_sea_session_id()
435435

@@ -504,11 +504,11 @@ def cancel_command(self, command_id: CommandId) -> None:
504504
command_id: Command identifier to cancel
505505
506506
Raises:
507-
ValueError: If the command ID is invalid
507+
ProgrammingError: If the command ID is invalid
508508
"""
509509

510510
if command_id.backend_type != BackendType.SEA:
511-
raise ValueError("Not a valid SEA command ID")
511+
raise ProgrammingError("Not a valid SEA command ID")
512512

513513
sea_statement_id = command_id.to_sea_statement_id()
514514

@@ -527,11 +527,11 @@ def close_command(self, command_id: CommandId) -> None:
527527
command_id: Command identifier to close
528528
529529
Raises:
530-
ValueError: If the command ID is invalid
530+
ProgrammingError: If the command ID is invalid
531531
"""
532532

533533
if command_id.backend_type != BackendType.SEA:
534-
raise ValueError("Not a valid SEA command ID")
534+
raise ProgrammingError("Not a valid SEA command ID")
535535

536536
sea_statement_id = command_id.to_sea_statement_id()
537537

@@ -553,7 +553,7 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
553553
CommandState: The current state of the command
554554
555555
Raises:
556-
ValueError: If the command ID is invalid
556+
ProgrammingError: If the command ID is invalid
557557
"""
558558

559559
if command_id.backend_type != BackendType.SEA:
@@ -592,7 +592,7 @@ def get_execution_result(
592592
"""
593593

594594
if command_id.backend_type != BackendType.SEA:
595-
raise ValueError("Not a valid SEA command ID")
595+
raise ProgrammingError("Not a valid SEA command ID")
596596

597597
sea_statement_id = command_id.to_sea_statement_id()
598598

@@ -658,7 +658,7 @@ def get_schemas(
658658
) -> SeaResultSet:
659659
"""Get schemas by executing 'SHOW SCHEMAS IN catalog [LIKE pattern]'."""
660660
if not catalog_name:
661-
raise ValueError("Catalog name is required for get_schemas")
661+
raise DatabaseError("Catalog name is required for get_schemas")
662662

663663
operation = MetadataCommands.SHOW_SCHEMAS.value.format(catalog_name)
664664

@@ -740,7 +740,7 @@ def get_columns(
740740
) -> SeaResultSet:
741741
"""Get columns by executing 'SHOW COLUMNS IN CATALOG catalog [SCHEMA LIKE pattern] [TABLE LIKE pattern] [LIKE pattern]'."""
742742
if not catalog_name:
743-
raise ValueError("Catalog name is required for get_columns")
743+
raise DatabaseError("Catalog name is required for get_columns")
744744

745745
operation = MetadataCommands.SHOW_COLUMNS.value.format(catalog_name)
746746

tests/unit/test_sea_backend.py

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
from databricks.sql.exc import (
1919
Error,
2020
NotSupportedError,
21+
ProgrammingError,
2122
ServerOperationError,
2223
DatabaseError,
2324
)
@@ -129,7 +130,7 @@ def test_initialization(self, mock_http_client):
129130
assert client3.max_download_threads == 5
130131

131132
# Test with invalid HTTP path
132-
with pytest.raises(ValueError) as excinfo:
133+
with pytest.raises(ProgrammingError) as excinfo:
133134
SeaDatabricksClient(
134135
server_hostname="test-server.databricks.com",
135136
port=443,
@@ -195,7 +196,7 @@ def test_session_management(self, sea_client, mock_http_client, thrift_session_i
195196
)
196197

197198
# Test close_session with invalid ID type
198-
with pytest.raises(ValueError) as excinfo:
199+
with pytest.raises(ProgrammingError) as excinfo:
199200
sea_client.close_session(thrift_session_id)
200201
assert "Not a valid SEA session ID" in str(excinfo.value)
201202

@@ -244,7 +245,7 @@ def test_command_execution_sync(
244245
assert cmd_id_arg.guid == "test-statement-123"
245246

246247
# Test with invalid session ID
247-
with pytest.raises(ValueError) as excinfo:
248+
with pytest.raises(ProgrammingError) as excinfo:
248249
mock_thrift_handle = MagicMock()
249250
mock_thrift_handle.sessionId.guid = b"guid"
250251
mock_thrift_handle.sessionId.secret = b"secret"
@@ -448,7 +449,7 @@ def test_command_management(
448449
)
449450

450451
# Test cancel_command with invalid ID
451-
with pytest.raises(ValueError) as excinfo:
452+
with pytest.raises(ProgrammingError) as excinfo:
452453
sea_client.cancel_command(thrift_command_id)
453454
assert "Not a valid SEA command ID" in str(excinfo.value)
454455

@@ -462,7 +463,7 @@ def test_command_management(
462463
)
463464

464465
# Test close_command with invalid ID
465-
with pytest.raises(ValueError) as excinfo:
466+
with pytest.raises(ProgrammingError) as excinfo:
466467
sea_client.close_command(thrift_command_id)
467468
assert "Not a valid SEA command ID" in str(excinfo.value)
468469

@@ -521,7 +522,7 @@ def test_command_management(
521522
assert result.status == CommandState.SUCCEEDED
522523

523524
# Test get_execution_result with invalid ID
524-
with pytest.raises(ValueError) as excinfo:
525+
with pytest.raises(ProgrammingError) as excinfo:
525526
sea_client.get_execution_result(thrift_command_id, mock_cursor)
526527
assert "Not a valid SEA command ID" in str(excinfo.value)
527528

@@ -717,7 +718,7 @@ def test_get_schemas(self, sea_client, sea_session_id, mock_cursor):
717718
)
718719

719720
# Case 3: Without catalog name (should raise ValueError)
720-
with pytest.raises(ValueError) as excinfo:
721+
with pytest.raises(DatabaseError) as excinfo:
721722
sea_client.get_schemas(
722723
session_id=sea_session_id,
723724
max_rows=100,
@@ -868,7 +869,7 @@ def test_get_columns(self, sea_client, sea_session_id, mock_cursor):
868869
)
869870

870871
# Case 3: Without catalog name (should raise ValueError)
871-
with pytest.raises(ValueError) as excinfo:
872+
with pytest.raises(DatabaseError) as excinfo:
872873
sea_client.get_columns(
873874
session_id=sea_session_id,
874875
max_rows=100,

0 commit comments

Comments
 (0)