filters and sea_result_set unit tests

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

Author: varun-edachali-dbx

src/databricks/sql/backend/filters.py

@@ -11,120 +11,128 @@
 from databricks.sql.backend.sea_result_set import SeaResultSet
 
 # Type variable for the result set type
-T = TypeVar('T')
+T = TypeVar("T")
 
 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_by_column_values(
         result_set: Any,
         column_index: int,
         allowed_values: List[str],
-        case_sensitive: bool = False
+        case_sensitive: bool = False,
     ) -> Any:
         """
         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, 
+                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
-                )
+                    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: Any,
-        table_types: Optional[List[str]] = None
+        result_set: Any, table_types: Optional[List[str]] = None
     ) -> Any:
         """
         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
-        
+        valid_types = (
+            table_types if table_types and len(table_types) > 0 else DEFAULT_TABLE_TYPES
+        )
+
         # Table type is typically in the 4th column (index 3)
-        return ResultSetFilter.filter_by_column_values(result_set, 3, valid_types, case_sensitive=False)
-    
+        return ResultSetFilter.filter_by_column_values(
+            result_set, 3, valid_types, case_sensitive=False
+        )
+
     @staticmethod
-    def _filter_sea_result_set(result_set: Any, filter_func: Callable[[List[Any]], bool]) -> Any:
+    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
         """
-        # Type checking for SeaResultSet
-        if not isinstance(result_set, SeaResultSet):
-            return result_set
-            
         # Create a filtered version of the result set
-        sea_result = cast(SeaResultSet, result_set)
-        filtered_response = sea_result._response.copy()  # type: ignore
-        
+        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"]:
+        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=sea_result.connection,
+            connection=result_set.connection,
             sea_response=filtered_response,
-            sea_client=sea_result._sea_client,  # type: ignore
-            buffer_size_bytes=sea_result._buffer_size_bytes,  # type: ignore
-            arraysize=sea_result._arraysize,  # type: ignore
-        )
+            sea_client=result_set.backend,
+            buffer_size_bytes=result_set.buffer_size_bytes,
+            arraysize=result_set.arraysize,
+        )

src/databricks/sql/backend/sea_backend.py

@@ -555,6 +555,7 @@ def get_tables(
 
         # Apply client-side filtering by table_types if specified
         from databricks.sql.backend.filters import ResultSetFilter
+
         result = ResultSetFilter.filter_tables_by_type(result, table_types)
 
         return result

src/databricks/sql/backend/sea_result_set.py

@@ -39,6 +39,12 @@ def __init__(
         """Initialize a SeaResultSet with the response from a SEA query execution."""
         super().__init__(connection, sea_client, arraysize, buffer_size_bytes)
 
+        # Store the original response for filtering operations
+        self._response = sea_response
+        self._sea_client = sea_client
+        self._buffer_size_bytes = buffer_size_bytes
+        self._arraysize = arraysize
+
         # Extract and store SEA-specific properties
         self.statement_id = sea_response.get("statement_id")
 
@@ -89,7 +95,7 @@ def __init__(
         result_data = sea_response.get("result")
         if result_data:
             self.result: Optional[ResultData] = ResultData(
-                data=result_data.get("data"),
+                data=result_data.get("data_array"),  # Changed from data to data_array
                 external_links=result_data.get("external_links"),
             )
         else:

tests/unit/test_filters.py

@@ -28,10 +28,10 @@ def setUp(self):
                     ["catalog1", "schema1", "table3", "SYSTEM TABLE", ""],
                     ["catalog1", "schema1", "table4", "EXTERNAL TABLE", ""],
                 ],
-                "row_count": 4
+                "row_count": 4,
             }
         }
-        
+
         # Set up the connection and other required attributes
         self.mock_sea_result_set.connection = MagicMock()
         self.mock_sea_result_set._sea_client = MagicMock()
@@ -44,24 +44,26 @@ def test_filter_tables_by_type(self, mock_sea_result_set_class):
         # Set up the mock to return a new mock when instantiated
         mock_instance = MagicMock()
         mock_sea_result_set_class.return_value = mock_instance
-        
+
         # Test with specific table types
         table_types = ["TABLE", "VIEW"]
-        
+
         # Make the mock_sea_result_set appear to be a SeaResultSet
         with patch("databricks.sql.backend.filters.isinstance", return_value=True):
-            result = ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, table_types)
-        
+            result = ResultSetFilter.filter_tables_by_type(
+                self.mock_sea_result_set, table_types
+            )
+
         # Verify the filter was applied correctly
         mock_sea_result_set_class.assert_called_once()
         call_kwargs = mock_sea_result_set_class.call_args[1]
-        
+
         # Check that the filtered response contains only TABLE and VIEW
         filtered_data = call_kwargs["sea_response"]["result"]["data_array"]
         self.assertEqual(len(filtered_data), 2)
         self.assertEqual(filtered_data[0][3], "TABLE")
         self.assertEqual(filtered_data[1][3], "VIEW")
-        
+
         # Check row count was updated
         self.assertEqual(call_kwargs["sea_response"]["result"]["row_count"], 2)
 
@@ -71,18 +73,20 @@ def test_filter_tables_by_type_case_insensitive(self, mock_sea_result_set_class)
         # Set up the mock to return a new mock when instantiated
         mock_instance = MagicMock()
         mock_sea_result_set_class.return_value = mock_instance
-        
+
         # Test with lowercase table types
         table_types = ["table", "view"]
-        
+
         # Make the mock_sea_result_set appear to be a SeaResultSet
         with patch("databricks.sql.backend.filters.isinstance", return_value=True):
-            result = ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, table_types)
-        
+            result = ResultSetFilter.filter_tables_by_type(
+                self.mock_sea_result_set, table_types
+            )
+
         # Verify the filter was applied correctly
         mock_sea_result_set_class.assert_called_once()
         call_kwargs = mock_sea_result_set_class.call_args[1]
-        
+
         # Check that the filtered response contains only TABLE and VIEW
         filtered_data = call_kwargs["sea_response"]["result"]["data_array"]
         self.assertEqual(len(filtered_data), 2)
@@ -95,15 +99,17 @@ def test_filter_tables_by_type_default(self, mock_sea_result_set_class):
         # Set up the mock to return a new mock when instantiated
         mock_instance = MagicMock()
         mock_sea_result_set_class.return_value = mock_instance
-        
+
         # Make the mock_sea_result_set appear to be a SeaResultSet
         with patch("databricks.sql.backend.filters.isinstance", return_value=True):
-            result = ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, None)
-        
+            result = ResultSetFilter.filter_tables_by_type(
+                self.mock_sea_result_set, None
+            )
+
         # Verify the filter was applied correctly
         mock_sea_result_set_class.assert_called_once()
         call_kwargs = mock_sea_result_set_class.call_args[1]
-        
+
         # Check that the filtered response contains TABLE, VIEW, and SYSTEM TABLE
         filtered_data = call_kwargs["sea_response"]["result"]["data_array"]
         self.assertEqual(len(filtered_data), 3)
@@ -113,4 +119,4 @@ def test_filter_tables_by_type_default(self, mock_sea_result_set_class):
 
 
 if __name__ == "__main__":
-    unittest.main()
+    unittest.main()