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

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -1,5 +1,20 @@
+"""
+Abstract client interface for interacting with Databricks SQL services.
+
+Implementations of this class are responsible for:
+- Managing connections to Databricks SQL services
+- Handling authentication
+- Executing SQL queries and commands
+- Retrieving query results
+- Fetching metadata about catalogs, schemas, tables, and columns
+- Managing error handling and retries
+"""
+
 from abc import ABC, abstractmethod
-from typing import Dict, Tuple, List, Optional, Any, Union
+from typing import Dict, Tuple, List, Optional, Any, Union, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from databricks.sql.client import Cursor
 
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.backend.types import SessionId, CommandId, CommandState
@@ -22,10 +37,42 @@ def open_session(
         catalog: Optional[str],
         schema: Optional[str],
     ) -> SessionId:
+        """
+        Opens a new session with the Databricks SQL service.
+
+        This method establishes a new session with the server and returns a session
+        identifier that can be used for subsequent operations.
+
+        Args:
+            session_configuration: Optional dictionary of configuration parameters for the session
+            catalog: Optional catalog name to use as the initial catalog for the session
+            schema: Optional schema name to use as the initial schema for the session
+
+        Returns:
+            SessionId: A session identifier object that can be used for subsequent operations
+
+        Raises:
+            Error: If the session configuration is invalid
+            OperationalError: If there's an error establishing the session
+            InvalidServerResponseError: If the server response is invalid or unexpected
+        """
         pass
 
     @abstractmethod
     def close_session(self, session_id: SessionId) -> None:
+        """
+        Closes an existing session with the Databricks SQL service.
+
+        This method terminates the session identified by the given session ID and
+        releases any resources associated with it.
+
+        Args:
+            session_id: The session identifier returned by open_session()
+
+        Raises:
+            ValueError: If the session ID is invalid
+            OperationalError: If there's an error closing the session
+        """
         pass
 
     # == Query Execution, Command Management ==
@@ -37,7 +84,7 @@ def execute_command(
         max_rows: int,
         max_bytes: int,
         lz4_compression: bool,
-        cursor: Any,
+        cursor: "Cursor",
         use_cloud_fetch: bool,
         parameters: List[ttypes.TSparkParameter],
         async_op: bool,
@@ -47,6 +94,19 @@ def execute_command(
 
     @abstractmethod
     def cancel_command(self, command_id: CommandId) -> None:
+        """
+        Cancels a running command or query.
+
+        This method attempts to cancel a command that is currently being executed.
+        It can be called from a different thread than the one executing the command.
+
+        Args:
+            command_id: The command identifier to cancel
+
+        Raises:
+            ValueError: If the command ID is invalid
+            OperationalError: If there's an error canceling the command
+        """
         pass
 
     @abstractmethod
@@ -82,7 +142,7 @@ def get_schemas(
         session_id: SessionId,
         max_rows: int,
         max_bytes: int,
-        cursor: Any,
+        cursor: "Cursor",
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
     ) -> "ResultSet":
@@ -94,7 +154,7 @@ def get_tables(
         session_id: SessionId,
         max_rows: int,
         max_bytes: int,
-        cursor: Any,
+        cursor: "Cursor",
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
@@ -108,35 +168,45 @@ def get_columns(
         session_id: SessionId,
         max_rows: int,
         max_bytes: int,
-        cursor: Any,
+        cursor: "Cursor",
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
         column_name: Optional[str] = None,
     ) -> "ResultSet":
         pass
 
-    # == Utility Methods ==
-    @abstractmethod
-    def handle_to_id(self, session_id: SessionId) -> Any:
-        pass
-
-    @abstractmethod
-    def handle_to_hex_id(self, session_id: SessionId) -> str:
-        pass
-
-    # Properties related to specific backend features
+    # == Properties ==
     @property
     @abstractmethod
     def staging_allowed_local_path(self) -> Union[None, str, List[str]]:
+        """
+        Gets the allowed local paths for staging operations.
+
+        Returns:
+            Union[None, str, List[str]]: The allowed local paths for staging operations,
+            or None if staging is not allowed
+        """
         pass
 
     @property
     @abstractmethod
     def ssl_options(self) -> SSLOptions:
+        """
+        Gets the SSL options for this client.
+
+        Returns:
+            SSLOptions: The SSL configuration options
+        """
         pass
 
     @property
     @abstractmethod
     def max_download_threads(self) -> int:
+        """
+        Gets the maximum number of download threads for cloud fetch operations.
+
+        Returns:
+            int: The maximum number of download threads
+        """
         pass

src/databricks/sql/backend/thrift_backend.py

@@ -5,7 +5,10 @@
 import time
 import uuid
 import threading
-from typing import List, Union, Any
+from typing import List, Union, Any, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from databricks.sql.client import Cursor
 
 from databricks.sql.thrift_api.TCLIService.ttypes import TOperationState
 from databricks.sql.backend.types import (
@@ -585,12 +588,12 @@ def open_session(self, session_configuration, catalog, schema) -> SessionId:
             response = self.make_request(self._client.OpenSession, open_session_req)
             self._check_initial_namespace(catalog, schema, response)
             self._check_protocol_version(response)
-            info = (
+            properties = (
                 {"serverProtocolVersion": response.serverProtocolVersion}
                 if response.serverProtocolVersion
                 else {}
             )
-            return SessionId.from_thrift_handle(response.sessionHandle, info)
+            return SessionId.from_thrift_handle(response.sessionHandle, properties)
         except:
             self._transport.close()
             raise
@@ -802,6 +805,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
             arrow_queue_opt = None
 
         command_id = CommandId.from_thrift_handle(resp.operationHandle)
+
         return ExecuteResponse(
             arrow_queue=arrow_queue_opt,
             status=operation_state,
@@ -814,7 +818,9 @@ def _results_message_to_execute_response(self, resp, operation_state):
             arrow_schema_bytes=schema_bytes,
         )
 
-    def get_execution_result(self, command_id: CommandId, cursor):
+    def get_execution_result(
+        self, command_id: CommandId, cursor: "Cursor"
+    ) -> ExecuteResponse:
         thrift_handle = command_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift command ID")
@@ -939,7 +945,7 @@ def execute_command(
         max_rows: int,
         max_bytes: int,
         lz4_compression: bool,
-        cursor: Any,
+        cursor: "Cursor",
         use_cloud_fetch=True,
         parameters=[],
         async_op=False,
@@ -1037,7 +1043,7 @@ def get_schemas(
         session_id: SessionId,
         max_rows: int,
         max_bytes: int,
-        cursor: Any,
+        cursor: "Cursor",
         catalog_name=None,
         schema_name=None,
     ) -> "ResultSet":
@@ -1071,7 +1077,7 @@ def get_tables(
         session_id: SessionId,
         max_rows: int,
         max_bytes: int,
-        cursor: Any,
+        cursor: "Cursor",
         catalog_name=None,
         schema_name=None,
         table_name=None,
@@ -1109,7 +1115,7 @@ def get_columns(
         session_id: SessionId,
         max_rows: int,
         max_bytes: int,
-        cursor: Any,
+        cursor: "Cursor",
         catalog_name=None,
         schema_name=None,
         table_name=None,
@@ -1230,18 +1236,4 @@ def close_command(self, command_id: CommandId) -> None:
         logger.debug("ThriftBackend.close_command(command_id=%s)", command_id)
         req = ttypes.TCloseOperationReq(operationHandle=thrift_handle)
         resp = self.make_request(self._client.CloseOperation, req)
-        logger.debug(
-            "ThriftBackend.close_command(command_id=%s) -> %s", command_id, resp
-        )
-
-    def handle_to_id(self, session_id: SessionId) -> Any:
-        """Get the raw session ID from a SessionId"""
-        if session_id.backend_type != BackendType.THRIFT:
-            raise ValueError("Not a valid Thrift session ID")
-        return session_id.guid
-
-    def handle_to_hex_id(self, session_id: SessionId) -> str:
-        """Get the hex representation of a session ID"""
-        if session_id.backend_type != BackendType.THRIFT:
-            raise ValueError("Not a valid Thrift session ID")
-        return guid_to_hex_id(session_id.guid)
+        return resp.status

src/databricks/sql/backend/types.py

@@ -78,7 +78,7 @@ def __init__(
         backend_type: BackendType,
         guid: Any,
         secret: Optional[Any] = None,
-        info: Optional[Dict[str, Any]] = None,
+        properties: Optional[Dict[str, Any]] = None,
     ):
         """
         Initialize a SessionId.
@@ -92,10 +92,28 @@ def __init__(
         self.backend_type = backend_type
         self.guid = guid
         self.secret = secret
-        self.info = info or {}
+        self.properties = properties or {}
+
+    def __str__(self) -> str:
+        """
+        Return a string representation of the SessionId.
+
+        For SEA backend, returns the guid.
+        For Thrift backend, returns a format like "guid|secret".
+
+        Returns:
+            A string representation of the session ID
+        """
+        if self.backend_type == BackendType.SEA:
+            return str(self.guid)
+        elif self.backend_type == BackendType.THRIFT:
+            return f"{self.get_hex_id()}|{guid_to_hex_id(self.secret) if isinstance(self.secret, bytes) else str(self.secret)}"
+        return str(self.guid)
 
     @classmethod
-    def from_thrift_handle(cls, session_handle, info: Optional[Dict[str, Any]] = None):
+    def from_thrift_handle(
+        cls, session_handle, properties: Optional[Dict[str, Any]] = None
+    ):
         """
         Create a SessionId from a Thrift session handle.
 
@@ -112,15 +130,15 @@ def from_thrift_handle(cls, session_handle, info: Optional[Dict[str, Any]] = Non
         secret_bytes = session_handle.sessionId.secret
 
         if session_handle.serverProtocolVersion is not None:
-            if info is None:
-                info = {}
-            info["serverProtocolVersion"] = session_handle.serverProtocolVersion
+            if properties is None:
+                properties = {}
+            properties["serverProtocolVersion"] = session_handle.serverProtocolVersion
 
-        return cls(BackendType.THRIFT, guid_bytes, secret_bytes, info)
+        return cls(BackendType.THRIFT, guid_bytes, secret_bytes, properties)
 
     @classmethod
     def from_sea_session_id(
-        cls, session_id: str, info: Optional[Dict[str, Any]] = None
+        cls, session_id: str, properties: Optional[Dict[str, Any]] = None
     ):
         """
         Create a SessionId from a SEA session ID.
@@ -131,7 +149,7 @@ def from_sea_session_id(
         Returns:
             A SessionId instance
         """
-        return cls(BackendType.SEA, session_id, info=info)
+        return cls(BackendType.SEA, session_id, properties=properties)
 
     def to_thrift_handle(self):
         """
@@ -146,7 +164,7 @@ def to_thrift_handle(self):
         from databricks.sql.thrift_api.TCLIService import ttypes
 
         handle_identifier = ttypes.THandleIdentifier(guid=self.guid, secret=self.secret)
-        server_protocol_version = self.info.get("serverProtocolVersion")
+        server_protocol_version = self.properties.get("serverProtocolVersion")
         return ttypes.TSessionHandle(
             sessionId=handle_identifier, serverProtocolVersion=server_protocol_version
         )
@@ -163,7 +181,13 @@ def to_sea_session_id(self):
 
         return self.guid
 
-    def to_hex_id(self) -> str:
+    def get_id(self) -> Any:
+        """
+        Get the ID of the session.
+        """
+        return self.guid
+
+    def get_hex_id(self) -> str:
         """
         Get a hexadecimal string representation of the session ID.
 
@@ -180,9 +204,10 @@ def get_protocol_version(self):
         Get the server protocol version for this session.
 
         Returns:
-            The server protocol version or None if this is not a Thrift session ID
+            The server protocol version or None if it does not exist
+            It is not expected to exist for SEA sessions.
         """
-        return self.info.get("serverProtocolVersion")
+        return self.properties.get("serverProtocolVersion")
 
 
 class CommandId:

src/databricks/sql/session.py

@@ -116,11 +116,11 @@ def get_session_id(self) -> SessionId:
 
     def get_id(self):
         """Get the raw session ID (backend-specific)"""
-        return self.backend.handle_to_id(self._session_id)
+        return self._session_id.get_id()
 
     def get_id_hex(self) -> str:
         """Get the session ID in hex format"""
-        return self.backend.handle_to_hex_id(self._session_id)
+        return self._session_id.get_hex_id()
 
     def close(self) -> None:
         """Close the underlying session."""