Merge branch 'sea-migration' into exec-sea

Author: varun-edachali-dbx

examples/experimental/sea_connector_test.py

@@ -159,4 +159,4 @@ def test_sea_session():
     test_sea_session()
 
     # Test query execution with compression
-    test_sea_query_execution_with_compression()
+    test_sea_query_execution_with_compression()

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
@@ -111,10 +109,40 @@ def cancel_command(self, command_id: CommandId) -> None:
 
     @abstractmethod
     def close_command(self, command_id: CommandId) -> None:
+        """
+        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
+
+        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:
+        """
+        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:
+            CommandState: 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
@@ -123,6 +151,23 @@ def get_execution_result(
         command_id: CommandId,
         cursor: "Cursor",
     ) -> "ResultSet":
+        """
+        Retrieves the results of a previously executed command.
+
+        This method fetches the results of a command that was executed asynchronously
+        or retrieves additional results from a command that has more rows available.
+
+        Args:
+            command_id: The command identifier for which to retrieve results
+            cursor: The cursor object that will handle the results
+
+        Returns:
+            ResultSet: An object containing the query results and metadata
+
+        Raises:
+            ValueError: If the command ID is invalid
+            OperationalError: If there's an error retrieving the results
+        """
         pass
 
     # == Metadata Operations ==
@@ -134,6 +179,25 @@ def get_catalogs(
         max_bytes: int,
         cursor: "Cursor",
     ) -> "ResultSet":
+        """
+        Retrieves a list of available catalogs.
+
+        This method fetches metadata about all catalogs available in the current
+        session's context.
+
+        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
+
+        Returns:
+            ResultSet: An object containing the catalog metadata
+
+        Raises:
+            ValueError: If the session ID is invalid
+            OperationalError: If there's an error retrieving the catalogs
+        """
         pass
 
     @abstractmethod
@@ -146,6 +210,27 @@ def get_schemas(
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
     ) -> "ResultSet":
+        """
+        Retrieves a list of schemas, optionally filtered by catalog and schema name patterns.
+
+        This method fetches metadata about schemas available in the specified catalog
+        or all catalogs if no catalog is specified.
+
+        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
+
+        Returns:
+            ResultSet: An object containing the schema metadata
+
+        Raises:
+            ValueError: If the session ID is invalid
+            OperationalError: If there's an error retrieving the schemas
+        """
         pass
 
     @abstractmethod
@@ -160,6 +245,29 @@ def get_tables(
         table_name: Optional[str] = None,
         table_types: Optional[List[str]] = None,
     ) -> "ResultSet":
+        """
+        Retrieves a list of tables, optionally filtered by catalog, schema, table name, and table types.
+
+        This method fetches metadata about tables available in the specified catalog
+        and schema, or all catalogs and schemas if not specified.
+
+        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
+            table_types: Optional list of table types to filter by (e.g., ['TABLE', 'VIEW'])
+
+        Returns:
+            ResultSet: An object containing the table metadata
+
+        Raises:
+            ValueError: If the session ID is invalid
+            OperationalError: If there's an error retrieving the tables
+        """
         pass
 
     @abstractmethod
@@ -174,29 +282,28 @@ def get_columns(
         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]]:
         """
-        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
+            ResultSet: An object containing the column metadata
+
+        Raises:
+            ValueError: If the session ID is invalid
+            OperationalError: If there's an error retrieving the columns
         """
         pass