Skip to content

Commit 2c22010

Browse files
remove SeaResultSet type assertion
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent a3ca7c7 commit 2c22010

File tree

1 file changed

+17
-20
lines changed

1 file changed

+17
-20
lines changed

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

Lines changed: 17 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
1+
from __future__ import annotations
2+
13
import logging
24
import time
35
import re
4-
from typing import Any, Dict, Tuple, List, Optional, Union, TYPE_CHECKING, Set
6+
from typing import Dict, Tuple, List, Optional, Union, TYPE_CHECKING, Set
57

68
from databricks.sql.backend.sea.models.base import ExternalLink, ResultManifest
79
from databricks.sql.backend.sea.utils.constants import (
@@ -12,11 +14,12 @@
1214
WaitTimeout,
1315
MetadataCommands,
1416
)
17+
1518
from databricks.sql.thrift_api.TCLIService import ttypes
1619

1720
if TYPE_CHECKING:
1821
from databricks.sql.client import Cursor
19-
from databricks.sql.result_set import ResultSet
22+
from databricks.sql.result_set import SeaResultSet
2023

2124
from databricks.sql.backend.databricks_client import DatabricksClient
2225
from databricks.sql.backend.types import (
@@ -409,7 +412,7 @@ def execute_command(
409412
parameters: List[ttypes.TSparkParameter],
410413
async_op: bool,
411414
enforce_embedded_schema_correctness: bool,
412-
) -> Union["ResultSet", None]:
415+
) -> Union[SeaResultSet, None]:
413416
"""
414417
Execute a SQL command using the SEA backend.
415418
@@ -426,7 +429,7 @@ def execute_command(
426429
enforce_embedded_schema_correctness: Whether to enforce schema correctness
427430
428431
Returns:
429-
ResultSet: A SeaResultSet instance for the executed command
432+
SeaResultSet: A SeaResultSet instance for the executed command
430433
"""
431434

432435
if session_id.backend_type != BackendType.SEA:
@@ -576,8 +579,8 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
576579
def get_execution_result(
577580
self,
578581
command_id: CommandId,
579-
cursor: "Cursor",
580-
) -> "ResultSet":
582+
cursor: Cursor,
583+
) -> SeaResultSet:
581584
"""
582585
Get the result of a command execution.
583586
@@ -586,7 +589,7 @@ def get_execution_result(
586589
cursor: Cursor executing the command
587590
588591
Returns:
589-
ResultSet: A SeaResultSet instance with the execution results
592+
SeaResultSet: A SeaResultSet instance with the execution results
590593
591594
Raises:
592595
ValueError: If the command ID is invalid
@@ -659,8 +662,8 @@ def get_catalogs(
659662
session_id: SessionId,
660663
max_rows: int,
661664
max_bytes: int,
662-
cursor: "Cursor",
663-
) -> "ResultSet":
665+
cursor: Cursor,
666+
) -> SeaResultSet:
664667
"""Get available catalogs by executing 'SHOW CATALOGS'."""
665668
result = self.execute_command(
666669
operation=MetadataCommands.SHOW_CATALOGS.value,
@@ -682,10 +685,10 @@ def get_schemas(
682685
session_id: SessionId,
683686
max_rows: int,
684687
max_bytes: int,
685-
cursor: "Cursor",
688+
cursor: Cursor,
686689
catalog_name: Optional[str] = None,
687690
schema_name: Optional[str] = None,
688-
) -> "ResultSet":
691+
) -> SeaResultSet:
689692
"""Get schemas by executing 'SHOW SCHEMAS IN catalog [LIKE pattern]'."""
690693
if not catalog_name:
691694
raise ValueError("Catalog name is required for get_schemas")
@@ -720,7 +723,7 @@ def get_tables(
720723
schema_name: Optional[str] = None,
721724
table_name: Optional[str] = None,
722725
table_types: Optional[List[str]] = None,
723-
) -> "ResultSet":
726+
) -> SeaResultSet:
724727
"""Get tables by executing 'SHOW TABLES IN catalog [SCHEMA LIKE pattern] [LIKE pattern]'."""
725728
operation = (
726729
MetadataCommands.SHOW_TABLES_ALL_CATALOGS.value
@@ -750,12 +753,6 @@ def get_tables(
750753
)
751754
assert result is not None, "execute_command returned None in synchronous mode"
752755

753-
from databricks.sql.result_set import SeaResultSet
754-
755-
assert isinstance(
756-
result, SeaResultSet
757-
), "execute_command returned a non-SeaResultSet"
758-
759756
# Apply client-side filtering by table_types
760757
from databricks.sql.backend.sea.utils.filters import ResultSetFilter
761758

@@ -768,12 +765,12 @@ def get_columns(
768765
session_id: SessionId,
769766
max_rows: int,
770767
max_bytes: int,
771-
cursor: "Cursor",
768+
cursor: Cursor,
772769
catalog_name: Optional[str] = None,
773770
schema_name: Optional[str] = None,
774771
table_name: Optional[str] = None,
775772
column_name: Optional[str] = None,
776-
) -> "ResultSet":
773+
) -> SeaResultSet:
777774
"""Get columns by executing 'SHOW COLUMNS IN CATALOG catalog [SCHEMA LIKE pattern] [TABLE LIKE pattern] [LIKE pattern]'."""
778775
if not catalog_name:
779776
raise ValueError("Catalog name is required for get_columns")

0 commit comments

Comments
 (0)