databricks
diff --git a/‎src/databricks/sql/backend/databricks_client.py
Lines changed: 22 additions & 19 deletions b/‎src/databricks/sql/backend/databricks_client.py
Lines changed: 22 additions & 19 deletions
@@ -15,10 +15,16 @@
     from databricks.sql.client import Cursor
 
 from databricks.sql.thrift_api.TCLIService import ttypes
-from databricks.sql.backend.types import SessionId, CommandId
+from databricks.sql.backend.types import SessionId, CommandId, CommandState
 from databricks.sql.utils import ExecuteResponse
 from databricks.sql.types import SSLOptions
 
+# Forward reference for type hints
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from databricks.sql.result_set import ResultSet
+
 
 class DatabricksClient(ABC):
     # == Connection and Session Management ==
@@ -81,7 +87,7 @@ def execute_command(
         parameters: List[ttypes.TSparkParameter],
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
-    ) -> Optional[ExecuteResponse]:
+    ) -> Union["ResultSet", None]:
         """
         Executes a SQL command or query within the specified session.
 
@@ -101,7 +107,7 @@ def execute_command(
             enforce_embedded_schema_correctness: Whether to enforce schema correctness
 
         Returns:
-            If async_op is False, returns an ExecuteResponse object containing the
+            If async_op is False, returns a ResultSet object containing the
             query results and metadata. If async_op is True, returns None and the
             results must be fetched later using get_execution_result().
 
@@ -130,7 +136,7 @@ def cancel_command(self, command_id: CommandId) -> None:
         pass
 
     @abstractmethod
-    def close_command(self, command_id: CommandId) -> ttypes.TStatus:
+    def close_command(self, command_id: CommandId) -> None:
         """
         Closes a command and releases associated resources.
 
@@ -140,17 +146,14 @@ def close_command(self, command_id: CommandId) -> ttypes.TStatus:
         Args:
             command_id: The command identifier to close
 
-        Returns:
-            ttypes.TStatus: The status of the close operation
-
         Raises:
             ValueError: If the command ID is invalid
             OperationalError: If there's an error closing the command
         """
         pass
 
     @abstractmethod
-    def get_query_state(self, command_id: CommandId) -> ttypes.TOperationState:
+    def get_query_state(self, command_id: CommandId) -> CommandState:
         """
         Gets the current state of a query or command.
 
@@ -160,7 +163,7 @@ def get_query_state(self, command_id: CommandId) -> ttypes.TOperationState:
             command_id: The command identifier to check
 
         Returns:
-            ttypes.TOperationState: The current state of the command
+            CommandState: The current state of the command
 
         Raises:
             ValueError: If the command ID is invalid
@@ -175,7 +178,7 @@ def get_execution_result(
         self,
         command_id: CommandId,
         cursor: "Cursor",
-    ) -> ExecuteResponse:
+    ) -> "ResultSet":
         """
         Retrieves the results of a previously executed command.
 
@@ -187,7 +190,7 @@ def get_execution_result(
             cursor: The cursor object that will handle the results
 
         Returns:
-            ExecuteResponse: An object containing the query results and metadata
+            ResultSet: An object containing the query results and metadata
 
         Raises:
             ValueError: If the command ID is invalid
@@ -203,7 +206,7 @@ def get_catalogs(
         max_rows: int,
         max_bytes: int,
         cursor: "Cursor",
-    ) -> ExecuteResponse:
+    ) -> "ResultSet":
         """
         Retrieves a list of available catalogs.
 
@@ -217,7 +220,7 @@ def get_catalogs(
             cursor: The cursor object that will handle the results
 
         Returns:
-            ExecuteResponse: An object containing the catalog metadata
+            ResultSet: An object containing the catalog metadata
 
         Raises:
             ValueError: If the session ID is invalid
@@ -234,7 +237,7 @@ def get_schemas(
         cursor: "Cursor",
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
-    ) -> ExecuteResponse:
+    ) -> "ResultSet":
         """
         Retrieves a list of schemas, optionally filtered by catalog and schema name patterns.
 
@@ -250,7 +253,7 @@ def get_schemas(
             schema_name: Optional schema name pattern to filter by
 
         Returns:
-            ExecuteResponse: An object containing the schema metadata
+            ResultSet: An object containing the schema metadata
 
         Raises:
             ValueError: If the session ID is invalid
@@ -269,7 +272,7 @@ def get_tables(
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
         table_types: Optional[List[str]] = None,
-    ) -> ExecuteResponse:
+    ) -> "ResultSet":
         """
         Retrieves a list of tables, optionally filtered by catalog, schema, table name, and table types.
 
@@ -287,7 +290,7 @@ def get_tables(
             table_types: Optional list of table types to filter by (e.g., ['TABLE', 'VIEW'])
 
         Returns:
-            ExecuteResponse: An object containing the table metadata
+            ResultSet: An object containing the table metadata
 
         Raises:
             ValueError: If the session ID is invalid
@@ -306,7 +309,7 @@ def get_columns(
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
         column_name: Optional[str] = None,
-    ) -> ExecuteResponse:
+    ) -> "ResultSet":
         """
         Retrieves a list of columns, optionally filtered by catalog, schema, table, and column name patterns.
 
@@ -324,7 +327,7 @@ def get_columns(
             column_name: Optional column name pattern to filter by
 
         Returns:
-            ExecuteResponse: An object containing the column metadata
+            ResultSet: An object containing the column metadata
 
         Raises:
             ValueError: If the session ID is invalid