1
+ from __future__ import annotations
2
+
1
3
import logging
2
4
import time
3
5
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
5
7
6
8
from databricks .sql .backend .sea .models .base import ExternalLink , ResultManifest
7
9
from databricks .sql .backend .sea .utils .constants import (
12
14
WaitTimeout ,
13
15
MetadataCommands ,
14
16
)
17
+
15
18
from databricks .sql .thrift_api .TCLIService import ttypes
16
19
17
20
if TYPE_CHECKING :
18
21
from databricks .sql .client import Cursor
19
- from databricks .sql .result_set import ResultSet
22
+ from databricks .sql .result_set import SeaResultSet
20
23
21
24
from databricks .sql .backend .databricks_client import DatabricksClient
22
25
from databricks .sql .backend .types import (
@@ -409,7 +412,7 @@ def execute_command(
409
412
parameters : List [ttypes .TSparkParameter ],
410
413
async_op : bool ,
411
414
enforce_embedded_schema_correctness : bool ,
412
- ) -> Union ["ResultSet" , None ]:
415
+ ) -> Union [SeaResultSet , None ]:
413
416
"""
414
417
Execute a SQL command using the SEA backend.
415
418
@@ -426,7 +429,7 @@ def execute_command(
426
429
enforce_embedded_schema_correctness: Whether to enforce schema correctness
427
430
428
431
Returns:
429
- ResultSet : A SeaResultSet instance for the executed command
432
+ SeaResultSet : A SeaResultSet instance for the executed command
430
433
"""
431
434
432
435
if session_id .backend_type != BackendType .SEA :
@@ -576,8 +579,8 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
576
579
def get_execution_result (
577
580
self ,
578
581
command_id : CommandId ,
579
- cursor : " Cursor" ,
580
- ) -> "ResultSet" :
582
+ cursor : Cursor ,
583
+ ) -> SeaResultSet :
581
584
"""
582
585
Get the result of a command execution.
583
586
@@ -586,7 +589,7 @@ def get_execution_result(
586
589
cursor: Cursor executing the command
587
590
588
591
Returns:
589
- ResultSet : A SeaResultSet instance with the execution results
592
+ SeaResultSet : A SeaResultSet instance with the execution results
590
593
591
594
Raises:
592
595
ValueError: If the command ID is invalid
@@ -659,8 +662,8 @@ def get_catalogs(
659
662
session_id : SessionId ,
660
663
max_rows : int ,
661
664
max_bytes : int ,
662
- cursor : " Cursor" ,
663
- ) -> "ResultSet" :
665
+ cursor : Cursor ,
666
+ ) -> SeaResultSet :
664
667
"""Get available catalogs by executing 'SHOW CATALOGS'."""
665
668
result = self .execute_command (
666
669
operation = MetadataCommands .SHOW_CATALOGS .value ,
@@ -682,10 +685,10 @@ def get_schemas(
682
685
session_id : SessionId ,
683
686
max_rows : int ,
684
687
max_bytes : int ,
685
- cursor : " Cursor" ,
688
+ cursor : Cursor ,
686
689
catalog_name : Optional [str ] = None ,
687
690
schema_name : Optional [str ] = None ,
688
- ) -> "ResultSet" :
691
+ ) -> SeaResultSet :
689
692
"""Get schemas by executing 'SHOW SCHEMAS IN catalog [LIKE pattern]'."""
690
693
if not catalog_name :
691
694
raise ValueError ("Catalog name is required for get_schemas" )
@@ -720,7 +723,7 @@ def get_tables(
720
723
schema_name : Optional [str ] = None ,
721
724
table_name : Optional [str ] = None ,
722
725
table_types : Optional [List [str ]] = None ,
723
- ) -> "ResultSet" :
726
+ ) -> SeaResultSet :
724
727
"""Get tables by executing 'SHOW TABLES IN catalog [SCHEMA LIKE pattern] [LIKE pattern]'."""
725
728
operation = (
726
729
MetadataCommands .SHOW_TABLES_ALL_CATALOGS .value
@@ -750,12 +753,6 @@ def get_tables(
750
753
)
751
754
assert result is not None , "execute_command returned None in synchronous mode"
752
755
753
- from databricks .sql .result_set import SeaResultSet
754
-
755
- assert isinstance (
756
- result , SeaResultSet
757
- ), "execute_command returned a non-SeaResultSet"
758
-
759
756
# Apply client-side filtering by table_types
760
757
from databricks .sql .backend .sea .utils .filters import ResultSetFilter
761
758
@@ -768,12 +765,12 @@ def get_columns(
768
765
session_id : SessionId ,
769
766
max_rows : int ,
770
767
max_bytes : int ,
771
- cursor : " Cursor" ,
768
+ cursor : Cursor ,
772
769
catalog_name : Optional [str ] = None ,
773
770
schema_name : Optional [str ] = None ,
774
771
table_name : Optional [str ] = None ,
775
772
column_name : Optional [str ] = None ,
776
- ) -> "ResultSet" :
773
+ ) -> SeaResultSet :
777
774
"""Get columns by executing 'SHOW COLUMNS IN CATALOG catalog [SCHEMA LIKE pattern] [TABLE LIKE pattern] [LIKE pattern]'."""
778
775
if not catalog_name :
779
776
raise ValueError ("Catalog name is required for get_columns" )
0 commit comments