[squashed from exec-sea] init execution func

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

examples/experimental/sea_connector_test.py

@@ -6,34 +6,122 @@
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
+
+def test_sea_query_exec():
+    """
+    Test executing a query using the SEA backend with result compression.
+
+    This function connects to a Databricks SQL endpoint using the SEA backend,
+    executes a simple query with result compression enabled and disabled,
+    and verifies that execution completes successfully.
+    """
+    server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME")
+    http_path = os.environ.get("DATABRICKS_HTTP_PATH")
+    access_token = os.environ.get("DATABRICKS_TOKEN")
+    catalog = os.environ.get("DATABRICKS_CATALOG")
+
+    if not all([server_hostname, http_path, access_token]):
+        logger.error("Missing required environment variables.")
+        logger.error(
+            "Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN."
+        )
+        sys.exit(1)
+
+    try:
+        # Test with compression enabled
+        logger.info("Creating connection with LZ4 compression enabled")
+        connection = Connection(
+            server_hostname=server_hostname,
+            http_path=http_path,
+            access_token=access_token,
+            catalog=catalog,
+            schema="default",
+            use_sea=True,
+            user_agent_entry="SEA-Test-Client",
+            use_cloud_fetch=True,  # Enable cloud fetch to use compression
+            enable_query_result_lz4_compression=True,  # Enable LZ4 compression
+        )
+
+        logger.info(
+            f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
+        )
+        logger.info(f"backend type: {type(connection.session.backend)}")
+
+        # Execute a simple query with compression enabled
+        cursor = connection.cursor(arraysize=0, buffer_size_bytes=0)
+        logger.info("Executing query with LZ4 compression: SELECT 1 as test_value")
+        cursor.execute("SELECT 1 as test_value")
+        logger.info("Query with compression executed successfully")
+        cursor.close()
+        connection.close()
+        logger.info("Successfully closed SEA session with compression enabled")
+
+        # Test with compression disabled
+        logger.info("Creating connection with LZ4 compression disabled")
+        connection = Connection(
+            server_hostname=server_hostname,
+            http_path=http_path,
+            access_token=access_token,
+            catalog=catalog,
+            schema="default",
+            use_sea=True,
+            user_agent_entry="SEA-Test-Client",
+            use_cloud_fetch=False,  # Enable cloud fetch
+            enable_query_result_lz4_compression=False,  # Disable LZ4 compression
+        )
+
+        logger.info(
+            f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
+        )
+
+        # Execute a simple query with compression disabled
+        cursor = connection.cursor(arraysize=0, buffer_size_bytes=0)
+        logger.info("Executing query without compression: SELECT 1 as test_value")
+        cursor.execute("SELECT 1 as test_value")
+        logger.info("Query without compression executed successfully")
+        cursor.close()
+        connection.close()
+        logger.info("Successfully closed SEA session with compression disabled")
+
+    except Exception as e:
+        logger.error(f"Error during SEA query execution test: {str(e)}")
+        import traceback
+
+        logger.error(traceback.format_exc())
+        sys.exit(1)
+
+    logger.info("SEA query execution test with compression completed successfully")
+
+
 def test_sea_session():
     """
     Test opening and closing a SEA session using the connector.
-    
+
     This function connects to a Databricks SQL endpoint using the SEA backend,
     opens a session, and then closes it.
-    
+
     Required environment variables:
     - DATABRICKS_SERVER_HOSTNAME: Databricks server hostname
     - DATABRICKS_HTTP_PATH: HTTP path for the SQL endpoint
     - DATABRICKS_TOKEN: Personal access token for authentication
     """
-
     server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME")
     http_path = os.environ.get("DATABRICKS_HTTP_PATH")
     access_token = os.environ.get("DATABRICKS_TOKEN")
     catalog = os.environ.get("DATABRICKS_CATALOG")
-    
+
     if not all([server_hostname, http_path, access_token]):
         logger.error("Missing required environment variables.")
-        logger.error("Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN.")
+        logger.error(
+            "Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN."
+        )
         sys.exit(1)
-    
+
     logger.info(f"Connecting to {server_hostname}")
     logger.info(f"HTTP Path: {http_path}")
     if catalog:
         logger.info(f"Using catalog: {catalog}")
-    
+
     try:
         logger.info("Creating connection with SEA backend...")
         connection = Connection(
@@ -42,25 +130,33 @@ def test_sea_session():
             access_token=access_token,
             catalog=catalog,
             schema="default",
-            use_sea=True, 
-            user_agent_entry="SEA-Test-Client" # add custom user agent
+            use_sea=True,
+            user_agent_entry="SEA-Test-Client",  # add custom user agent
+        )
+
+        logger.info(
+            f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
         )
-        
-        logger.info(f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}")
         logger.info(f"backend type: {type(connection.session.backend)}")
-        
+
         # Close the connection
         logger.info("Closing the SEA session...")
         connection.close()
         logger.info("Successfully closed SEA session")
-        
+
     except Exception as e:
         logger.error(f"Error testing SEA session: {str(e)}")
         import traceback
+
         logger.error(traceback.format_exc())
         sys.exit(1)
-    
+
     logger.info("SEA session test completed successfully")
 
+
 if __name__ == "__main__":
+    # Test session management
     test_sea_session()
+
+    # Test query execution with compression
+    test_sea_query_exec()

src/databricks/sql/backend/databricks_client.py

@@ -86,34 +86,6 @@ def execute_command(
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
     ) -> Union["ResultSet", None]:
-        """
-        Executes a SQL command or query within the specified session.
-
-        This method sends a SQL command to the server for execution and handles
-        the response. It can operate in both synchronous and asynchronous modes.
-
-        Args:
-            operation: The SQL command or query to execute
-            session_id: The session identifier in which to execute the command
-            max_rows: Maximum number of rows to fetch in a single fetch batch
-            max_bytes: Maximum number of bytes to fetch in a single fetch batch
-            lz4_compression: Whether to use LZ4 compression for result data
-            cursor: The cursor object that will handle the results
-            use_cloud_fetch: Whether to use cloud fetch for retrieving large result sets
-            parameters: List of parameters to bind to the query
-            async_op: Whether to execute the command asynchronously
-            enforce_embedded_schema_correctness: Whether to enforce schema correctness
-
-        Returns:
-            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().
-
-        Raises:
-            ValueError: If the session ID is invalid
-            OperationalError: If there's an error executing the command
-            ServerOperationError: If the server encounters an error during execution
-        """
         pass
 
     @abstractmethod

src/databricks/sql/backend/filters.py

@@ -0,0 +1,143 @@
+"""
+Client-side filtering utilities for Databricks SQL connector.
+
+This module provides filtering capabilities for result sets returned by different backends.
+"""
+
+import logging
+from typing import (
+    List,
+    Optional,
+    Any,
+    Callable,
+    TYPE_CHECKING,
+)
+
+if TYPE_CHECKING:
+    from databricks.sql.result_set import ResultSet
+
+from databricks.sql.result_set import SeaResultSet
+
+logger = logging.getLogger(__name__)
+
+
+class ResultSetFilter:
+    """
+    A general-purpose filter for result sets that can be applied to any backend.
+
+    This class provides methods to filter result sets based on various criteria,
+    similar to the client-side filtering in the JDBC connector.
+    """
+
+    @staticmethod
+    def _filter_sea_result_set(
+        result_set: "SeaResultSet", filter_func: Callable[[List[Any]], bool]
+    ) -> "SeaResultSet":
+        """
+        Filter a SEA result set using the provided filter function.
+
+        Args:
+            result_set: The SEA result set to filter
+            filter_func: Function that takes a row and returns True if the row should be included
+
+        Returns:
+            A filtered SEA result set
+        """
+        # Create a filtered version of the result set
+        filtered_response = result_set._response.copy()
+
+        # If there's a result with rows, filter them
+        if (
+            "result" in filtered_response
+            and "data_array" in filtered_response["result"]
+        ):
+            rows = filtered_response["result"]["data_array"]
+            filtered_rows = [row for row in rows if filter_func(row)]
+            filtered_response["result"]["data_array"] = filtered_rows
+
+            # Update row count if present
+            if "row_count" in filtered_response["result"]:
+                filtered_response["result"]["row_count"] = len(filtered_rows)
+
+        # Create a new result set with the filtered data
+        return SeaResultSet(
+            connection=result_set.connection,
+            sea_response=filtered_response,
+            sea_client=result_set.backend,
+            buffer_size_bytes=result_set.buffer_size_bytes,
+            arraysize=result_set.arraysize,
+        )
+
+    @staticmethod
+    def filter_by_column_values(
+        result_set: "ResultSet",
+        column_index: int,
+        allowed_values: List[str],
+        case_sensitive: bool = False,
+    ) -> "ResultSet":
+        """
+        Filter a result set by values in a specific column.
+
+        Args:
+            result_set: The result set to filter
+            column_index: The index of the column to filter on
+            allowed_values: List of allowed values for the column
+            case_sensitive: Whether to perform case-sensitive comparison
+
+        Returns:
+            A filtered result set
+        """
+        # Convert to uppercase for case-insensitive comparison if needed
+        if not case_sensitive:
+            allowed_values = [v.upper() for v in allowed_values]
+
+        # Determine the type of result set and apply appropriate filtering
+        if isinstance(result_set, SeaResultSet):
+            return ResultSetFilter._filter_sea_result_set(
+                result_set,
+                lambda row: (
+                    len(row) > column_index
+                    and isinstance(row[column_index], str)
+                    and (
+                        row[column_index].upper()
+                        if not case_sensitive
+                        else row[column_index]
+                    )
+                    in allowed_values
+                ),
+            )
+
+        # For other result set types, return the original (should be handled by specific implementations)
+        logger.warning(
+            f"Filtering not implemented for result set type: {type(result_set).__name__}"
+        )
+        return result_set
+
+    @staticmethod
+    def filter_tables_by_type(
+        result_set: "ResultSet", table_types: Optional[List[str]] = None
+    ) -> "ResultSet":
+        """
+        Filter a result set of tables by the specified table types.
+
+        This is a client-side filter that processes the result set after it has been
+        retrieved from the server. It filters out tables whose type does not match
+        any of the types in the table_types list.
+
+        Args:
+            result_set: The original result set containing tables
+            table_types: List of table types to include (e.g., ["TABLE", "VIEW"])
+
+        Returns:
+            A filtered result set containing only tables of the specified types
+        """
+        # Default table types if none specified
+        DEFAULT_TABLE_TYPES = ["TABLE", "VIEW", "SYSTEM TABLE"]
+        valid_types = (
+            table_types if table_types and len(table_types) > 0 else DEFAULT_TABLE_TYPES
+        )
+
+        # Table type is typically in the 6th column (index 5)
+        return ResultSetFilter.filter_by_column_values(
+            result_set, 5, valid_types, case_sensitive=False
+        )