Merge remote-tracking branch 'origin/sea-migration' into fetch-interface

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -3,11 +3,9 @@
 
 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
@@ -110,11 +108,44 @@ def cancel_command(self, command_id: CommandId) -> None:
         pass
 
     @abstractmethod
-    def close_command(self, command_id: CommandId) -> None:
+    def close_command(self, command_id: CommandId) -> ttypes.TStatus:
+        """
+        Closes a command and releases associated resources.
+
+        This method informs the server that the client is done with the command
+        and any resources associated with it can be released.
+
+        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) -> CommandState:
+    def get_query_state(self, command_id: CommandId) -> ttypes.TOperationState:
+        """
+        Gets the current state of a query or command.
+
+        This method retrieves the current execution state of a command from the server.
+
+        Args:
+            command_id: The command identifier to check
+
+        Returns:
+            ttypes.TOperationState: The current state of the command
+
+        Raises:
+            ValueError: If the command ID is invalid
+            OperationalError: If there's an error retrieving the state
+            ServerOperationError: If the command is in an error state
+            DatabaseError: If the command has been closed unexpectedly
+        """
         pass
 
     @abstractmethod
@@ -173,30 +204,29 @@ def get_columns(
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
         column_name: Optional[str] = None,
-    ) -> "ResultSet":
-        pass
-
-    # == Properties ==
-    @property
-    @abstractmethod
-    def staging_allowed_local_path(self) -> Union[None, str, List[str]]:
+    ) -> ExecuteResponse:
         """
-        Gets the allowed local paths for staging operations.
+        Retrieves a list of columns, optionally filtered by catalog, schema, table, and column name patterns.
 
-        Returns:
-            Union[None, str, List[str]]: The allowed local paths for staging operations,
-            or None if staging is not allowed
-        """
-        pass
+        This method fetches metadata about columns available in the specified table,
+        or all tables if not specified.
 
-    @property
-    @abstractmethod
-    def ssl_options(self) -> SSLOptions:
-        """
-        Gets the SSL options for this client.
+        Args:
+            session_id: The session identifier
+            max_rows: Maximum number of rows to fetch in a single batch
+            max_bytes: Maximum number of bytes to fetch in a single batch
+            cursor: The cursor object that will handle the results
+            catalog_name: Optional catalog name pattern to filter by
+            schema_name: Optional schema name pattern to filter by
+            table_name: Optional table name pattern to filter by
+            column_name: Optional column name pattern to filter by
 
         Returns:
-            SSLOptions: The SSL configuration options
+            ExecuteResponse: An object containing the column metadata
+
+        Raises:
+            ValueError: If the session ID is invalid
+            OperationalError: If there's an error retrieving the columns
         """
         pass
 

src/databricks/sql/backend/thrift_backend.py

@@ -5,19 +5,18 @@
 import time
 import uuid
 import threading
-from typing import List, Union, Any, TYPE_CHECKING
+from typing import List, Optional, 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 (
-    CommandState,
     SessionId,
     CommandId,
     BackendType,
-    guid_to_hex_id,
 )
+from databricks.sql.backend.utils import guid_to_hex_id
 
 try:
     import pyarrow
@@ -103,7 +102,6 @@ def __init__(
         http_headers,
         auth_provider: AuthProvider,
         ssl_options: SSLOptions,
-        staging_allowed_local_path: Union[None, str, List[str]] = None,
         **kwargs,
     ):
         # Internal arguments in **kwargs:
@@ -162,7 +160,6 @@ def __init__(
         else:
             raise ValueError("No valid connection settings.")
 
-        self._staging_allowed_local_path = staging_allowed_local_path
         self._initialize_retry_args(kwargs)
         self._use_arrow_native_complex_types = kwargs.get(
             "_use_arrow_native_complex_types", True
@@ -236,14 +233,6 @@ def __init__(
 
         self._request_lock = threading.RLock()
 
-    @property
-    def staging_allowed_local_path(self) -> Union[None, str, List[str]]:
-        return self._staging_allowed_local_path
-
-    @property
-    def ssl_options(self) -> SSLOptions:
-        return self._ssl_options
-
     @property
     def max_download_threads(self) -> int:
         return self._max_download_threads

src/databricks/sql/backend/types.py

@@ -1,65 +1,16 @@
 from enum import Enum
 from typing import Dict, Optional, Any, Union
-import uuid
 import logging
 
-from databricks.sql.thrift_api.TCLIService import ttypes
+from databricks.sql.backend.utils import guid_to_hex_id
 
 logger = logging.getLogger(__name__)
 
 
-class CommandState(Enum):
-    PENDING = "PENDING"
-    RUNNING = "RUNNING"
-    SUCCEEDED = "SUCCEEDED"
-    FAILED = "FAILED"
-    CLOSED = "CLOSED"
-    CANCELLED = "CANCELLED"
-
-    @classmethod
-    def from_thrift_state(cls, state: ttypes.TOperationState) -> "CommandState":
-        if state in (
-            ttypes.TOperationState.INITIALIZED_STATE,
-            ttypes.TOperationState.PENDING_STATE,
-        ):
-            return cls.PENDING
-        elif state == ttypes.TOperationState.RUNNING_STATE:
-            return cls.RUNNING
-        elif state == ttypes.TOperationState.FINISHED_STATE:
-            return cls.SUCCEEDED
-        elif state in (
-            ttypes.TOperationState.ERROR_STATE,
-            ttypes.TOperationState.TIMEDOUT_STATE,
-            ttypes.TOperationState.UKNOWN_STATE,
-        ):
-            return cls.FAILED
-        elif state == ttypes.TOperationState.CLOSED_STATE:
-            return cls.CLOSED
-        elif state == ttypes.TOperationState.CANCELED_STATE:
-            return cls.CANCELLED
-        else:
-            raise ValueError(f"Unknown command state: {state}")
-
-
-def guid_to_hex_id(guid: bytes) -> str:
-    """Return a hexadecimal string instead of bytes
-
-    Example:
-        IN   b'\x01\xee\x1d)\xa4\x19\x1d\xb6\xa9\xc0\x8d\xf1\xfe\xbaB\xdd'
-        OUT  '01ee1d29-a419-1db6-a9c0-8df1feba42dd'
-
-    If conversion to hexadecimal fails, the original bytes are returned
-    """
-    try:
-        this_uuid = uuid.UUID(bytes=guid)
-    except Exception as e:
-        logger.debug(f"Unable to convert bytes to UUID: {guid!r} -- {str(e)}")
-        return str(guid)
-    return str(this_uuid)
-
-
 class BackendType(Enum):
-    """Enum representing the type of backend."""
+    """
+    Enum representing the type of backend
+    """
 
     THRIFT = "thrift"
     SEA = "sea"
@@ -87,7 +38,7 @@ def __init__(
             backend_type: The type of backend (THRIFT or SEA)
             guid: The primary identifier for the session
             secret: The secret part of the identifier (only used for Thrift)
-            info: Additional information about the session
+            properties: Additional information about the session
         """
         self.backend_type = backend_type
         self.guid = guid
@@ -107,7 +58,12 @@ def __str__(self) -> str:
         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)}"
+            secret_hex = (
+                guid_to_hex_id(self.secret)
+                if isinstance(self.secret, bytes)
+                else str(self.secret)
+            )
+            return f"{self.get_hex_guid()}|{secret_hex}"
         return str(self.guid)
 
     @classmethod
@@ -181,13 +137,13 @@ def to_sea_session_id(self):
 
         return self.guid
 
-    def get_id(self) -> Any:
+    def get_guid(self) -> Any:
         """
         Get the ID of the session.
         """
         return self.guid
 
-    def get_hex_id(self) -> str:
+    def get_hex_guid(self) -> str:
         """
         Get a hexadecimal string representation of the session ID.
 
@@ -316,7 +272,7 @@ def to_sea_statement_id(self):
 
         return self.guid
 
-    def to_hex_id(self) -> str:
+    def to_hex_guid(self) -> str:
         """
         Get a hexadecimal string representation of the command ID.
 

src/databricks/sql/backend/utils/__init__.py

@@ -0,0 +1,3 @@
+from .guid_utils import guid_to_hex_id
+
+__all__ = ["guid_to_hex_id"]

src/databricks/sql/backend/utils/guid_utils.py

@@ -0,0 +1,22 @@
+import uuid
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def guid_to_hex_id(guid: bytes) -> str:
+    """Return a hexadecimal string instead of bytes
+
+    Example:
+        IN   b'\x01\xee\x1d)\xa4\x19\x1d\xb6\xa9\xc0\x8d\xf1\xfe\xbaB\xdd'
+        OUT  '01ee1d29-a419-1db6-a9c0-8df1feba42dd'
+
+    If conversion to hexadecimal fails, a string representation of the original
+    bytes is returned
+    """
+    try:
+        this_uuid = uuid.UUID(bytes=guid)
+    except Exception as e:
+        logger.debug(f"Unable to convert bytes to UUID: {guid!r} -- {str(e)}")
+        return str(guid)
+    return str(this_uuid)

src/databricks/sql/client.py

@@ -247,6 +247,7 @@ def read(self) -> Optional[OAuthToken]:
         self.use_inline_params = self._set_use_inline_params_with_warning(
             kwargs.get("use_inline_params", False)
         )
+        self.staging_allowed_local_path = kwargs.get("staging_allowed_local_path", None)
 
     def _set_use_inline_params_with_warning(self, value: Union[bool, str]):
         """Valid values are True, False, and "silent"
@@ -785,7 +786,7 @@ def execute(
 
         if self.active_result_set and self.active_result_set.is_staging_operation:
             self._handle_staging_operation(
-                staging_allowed_local_path=self.backend.staging_allowed_local_path
+                staging_allowed_local_path=self.connection.staging_allowed_local_path
             )
 
         return self
@@ -881,7 +882,7 @@ def get_async_execution_result(self):
 
             if self.active_result_set and self.active_result_set.is_staging_operation:
                 self._handle_staging_operation(
-                    staging_allowed_local_path=self.backend.staging_allowed_local_path
+                    staging_allowed_local_path=self.connection.staging_allowed_local_path
                 )
 
             return self
@@ -1106,7 +1107,7 @@ def query_id(self) -> Optional[str]:
         invoked via the execute method yet, or if cursor was closed.
         """
         if self.active_command_id is not None:
-            return self.active_command_id.to_hex_id()
+            return self.active_command_id.to_hex_guid()
         return None
 
     @property

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._session_id.get_id()
+        return self._session_id.get_guid()
 
     def get_id_hex(self) -> str:
         """Get the session ID in hex format"""
-        return self._session_id.get_hex_id()
+        return self._session_id.get_hex_guid()
 
     def close(self) -> None:
         """Close the underlying session."""

tests/unit/test_client.py

@@ -133,6 +133,7 @@ def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
 
         result_set = ThriftResultSet(
             connection=mock_connection,
+            backend=mock_backend,
             execute_response=Mock(),
             thrift_client=mock_backend,
         )

tests/unit/test_thrift_backend.py

@@ -53,6 +53,7 @@ class ThriftBackendTestSuite(unittest.TestCase):
     open_session_resp = ttypes.TOpenSessionResp(
         status=okay_status,
         serverProtocolVersion=ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V4,
+        sessionHandle=session_handle,
     )
 
     metadata_resp = ttypes.TGetResultSetMetadataResp(