Skip to content

Commit e9b1314

Browse files
make SEA backend methods return SeaResultSet
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent a4d5bdb commit e9b1314

File tree

1 file changed

+16
-20
lines changed

1 file changed

+16
-20
lines changed

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

Lines changed: 16 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from __future__ import annotations
2+
13
import logging
24
import time
35
import re
@@ -15,7 +17,7 @@
1517

1618
if TYPE_CHECKING:
1719
from databricks.sql.client import Cursor
18-
from databricks.sql.result_set import ResultSet
20+
from databricks.sql.result_set import SeaResultSet
1921

2022
from databricks.sql.backend.databricks_client import DatabricksClient
2123
from databricks.sql.backend.types import (
@@ -401,12 +403,12 @@ def execute_command(
401403
max_rows: int,
402404
max_bytes: int,
403405
lz4_compression: bool,
404-
cursor: "Cursor",
406+
cursor: Cursor,
405407
use_cloud_fetch: bool,
406408
parameters: List[Dict[str, Any]],
407409
async_op: bool,
408410
enforce_embedded_schema_correctness: bool,
409-
) -> Union["ResultSet", None]:
411+
) -> Union[SeaResultSet, None]:
410412
"""
411413
Execute a SQL command using the SEA backend.
412414
@@ -573,8 +575,8 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
573575
def get_execution_result(
574576
self,
575577
command_id: CommandId,
576-
cursor: "Cursor",
577-
) -> "ResultSet":
578+
cursor: Cursor,
579+
) -> SeaResultSet:
578580
"""
579581
Get the result of a command execution.
580582
@@ -583,7 +585,7 @@ def get_execution_result(
583585
cursor: Cursor executing the command
584586
585587
Returns:
586-
ResultSet: A SeaResultSet instance with the execution results
588+
SeaResultSet: A SeaResultSet instance with the execution results
587589
588590
Raises:
589591
ValueError: If the command ID is invalid
@@ -627,8 +629,8 @@ def get_catalogs(
627629
session_id: SessionId,
628630
max_rows: int,
629631
max_bytes: int,
630-
cursor: "Cursor",
631-
) -> "ResultSet":
632+
cursor: Cursor,
633+
) -> SeaResultSet:
632634
"""Get available catalogs by executing 'SHOW CATALOGS'."""
633635
result = self.execute_command(
634636
operation=MetadataCommands.SHOW_CATALOGS.value,
@@ -650,10 +652,10 @@ def get_schemas(
650652
session_id: SessionId,
651653
max_rows: int,
652654
max_bytes: int,
653-
cursor: "Cursor",
655+
cursor: Cursor,
654656
catalog_name: Optional[str] = None,
655657
schema_name: Optional[str] = None,
656-
) -> "ResultSet":
658+
) -> SeaResultSet:
657659
"""Get schemas by executing 'SHOW SCHEMAS IN catalog [LIKE pattern]'."""
658660
if not catalog_name:
659661
raise ValueError("Catalog name is required for get_schemas")
@@ -683,12 +685,12 @@ def get_tables(
683685
session_id: SessionId,
684686
max_rows: int,
685687
max_bytes: int,
686-
cursor: "Cursor",
688+
cursor: Cursor,
687689
catalog_name: Optional[str] = None,
688690
schema_name: Optional[str] = None,
689691
table_name: Optional[str] = None,
690692
table_types: Optional[List[str]] = None,
691-
) -> "ResultSet":
693+
) -> SeaResultSet:
692694
"""Get tables by executing 'SHOW TABLES IN catalog [SCHEMA LIKE pattern] [LIKE pattern]'."""
693695
operation = (
694696
MetadataCommands.SHOW_TABLES_ALL_CATALOGS.value
@@ -718,12 +720,6 @@ def get_tables(
718720
)
719721
assert result is not None, "execute_command returned None in synchronous mode"
720722

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

@@ -736,12 +732,12 @@ def get_columns(
736732
session_id: SessionId,
737733
max_rows: int,
738734
max_bytes: int,
739-
cursor: "Cursor",
735+
cursor: Cursor,
740736
catalog_name: Optional[str] = None,
741737
schema_name: Optional[str] = None,
742738
table_name: Optional[str] = None,
743739
column_name: Optional[str] = None,
744-
) -> "ResultSet":
740+
) -> SeaResultSet:
745741
"""Get columns by executing 'SHOW COLUMNS IN CATALOG catalog [SCHEMA LIKE pattern] [TABLE LIKE pattern] [LIKE pattern]'."""
746742
if not catalog_name:
747743
raise ValueError("Catalog name is required for get_columns")

0 commit comments

Comments
 (0)