[squashed from sea-exec] merge sea stuffs

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

Author: varun-edachali-dbx

examples/experimental/sea_connector_test.py

@@ -22,99 +22,90 @@
     "test_sea_metadata",
 ]
 
-
 def load_test_function(module_name: str) -> Callable:
     """Load a test function from a module."""
     module_path = os.path.join(
-        os.path.dirname(os.path.abspath(__file__)), "tests", f"{module_name}.py"
+        os.path.dirname(os.path.abspath(__file__)),
+        "tests",
+        f"{module_name}.py"
     )
-
+    
     spec = importlib.util.spec_from_file_location(module_name, module_path)
     module = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(module)
-
+    
     # Get the main test function (assuming it starts with "test_")
     for name in dir(module):
         if name.startswith("test_") and callable(getattr(module, name)):
             # For sync and async query modules, we want the main function that runs both tests
             if name == f"test_sea_{module_name.replace('test_sea_', '')}_exec":
                 return getattr(module, name)
-
+    
     # Fallback to the first test function found
     for name in dir(module):
         if name.startswith("test_") and callable(getattr(module, name)):
             return getattr(module, name)
-
+    
     raise ValueError(f"No test function found in module {module_name}")
 
-
 def run_tests() -> List[Tuple[str, bool]]:
     """Run all tests and return results."""
     results = []
-
+    
     for module_name in TEST_MODULES:
         try:
             test_func = load_test_function(module_name)
             logger.info(f"\n{'=' * 50}")
             logger.info(f"Running test: {module_name}")
             logger.info(f"{'-' * 50}")
-
+            
             success = test_func()
             results.append((module_name, success))
-
+            
             status = "✅ PASSED" if success else "❌ FAILED"
             logger.info(f"Test {module_name}: {status}")
-
+            
         except Exception as e:
             logger.error(f"Error loading or running test {module_name}: {str(e)}")
             import traceback
-
             logger.error(traceback.format_exc())
             results.append((module_name, False))
-
+    
     return results
 
-
 def print_summary(results: List[Tuple[str, bool]]) -> None:
     """Print a summary of test results."""
     logger.info(f"\n{'=' * 50}")
     logger.info("TEST SUMMARY")
     logger.info(f"{'-' * 50}")
-
+    
     passed = sum(1 for _, success in results if success)
     total = len(results)
-
+    
     for module_name, success in results:
         status = "✅ PASSED" if success else "❌ FAILED"
         logger.info(f"{status} - {module_name}")
-
+    
     logger.info(f"{'-' * 50}")
     logger.info(f"Total: {total} | Passed: {passed} | Failed: {total - passed}")
     logger.info(f"{'=' * 50}")
 
-
 if __name__ == "__main__":
     # Check if required environment variables are set
-    required_vars = [
-        "DATABRICKS_SERVER_HOSTNAME",
-        "DATABRICKS_HTTP_PATH",
-        "DATABRICKS_TOKEN",
-    ]
+    required_vars = ["DATABRICKS_SERVER_HOSTNAME", "DATABRICKS_HTTP_PATH", "DATABRICKS_TOKEN"]
     missing_vars = [var for var in required_vars if not os.environ.get(var)]
-
+    
     if missing_vars:
-        logger.error(
-            f"Missing required environment variables: {', '.join(missing_vars)}"
-        )
+        logger.error(f"Missing required environment variables: {', '.join(missing_vars)}")
         logger.error("Please set these variables before running the tests.")
         sys.exit(1)
-
+    
     # Run all tests
     results = run_tests()
-
+    
     # Print summary
     print_summary(results)
-
+    
     # Exit with appropriate status code
     all_passed = all(success for _, success in results)
     sys.exit(0 if all_passed else 1)

examples/experimental/tests/__init__.py

@@ -0,0 +1 @@
+# This file makes the tests directory a Python package

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
+        )