|
5 | 5 | import unittest |
6 | 6 | from unittest.mock import MagicMock, patch |
7 | 7 |
|
| 8 | +import pytest |
| 9 | + |
8 | 10 | from databricks.sql.backend.sea.utils.filters import ResultSetFilter |
9 | 11 |
|
10 | 12 |
|
@@ -153,5 +155,121 @@ def test_filter_tables_by_type(self): |
153 | 155 | ) |
154 | 156 |
|
155 | 157 |
|
| 158 | +class TestResultSetFilterWithFetchoneBuffer(unittest.TestCase): |
| 159 | + """Rows buffered by fetchone() on a real SeaResultSet must reach the filter.""" |
| 160 | + |
| 161 | + N_ROWS = 12 |
| 162 | + |
| 163 | + def _row_type(self, i): |
| 164 | + return "TABLE" if i % 2 == 0 else "VIEW" |
| 165 | + |
| 166 | + def _make_result_set(self, queue, result_format): |
| 167 | + from databricks.sql.backend.sea.models.base import ResultData, ResultManifest |
| 168 | + from databricks.sql.backend.sea.result_set import SeaResultSet |
| 169 | + from databricks.sql.backend.types import CommandId, CommandState |
| 170 | + |
| 171 | + connection = MagicMock() |
| 172 | + connection.open = True |
| 173 | + connection.disable_pandas = True |
| 174 | + sea_client = MagicMock() |
| 175 | + sea_client.max_download_threads = 10 |
| 176 | + |
| 177 | + execute_response = MagicMock() |
| 178 | + execute_response.command_id = CommandId.from_sea_statement_id("stmt-1") |
| 179 | + execute_response.status = CommandState.SUCCEEDED |
| 180 | + execute_response.has_been_closed_server_side = False |
| 181 | + execute_response.description = [ |
| 182 | + ("table_name", "string", None, None, None, None, None), |
| 183 | + ("table_type", "string", None, None, None, None, None), |
| 184 | + ] |
| 185 | + execute_response.is_staging_operation = False |
| 186 | + execute_response.lz4_compressed = False |
| 187 | + execute_response.arrow_schema_bytes = None |
| 188 | + |
| 189 | + manifest = ResultManifest( |
| 190 | + format=result_format.value, |
| 191 | + schema={}, |
| 192 | + total_row_count=self.N_ROWS, |
| 193 | + total_byte_count=-1, |
| 194 | + total_chunk_count=-1, |
| 195 | + ) |
| 196 | + with patch( |
| 197 | + "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue", |
| 198 | + return_value=queue, |
| 199 | + ): |
| 200 | + return SeaResultSet( |
| 201 | + connection=connection, |
| 202 | + execute_response=execute_response, |
| 203 | + sea_client=sea_client, |
| 204 | + result_data=ResultData(data=None, external_links=None), |
| 205 | + manifest=manifest, |
| 206 | + buffer_size_bytes=1000, |
| 207 | + arraysize=100, |
| 208 | + ) |
| 209 | + |
| 210 | + def _consume_two_rows(self, rs): |
| 211 | + # 1st fetchone() refills 1 row, 2nd refills 10: rows 2..10 stay buffered |
| 212 | + # and row 11 is still in the queue. |
| 213 | + self.assertEqual(rs.fetchone()[0], "t0") |
| 214 | + self.assertEqual(rs.fetchone()[0], "t1") |
| 215 | + self.assertGreater(rs._row_buffer_count(), 0) |
| 216 | + |
| 217 | + def _expected_names(self): |
| 218 | + return [ |
| 219 | + "t{}".format(i) |
| 220 | + for i in range(2, self.N_ROWS) |
| 221 | + if self._row_type(i) == "TABLE" |
| 222 | + ] |
| 223 | + |
| 224 | + def test_json_filter_includes_rows_buffered_by_fetchone(self): |
| 225 | + from databricks.sql.backend.sea.queue import JsonQueue |
| 226 | + from databricks.sql.backend.sea.utils.constants import ResultFormat |
| 227 | + |
| 228 | + data = [["t{}".format(i), self._row_type(i)] for i in range(self.N_ROWS)] |
| 229 | + rs = self._make_result_set(JsonQueue(data), ResultFormat.JSON_ARRAY) |
| 230 | + self._consume_two_rows(rs) |
| 231 | + |
| 232 | + with patch( |
| 233 | + "databricks.sql.backend.sea.result_set.SeaResultSet" |
| 234 | + ) as mock_sea_result_set_class: |
| 235 | + ResultSetFilter._filter_json_result_set( |
| 236 | + rs, 1, ["TABLE"], case_sensitive=True |
| 237 | + ) |
| 238 | + result_data = mock_sea_result_set_class.call_args.kwargs["result_data"] |
| 239 | + |
| 240 | + self.assertEqual([r[0] for r in result_data.data], self._expected_names()) |
| 241 | + self.assertEqual(rs.rownumber, self.N_ROWS) |
| 242 | + |
| 243 | + def test_arrow_filter_includes_rows_buffered_by_fetchone(self): |
| 244 | + pyarrow = pytest.importorskip("pyarrow") |
| 245 | + from databricks.sql.backend.sea.utils.constants import ResultFormat |
| 246 | + from databricks.sql.utils import ArrowQueue |
| 247 | + |
| 248 | + table = pyarrow.Table.from_pydict( |
| 249 | + { |
| 250 | + "table_name": ["t{}".format(i) for i in range(self.N_ROWS)], |
| 251 | + "table_type": [self._row_type(i) for i in range(self.N_ROWS)], |
| 252 | + } |
| 253 | + ) |
| 254 | + rs = self._make_result_set( |
| 255 | + ArrowQueue(table, self.N_ROWS), ResultFormat.ARROW_STREAM |
| 256 | + ) |
| 257 | + self._consume_two_rows(rs) |
| 258 | + |
| 259 | + with patch( |
| 260 | + "databricks.sql.backend.sea.result_set.SeaResultSet" |
| 261 | + ) as mock_sea_result_set_class: |
| 262 | + ResultSetFilter._filter_arrow_result_set( |
| 263 | + rs, 1, ["TABLE"], case_sensitive=True |
| 264 | + ) |
| 265 | + result_data = mock_sea_result_set_class.call_args.kwargs["result_data"] |
| 266 | + |
| 267 | + filtered = pyarrow.ipc.open_stream(result_data.attachment).read_all() |
| 268 | + self.assertEqual( |
| 269 | + filtered.column("table_name").to_pylist(), self._expected_names() |
| 270 | + ) |
| 271 | + self.assertEqual(rs.rownumber, self.N_ROWS) |
| 272 | + |
| 273 | + |
156 | 274 | if __name__ == "__main__": |
157 | 275 | unittest.main() |
0 commit comments