Merge remote-tracking branch 'origin/backend-interface' into fetch-interface

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -2,15 +2,10 @@
 from typing import Dict, Tuple, List, Optional, Any, Union
 
 from databricks.sql.thrift_api.TCLIService import ttypes
+from databricks.sql.backend.types import SessionId, CommandId
 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 ==
@@ -20,19 +15,19 @@ def open_session(
         session_configuration: Optional[Dict[str, Any]],
         catalog: Optional[str],
         schema: Optional[str],
-    ) -> ttypes.TOpenSessionResp:
+    ) -> SessionId:
         pass
 
     @abstractmethod
-    def close_session(self, session_handle: ttypes.TSessionHandle) -> None:
+    def close_session(self, session_id: SessionId) -> None:
         pass
 
     # == Query Execution, Command Management ==
     @abstractmethod
     def execute_command(
         self,
         operation: str,
-        session_handle: ttypes.TSessionHandle,
+        session_id: SessionId,
         max_rows: int,
         max_bytes: int,
         lz4_compression: bool,
@@ -41,29 +36,25 @@ def execute_command(
         parameters: List[ttypes.TSparkParameter],
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
-    ) -> "ResultSet":  # Changed return type to ResultSet
+    ) -> Any:
         pass
 
     @abstractmethod
-    def cancel_command(self, operation_handle: ttypes.TOperationHandle) -> None:
+    def cancel_command(self, command_id: CommandId) -> None:
         pass
 
     @abstractmethod
-    def close_command(
-        self, operation_handle: ttypes.TOperationHandle
-    ) -> ttypes.TStatus:
+    def close_command(self, command_id: CommandId) -> ttypes.TStatus:
         pass
 
     @abstractmethod
-    def get_query_state(
-        self, operation_handle: ttypes.TOperationHandle
-    ) -> ttypes.TOperationState:
+    def get_query_state(self, command_id: CommandId) -> ttypes.TOperationState:
         pass
 
     @abstractmethod
     def get_execution_result(
         self,
-        operation_handle: ttypes.TOperationHandle,
+        command_id: CommandId,
         cursor: Any,
     ) -> ExecuteResponse:
         pass
@@ -72,7 +63,7 @@ def get_execution_result(
     @abstractmethod
     def get_catalogs(
         self,
-        session_handle: ttypes.TSessionHandle,
+        session_id: SessionId,
         max_rows: int,
         max_bytes: int,
         cursor: Any,
@@ -82,7 +73,7 @@ def get_catalogs(
     @abstractmethod
     def get_schemas(
         self,
-        session_handle: ttypes.TSessionHandle,
+        session_id: SessionId,
         max_rows: int,
         max_bytes: int,
         cursor: Any,
@@ -94,7 +85,7 @@ def get_schemas(
     @abstractmethod
     def get_tables(
         self,
-        session_handle: ttypes.TSessionHandle,
+        session_id: SessionId,
         max_rows: int,
         max_bytes: int,
         cursor: Any,
@@ -108,7 +99,7 @@ def get_tables(
     @abstractmethod
     def get_columns(
         self,
-        session_handle: ttypes.TSessionHandle,
+        session_id: SessionId,
         max_rows: int,
         max_bytes: int,
         cursor: Any,
@@ -121,11 +112,11 @@ def get_columns(
 
     # == Utility Methods ==
     @abstractmethod
-    def handle_to_id(self, handle: ttypes.TSessionHandle) -> bytes:
+    def handle_to_id(self, session_id: SessionId) -> Any:
         pass
 
     @abstractmethod
-    def handle_to_hex_id(self, handle: ttypes.TSessionHandle) -> str:
+    def handle_to_hex_id(self, session_id: SessionId) -> str:
         pass
 
     # Properties related to specific backend features