|
30 | 30 | from google.cloud.bigquery.enums import TimestampPrecision |
31 | 31 | from google.cloud.bigquery.enums import WriteDisposition |
32 | 32 | from google.cloud.bigquery.job import QueryJobConfig |
33 | | -from google.cloud.bigquery.table import RowIterator, _EmptyRowIterator |
| 33 | +from google.cloud.bigquery.table import RowIterator, TableReference, _EmptyRowIterator |
34 | 34 |
|
35 | 35 | from .helpers import make_client |
36 | 36 |
|
@@ -954,6 +954,7 @@ def test_query_and_wait_arrow_jobs_insert_omits_arrow_fields(): |
954 | 954 | max_results=mock.ANY, |
955 | 955 | query_results_format="ARROW", |
956 | 956 | arrow_serialization_options=_DEFAULT_QUERY_OPTIONS, |
| 957 | + explicit_destination=True, |
957 | 958 | callback=mock.ANY, |
958 | 959 | ) |
959 | 960 | assert wait_or_cancel.call_args.kwargs["page_size"] == 0 |
@@ -1356,6 +1357,127 @@ def test_to_arrow_with_arrow_format_reads_with_the_shared_client(): |
1356 | 1357 | assert to_arrow_iterable.call_args.kwargs["bqstorage_client"] is None |
1357 | 1358 |
|
1358 | 1359 |
|
| 1360 | +@pytest.mark.parametrize( |
| 1361 | + ("job_config", "expected"), |
| 1362 | + [ |
| 1363 | + (QueryJobConfig(destination="p.d.t"), True), |
| 1364 | + # jobs.query doesn't accept a priority either, so this also runs |
| 1365 | + # through jobs.insert, without a destination. |
| 1366 | + (QueryJobConfig(priority="INTERACTIVE"), False), |
| 1367 | + ], |
| 1368 | + ids=["destination", "no_destination"], |
| 1369 | +) |
| 1370 | +def test_query_and_wait_arrow_marks_explicit_destination(job_config, expected): |
| 1371 | + client = mock.MagicMock(spec=Client) |
| 1372 | + |
| 1373 | + with ( |
| 1374 | + mock.patch.object(_job_helpers, "query_jobs_insert") as query_jobs_insert, |
| 1375 | + mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel, |
| 1376 | + ): |
| 1377 | + _query_and_wait(client, job_config=job_config, query_results_format="ARROW") |
| 1378 | + |
| 1379 | + wait_or_cancel.assert_called_once_with( |
| 1380 | + query_jobs_insert.return_value, |
| 1381 | + api_timeout=mock.ANY, |
| 1382 | + wait_timeout=mock.ANY, |
| 1383 | + retry=mock.ANY, |
| 1384 | + page_size=mock.ANY, |
| 1385 | + max_results=mock.ANY, |
| 1386 | + query_results_format="ARROW", |
| 1387 | + arrow_serialization_options=mock.ANY, |
| 1388 | + explicit_destination=expected, |
| 1389 | + callback=mock.ANY, |
| 1390 | + ) |
| 1391 | + |
| 1392 | + |
| 1393 | +def test_wait_or_cancel_sets_explicit_destination(): |
| 1394 | + job = mock.MagicMock(dry_run=False) |
| 1395 | + job.result.return_value = RowIterator( |
| 1396 | + client=mock.MagicMock(), api_request=mock.MagicMock(), path=None, schema=() |
| 1397 | + ) |
| 1398 | + assert not job.result.return_value._explicit_destination |
| 1399 | + |
| 1400 | + rows = _job_helpers._wait_or_cancel( |
| 1401 | + job, |
| 1402 | + api_timeout=None, |
| 1403 | + wait_timeout=None, |
| 1404 | + retry=None, |
| 1405 | + page_size=0, |
| 1406 | + max_results=None, |
| 1407 | + explicit_destination=True, |
| 1408 | + ) |
| 1409 | + |
| 1410 | + assert rows._explicit_destination |
| 1411 | + |
| 1412 | + |
| 1413 | +def _destination_iterator(row_counts, serialized_schema=None): |
| 1414 | + """A RowIterator for a query with an explicit destination table, plus its |
| 1415 | + mock BigQuery Storage client. |
| 1416 | + """ |
| 1417 | + bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") |
| 1418 | + options = bigquery_storage.types.ArrowSerializationOptions |
| 1419 | + if "picos_timestamp_precision" not in options.meta.fields: |
| 1420 | + # query_and_wait() rejects ARROW with this release before it runs. |
| 1421 | + pytest.skip("google-cloud-bigquery-storage has no timestamp precision") |
| 1422 | + schema, batches = _arrow_pages(*row_counts) |
| 1423 | + reader = mock.MagicMock() |
| 1424 | + reader.__iter__.return_value = iter([_stream_response(batch) for batch in batches]) |
| 1425 | + stream = mock.Mock() |
| 1426 | + stream.name = "projects/p/locations/us/sessions/s/streams/0" |
| 1427 | + session = mock.Mock(streams=[stream] if row_counts else []) |
| 1428 | + session.arrow_schema.serialized_schema = ( |
| 1429 | + schema if serialized_schema is None else serialized_schema |
| 1430 | + ) |
| 1431 | + bqstorage_client = mock.MagicMock() |
| 1432 | + bqstorage_client.create_read_session.return_value = session |
| 1433 | + bqstorage_client.read_rows.return_value = reader |
| 1434 | + iterator = RowIterator( |
| 1435 | + client=mock.MagicMock(project="billing-project"), |
| 1436 | + api_request=mock.MagicMock(), |
| 1437 | + path=None, |
| 1438 | + schema=(), |
| 1439 | + table=TableReference.from_string("p.d.t"), |
| 1440 | + project="p", |
| 1441 | + location="US", |
| 1442 | + job_id="j", |
| 1443 | + query_results_format="ARROW", |
| 1444 | + ) |
| 1445 | + iterator._explicit_destination = True |
| 1446 | + return iterator, bqstorage_client |
| 1447 | + |
| 1448 | + |
| 1449 | +def test_download_arrow_explicit_destination_reads_a_read_session(): |
| 1450 | + iterator, bqstorage_client = _destination_iterator([10, 15]) |
| 1451 | + from google.cloud import bigquery_storage |
| 1452 | + |
| 1453 | + batches = list(iterator._download_arrow_from_job_id(bqstorage_client)) |
| 1454 | + |
| 1455 | + assert _values(batches) == list(range(1, 26)) |
| 1456 | + bqstorage_client.create_read_session.assert_called_once_with( |
| 1457 | + parent="projects/billing-project", |
| 1458 | + read_session=bigquery_storage.types.stream.ReadSession( |
| 1459 | + table="projects/p/datasets/d/tables/t", |
| 1460 | + data_format=bigquery_storage.types.stream.DataFormat.ARROW, |
| 1461 | + read_options=bigquery_storage.types.stream.ReadSession.TableReadOptions( |
| 1462 | + arrow_serialization_options=_DEFAULT_READ_OPTIONS |
| 1463 | + ), |
| 1464 | + ), |
| 1465 | + max_stream_count=1, |
| 1466 | + timeout=None, |
| 1467 | + ) |
| 1468 | + # The read session carries the serialization options, so ReadRows doesn't. |
| 1469 | + bqstorage_client.read_rows.assert_called_once_with( |
| 1470 | + "projects/p/locations/us/sessions/s/streams/0", offset=0, timeout=None |
| 1471 | + ) |
| 1472 | + |
| 1473 | + |
| 1474 | +def test_download_arrow_explicit_destination_without_streams_reads_nothing(): |
| 1475 | + iterator, bqstorage_client = _destination_iterator([], serialized_schema=b"") |
| 1476 | + |
| 1477 | + assert list(iterator._download_arrow_from_job_id(bqstorage_client)) == [] |
| 1478 | + bqstorage_client.read_rows.assert_not_called() |
| 1479 | + |
| 1480 | + |
1359 | 1481 | def test_to_arrow_iterable_reads_inline_result_twice(): |
1360 | 1482 | # With JOB_CREATION_OPTIONAL and every row inline there's no job to |
1361 | 1483 | # stream from, so the first page has to stay on the iterator for the |
|
0 commit comments