1
+ from __future__ import annotations
2
+
1
3
import logging
2
4
import time
3
5
import re
15
17
16
18
if TYPE_CHECKING :
17
19
from databricks .sql .client import Cursor
18
- from databricks .sql .result_set import ResultSet
20
+ from databricks .sql .result_set import SeaResultSet
19
21
20
22
from databricks .sql .backend .databricks_client import DatabricksClient
21
23
from databricks .sql .backend .types import (
@@ -401,12 +403,12 @@ def execute_command(
401
403
max_rows : int ,
402
404
max_bytes : int ,
403
405
lz4_compression : bool ,
404
- cursor : " Cursor" ,
406
+ cursor : Cursor ,
405
407
use_cloud_fetch : bool ,
406
408
parameters : List [Dict [str , Any ]],
407
409
async_op : bool ,
408
410
enforce_embedded_schema_correctness : bool ,
409
- ) -> Union ["ResultSet" , None ]:
411
+ ) -> Union [SeaResultSet , None ]:
410
412
"""
411
413
Execute a SQL command using the SEA backend.
412
414
@@ -573,8 +575,8 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
573
575
def get_execution_result (
574
576
self ,
575
577
command_id : CommandId ,
576
- cursor : " Cursor" ,
577
- ) -> "ResultSet" :
578
+ cursor : Cursor ,
579
+ ) -> SeaResultSet :
578
580
"""
579
581
Get the result of a command execution.
580
582
@@ -583,7 +585,7 @@ def get_execution_result(
583
585
cursor: Cursor executing the command
584
586
585
587
Returns:
586
- ResultSet : A SeaResultSet instance with the execution results
588
+ SeaResultSet : A SeaResultSet instance with the execution results
587
589
588
590
Raises:
589
591
ValueError: If the command ID is invalid
@@ -627,8 +629,8 @@ def get_catalogs(
627
629
session_id : SessionId ,
628
630
max_rows : int ,
629
631
max_bytes : int ,
630
- cursor : " Cursor" ,
631
- ) -> "ResultSet" :
632
+ cursor : Cursor ,
633
+ ) -> SeaResultSet :
632
634
"""Get available catalogs by executing 'SHOW CATALOGS'."""
633
635
result = self .execute_command (
634
636
operation = MetadataCommands .SHOW_CATALOGS .value ,
@@ -650,10 +652,10 @@ def get_schemas(
650
652
session_id : SessionId ,
651
653
max_rows : int ,
652
654
max_bytes : int ,
653
- cursor : " Cursor" ,
655
+ cursor : Cursor ,
654
656
catalog_name : Optional [str ] = None ,
655
657
schema_name : Optional [str ] = None ,
656
- ) -> "ResultSet" :
658
+ ) -> SeaResultSet :
657
659
"""Get schemas by executing 'SHOW SCHEMAS IN catalog [LIKE pattern]'."""
658
660
if not catalog_name :
659
661
raise ValueError ("Catalog name is required for get_schemas" )
@@ -683,12 +685,12 @@ def get_tables(
683
685
session_id : SessionId ,
684
686
max_rows : int ,
685
687
max_bytes : int ,
686
- cursor : " Cursor" ,
688
+ cursor : Cursor ,
687
689
catalog_name : Optional [str ] = None ,
688
690
schema_name : Optional [str ] = None ,
689
691
table_name : Optional [str ] = None ,
690
692
table_types : Optional [List [str ]] = None ,
691
- ) -> "ResultSet" :
693
+ ) -> SeaResultSet :
692
694
"""Get tables by executing 'SHOW TABLES IN catalog [SCHEMA LIKE pattern] [LIKE pattern]'."""
693
695
operation = (
694
696
MetadataCommands .SHOW_TABLES_ALL_CATALOGS .value
@@ -718,12 +720,6 @@ def get_tables(
718
720
)
719
721
assert result is not None , "execute_command returned None in synchronous mode"
720
722
721
- from databricks .sql .result_set import SeaResultSet
722
-
723
- assert isinstance (
724
- result , SeaResultSet
725
- ), "execute_command returned a non-SeaResultSet"
726
-
727
723
# Apply client-side filtering by table_types
728
724
from databricks .sql .backend .sea .utils .filters import ResultSetFilter
729
725
@@ -736,12 +732,12 @@ def get_columns(
736
732
session_id : SessionId ,
737
733
max_rows : int ,
738
734
max_bytes : int ,
739
- cursor : " Cursor" ,
735
+ cursor : Cursor ,
740
736
catalog_name : Optional [str ] = None ,
741
737
schema_name : Optional [str ] = None ,
742
738
table_name : Optional [str ] = None ,
743
739
column_name : Optional [str ] = None ,
744
- ) -> "ResultSet" :
740
+ ) -> SeaResultSet :
745
741
"""Get columns by executing 'SHOW COLUMNS IN CATALOG catalog [SCHEMA LIKE pattern] [TABLE LIKE pattern] [LIKE pattern]'."""
746
742
if not catalog_name :
747
743
raise ValueError ("Catalog name is required for get_columns" )
0 commit comments