formatting (black)

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

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -801,16 +801,19 @@ def _results_message_to_execute_response(self, resp, operation_state):
         if status is None:
             raise ValueError(f"Invalid operation state: {operation_state}")
 
-        return ExecuteResponse(
-            command_id=command_id,
-            status=status,
-            description=description,
-            has_more_rows=has_more_rows,
-            results_queue=arrow_queue_opt,
-            has_been_closed_server_side=has_been_closed_server_side,
-            lz4_compressed=lz4_compressed,
-            is_staging_operation=is_staging_operation,
-        ), schema_bytes
+        return (
+            ExecuteResponse(
+                command_id=command_id,
+                status=status,
+                description=description,
+                has_more_rows=has_more_rows,
+                results_queue=arrow_queue_opt,
+                has_been_closed_server_side=has_been_closed_server_side,
+                lz4_compressed=lz4_compressed,
+                is_staging_operation=is_staging_operation,
+            ),
+            schema_bytes,
+        )
 
     def get_execution_result(
         self, command_id: CommandId, cursor: "Cursor"
@@ -881,7 +884,7 @@ def get_execution_result(
             buffer_size_bytes=cursor.buffer_size_bytes,
             arraysize=cursor.arraysize,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
-            arrow_schema_bytes=schema_bytes
+            arrow_schema_bytes=schema_bytes,
         )
 
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
@@ -991,7 +994,9 @@ def execute_command(
             self._handle_execute_response_async(resp, cursor)
             return None
         else:
-            execute_response, arrow_schema_bytes = self._handle_execute_response(resp, cursor)
+            execute_response, arrow_schema_bytes = self._handle_execute_response(
+                resp, cursor
+            )
 
             return ThriftResultSet(
                 connection=cursor.connection,
@@ -1000,7 +1005,7 @@ def execute_command(
                 buffer_size_bytes=max_bytes,
                 arraysize=max_rows,
                 use_cloud_fetch=use_cloud_fetch,
-                arrow_schema_bytes=arrow_schema_bytes
+                arrow_schema_bytes=arrow_schema_bytes,
             )
 
     def get_catalogs(
@@ -1022,7 +1027,9 @@ def get_catalogs(
         )
         resp = self.make_request(self._client.GetCatalogs, req)
 
-        execute_response, arrow_schema_bytes = self._handle_execute_response(resp, cursor)
+        execute_response, arrow_schema_bytes = self._handle_execute_response(
+            resp, cursor
+        )
 
         return ThriftResultSet(
             connection=cursor.connection,
@@ -1031,7 +1038,7 @@ def get_catalogs(
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
-            arrow_schema_bytes=arrow_schema_bytes
+            arrow_schema_bytes=arrow_schema_bytes,
         )
 
     def get_schemas(
@@ -1057,7 +1064,9 @@ def get_schemas(
         )
         resp = self.make_request(self._client.GetSchemas, req)
 
-        execute_response, arrow_schema_bytes = self._handle_execute_response(resp, cursor)
+        execute_response, arrow_schema_bytes = self._handle_execute_response(
+            resp, cursor
+        )
 
         return ThriftResultSet(
             connection=cursor.connection,
@@ -1066,7 +1075,7 @@ def get_schemas(
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
-            arrow_schema_bytes=arrow_schema_bytes
+            arrow_schema_bytes=arrow_schema_bytes,
         )
 
     def get_tables(
@@ -1096,7 +1105,9 @@ def get_tables(
         )
         resp = self.make_request(self._client.GetTables, req)
 
-        execute_response, arrow_schema_bytes = self._handle_execute_response(resp, cursor)
+        execute_response, arrow_schema_bytes = self._handle_execute_response(
+            resp, cursor
+        )
 
         return ThriftResultSet(
             connection=cursor.connection,
@@ -1105,7 +1116,7 @@ def get_tables(
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
-            arrow_schema_bytes=arrow_schema_bytes
+            arrow_schema_bytes=arrow_schema_bytes,
         )
 
     def get_columns(
@@ -1135,7 +1146,9 @@ def get_columns(
         )
         resp = self.make_request(self._client.GetColumns, req)
 
-        execute_response, arrow_schema_bytes = self._handle_execute_response(resp, cursor)
+        execute_response, arrow_schema_bytes = self._handle_execute_response(
+            resp, cursor
+        )
 
         return ThriftResultSet(
             connection=cursor.connection,
@@ -1144,7 +1157,7 @@ def get_columns(
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
-            arrow_schema_bytes=arrow_schema_bytes
+            arrow_schema_bytes=arrow_schema_bytes,
         )
 
     def _handle_execute_response(self, resp, cursor):
@@ -1158,9 +1171,10 @@ def _handle_execute_response(self, resp, cursor):
             resp.directResults and resp.directResults.operationStatus,
         )
 
-        execute_response, arrow_schema_bytes = self._results_message_to_execute_response(
-            resp, final_operation_state
-        )
+        (
+            execute_response,
+            arrow_schema_bytes,
+        ) = self._results_message_to_execute_response(resp, final_operation_state)
         execute_response.command_id = command_id
         return execute_response, arrow_schema_bytes
 

tests/unit/test_client.py

@@ -121,7 +121,9 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
 
                 # Verify initial state
                 self.assertEqual(real_result_set.has_been_closed_server_side, closed)
-                expected_status = CommandState.CLOSED if closed else CommandState.SUCCEEDED
+                expected_status = (
+                    CommandState.CLOSED if closed else CommandState.SUCCEEDED
+                )
                 self.assertEqual(real_result_set.status, expected_status)
 
                 # Mock execute_command to return our real result set

tests/unit/test_result_set_filter.py

@@ -54,19 +54,30 @@ def sea_response_with_tables(self):
                 "data_array": [
                     ["schema1", "table1", False, None, "catalog1", "TABLE", None],
                     ["schema1", "table2", False, None, "catalog1", "VIEW", None],
-                    ["schema1", "table3", False, None, "catalog1", "SYSTEM TABLE", None],
+                    [
+                        "schema1",
+                        "table3",
+                        False,
+                        None,
+                        "catalog1",
+                        "SYSTEM TABLE",
+                        None,
+                    ],
                     ["schema1", "table4", False, None, "catalog1", "EXTERNAL", None],
                 ]
             },
         }
 
     @pytest.fixture
-    def sea_result_set_with_tables(self, mock_connection, mock_sea_client, sea_response_with_tables):
+    def sea_result_set_with_tables(
+        self, mock_connection, mock_sea_client, sea_response_with_tables
+    ):
         """Create a SeaResultSet with table data."""
         # Create a deep copy of the response to avoid test interference
         import copy
+
         sea_response_copy = copy.deepcopy(sea_response_with_tables)
-        
+
         return SeaResultSet(
             connection=mock_connection,
             sea_client=mock_sea_client,
@@ -78,11 +89,15 @@ def sea_result_set_with_tables(self, mock_connection, mock_sea_client, sea_respo
     def test_filter_tables_by_type_default(self, sea_result_set_with_tables):
         """Test filtering tables by type with default types."""
         # Default types are TABLE, VIEW, SYSTEM TABLE
-        filtered_result = ResultSetFilter.filter_tables_by_type(sea_result_set_with_tables)
+        filtered_result = ResultSetFilter.filter_tables_by_type(
+            sea_result_set_with_tables
+        )
 
         # Verify that only the default types are included
         assert len(filtered_result._response["result"]["data_array"]) == 3
-        table_types = [row[5] for row in filtered_result._response["result"]["data_array"]]
+        table_types = [
+            row[5] for row in filtered_result._response["result"]["data_array"]
+        ]
         assert "TABLE" in table_types
         assert "VIEW" in table_types
         assert "SYSTEM TABLE" in table_types
@@ -97,7 +112,9 @@ def test_filter_tables_by_type_custom(self, sea_result_set_with_tables):
 
         # Verify that only the specified types are included
         assert len(filtered_result._response["result"]["data_array"]) == 2
-        table_types = [row[5] for row in filtered_result._response["result"]["data_array"]]
+        table_types = [
+            row[5] for row in filtered_result._response["result"]["data_array"]
+        ]
         assert "TABLE" in table_types
         assert "EXTERNAL" in table_types
         assert "VIEW" not in table_types
@@ -112,7 +129,9 @@ def test_filter_tables_by_type_case_insensitive(self, sea_result_set_with_tables
 
         # Verify that the matching types are included despite case differences
         assert len(filtered_result._response["result"]["data_array"]) == 2
-        table_types = [row[5] for row in filtered_result._response["result"]["data_array"]]
+        table_types = [
+            row[5] for row in filtered_result._response["result"]["data_array"]
+        ]
         assert "TABLE" in table_types
         assert "VIEW" in table_types
         assert "SYSTEM TABLE" not in table_types
@@ -126,7 +145,9 @@ def test_filter_tables_by_type_empty_list(self, sea_result_set_with_tables):
 
         # Verify that default types are used
         assert len(filtered_result._response["result"]["data_array"]) == 3
-        table_types = [row[5] for row in filtered_result._response["result"]["data_array"]]
+        table_types = [
+            row[5] for row in filtered_result._response["result"]["data_array"]
+        ]
         assert "TABLE" in table_types
         assert "VIEW" in table_types
         assert "SYSTEM TABLE" in table_types
@@ -144,21 +165,27 @@ def test_filter_by_column_values(self, sea_result_set_with_tables):
 
         # Filter by table name in column index 1
         filtered_result = ResultSetFilter.filter_by_column_values(
-            sea_result_set_with_tables, column_index=1, allowed_values=["table1", "table3"]
+            sea_result_set_with_tables,
+            column_index=1,
+            allowed_values=["table1", "table3"],
         )
 
         # Only rows with table1 or table3 should be included
         assert len(filtered_result._response["result"]["data_array"]) == 2
-        table_names = [row[1] for row in filtered_result._response["result"]["data_array"]]
+        table_names = [
+            row[1] for row in filtered_result._response["result"]["data_array"]
+        ]
         assert "table1" in table_names
         assert "table3" in table_names
         assert "table2" not in table_names
         assert "table4" not in table_names
 
-    def test_filter_by_column_values_case_sensitive(self, mock_connection, mock_sea_client, sea_response_with_tables):
+    def test_filter_by_column_values_case_sensitive(
+        self, mock_connection, mock_sea_client, sea_response_with_tables
+    ):
         """Test case-sensitive filtering by column values."""
         import copy
-        
+
         # Create a fresh result set for the first test
         sea_response_copy1 = copy.deepcopy(sea_response_with_tables)
         result_set1 = SeaResultSet(
@@ -168,18 +195,18 @@ def test_filter_by_column_values_case_sensitive(self, mock_connection, mock_sea_
             buffer_size_bytes=1000,
             arraysize=100,
         )
-        
+
         # First test: Case-sensitive filtering with lowercase values (should find no matches)
         filtered_result = ResultSetFilter.filter_by_column_values(
             result_set1,
             column_index=5,  # tableType column
             allowed_values=["table", "view"],  # lowercase
             case_sensitive=True,
         )
-        
+
         # Verify no matches with lowercase values
         assert len(filtered_result._response["result"]["data_array"]) == 0
-        
+
         # Create a fresh result set for the second test
         sea_response_copy2 = copy.deepcopy(sea_response_with_tables)
         result_set2 = SeaResultSet(
@@ -189,20 +216,22 @@ def test_filter_by_column_values_case_sensitive(self, mock_connection, mock_sea_
             buffer_size_bytes=1000,
             arraysize=100,
         )
-        
+
         # Second test: Case-sensitive filtering with correct case (should find matches)
         filtered_result = ResultSetFilter.filter_by_column_values(
             result_set2,
             column_index=5,  # tableType column
             allowed_values=["TABLE", "VIEW"],  # correct case
             case_sensitive=True,
         )
-        
+
         # Verify matches with correct case
         assert len(filtered_result._response["result"]["data_array"]) == 2
-        
+
         # Extract the table types from the filtered results
-        table_types = [row[5] for row in filtered_result._response["result"]["data_array"]]
+        table_types = [
+            row[5] for row in filtered_result._response["result"]["data_array"]
+        ]
         assert "TABLE" in table_types
         assert "VIEW" in table_types
 
@@ -214,4 +243,4 @@ def test_filter_by_column_values_out_of_bounds(self, sea_result_set_with_tables)
         )
 
         # No rows should match
-        assert len(filtered_result._response["result"]["data_array"]) == 0
+        assert len(filtered_result._response["result"]["data_array"]) == 0

tests/unit/test_sea_backend.py

@@ -54,7 +54,7 @@ def sea_client(self, mock_http_client):
     def sea_session_id(self):
         """Create a SEA session ID."""
         return SessionId.from_sea_session_id("test-session-123")
-        
+
     @pytest.fixture
     def sea_command_id(self):
         """Create a SEA command ID."""
@@ -538,7 +538,7 @@ def test_get_execution_result(
         # Verify basic properties of the result
         assert result.statement_id == "test-statement-123"
         assert result.status == CommandState.SUCCEEDED
-        
+
         # Verify the HTTP request
         mock_http_client._make_request.assert_called_once()
         args, kwargs = mock_http_client._make_request.call_args