Skip to content

Commit a4d5bdb

Browse files
remove cast, throw error if not SeaResultSet
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 14057ac commit a4d5bdb

File tree

2 files changed

+11
-4
lines changed

2 files changed

+11
-4
lines changed

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import logging
22
import time
33
import re
4-
from typing import Any, Dict, Tuple, List, Optional, Union, TYPE_CHECKING, Set, cast
4+
from typing import Any, Dict, Tuple, List, Optional, Union, TYPE_CHECKING, Set
55

66
from databricks.sql.backend.sea.models.base import ResultManifest
77
from databricks.sql.backend.sea.utils.constants import (
@@ -718,11 +718,15 @@ def get_tables(
718718
)
719719
assert result is not None, "execute_command returned None in synchronous mode"
720720

721+
from databricks.sql.result_set import SeaResultSet
722+
723+
assert isinstance(
724+
result, SeaResultSet
725+
), "execute_command returned a non-SeaResultSet"
726+
721727
# Apply client-side filtering by table_types
722728
from databricks.sql.backend.sea.utils.filters import ResultSetFilter
723-
from databricks.sql.result_set import SeaResultSet
724729

725-
result = cast(SeaResultSet, result)
726730
result = ResultSetFilter.filter_tables_by_type(result, table_types)
727731

728732
return result

tests/unit/test_sea_backend.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -729,7 +729,10 @@ def test_get_schemas(self, sea_client, sea_session_id, mock_cursor):
729729
def test_get_tables(self, sea_client, sea_session_id, mock_cursor):
730730
"""Test the get_tables method with various parameter combinations."""
731731
# Mock the execute_command method
732-
mock_result_set = Mock()
732+
from databricks.sql.result_set import SeaResultSet
733+
734+
mock_result_set = Mock(spec=SeaResultSet)
735+
733736
with patch.object(
734737
sea_client, "execute_command", return_value=mock_result_set
735738
) as mock_execute:

0 commit comments

Comments
 (0)