1+ from __future__ import annotations
2+
13import logging
24import time
35import re
1517
1618if 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
2022from databricks .sql .backend .databricks_client import DatabricksClient
2123from 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