chore: import upstream snapshot with attribution

python/ray/data/tests/datasource/test_bigquery.py

@@ -0,0 +1,284 @@
+from typing import Iterator
+from unittest import mock
+
+import pandas as pd
+import pyarrow as pa
+import pytest
+from google.api_core import exceptions, operation
+from google.cloud import bigquery, bigquery_storage
+from google.cloud.bigquery import job
+from google.cloud.bigquery_storage_v1.types import stream as gcbqs_stream
+
+import ray
+from ray.data._internal.datasource.bigquery_datasink import BigQueryDatasink
+from ray.data._internal.datasource.bigquery_datasource import BigQueryDatasource
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.planner.plan_write_op import generate_collect_write_stats_fn
+from ray.data.block import Block
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.mock_http_server import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+_TEST_GCP_PROJECT_ID = "mock-test-project-id"
+_TEST_BQ_DATASET_ID = "mockdataset"
+_TEST_BQ_TABLE_ID = "mocktable"
+_TEST_BQ_DATASET = _TEST_BQ_DATASET_ID + "." + _TEST_BQ_TABLE_ID
+_TEST_BQ_TEMP_DESTINATION = _TEST_GCP_PROJECT_ID + ".tempdataset.temptable"
+
+
+@pytest.fixture(autouse=True)
+def bq_client_full_mock(monkeypatch):
+    client_mock = mock.create_autospec(bigquery.Client)
+    client_mock.return_value = client_mock
+
+    def bq_get_dataset_mock(dataset_id):
+        if dataset_id != _TEST_BQ_DATASET_ID:
+            raise exceptions.NotFound(
+                "Dataset {} is not found. Please ensure that it exists.".format(
+                    _TEST_BQ_DATASET
+                )
+            )
+
+    def bq_get_table_mock(table_id):
+        if table_id != _TEST_BQ_DATASET:
+            raise exceptions.NotFound(
+                "Table {} is not found. Please ensure that it exists.".format(
+                    _TEST_BQ_DATASET
+                )
+            )
+
+    def bq_create_dataset_mock(dataset_id, **kwargs):
+        if dataset_id == "existingdataset":
+            raise exceptions.Conflict("Dataset already exists")
+        return mock.Mock(operation.Operation)
+
+    def bq_delete_table_mock(table, **kwargs):
+        return None
+
+    def bq_query_mock(query):
+        fake_job_ref = job._JobReference(
+            "fake_job_id", _TEST_GCP_PROJECT_ID, "us-central1"
+        )
+        fake_query_job = job.QueryJob(fake_job_ref, query, None)
+        fake_query_job.configuration.destination = _TEST_BQ_TEMP_DESTINATION
+        return fake_query_job
+
+    client_mock.get_dataset = bq_get_dataset_mock
+    client_mock.get_table = bq_get_table_mock
+    client_mock.create_dataset = bq_create_dataset_mock
+    client_mock.delete_table = bq_delete_table_mock
+    client_mock.query = bq_query_mock
+
+    monkeypatch.setattr(bigquery, "Client", client_mock)
+    return client_mock
+
+
+@pytest.fixture(autouse=True)
+def bqs_client_full_mock(monkeypatch):
+    client_mock = mock.create_autospec(bigquery_storage.BigQueryReadClient)
+    client_mock.return_value = client_mock
+
+    def bqs_create_read_session(max_stream_count=0, **kwargs):
+        read_session_proto = gcbqs_stream.ReadSession()
+        read_session_proto.streams = [
+            gcbqs_stream.ReadStream() for _ in range(max_stream_count)
+        ]
+        return read_session_proto
+
+    client_mock.create_read_session = bqs_create_read_session
+
+    monkeypatch.setattr(bigquery_storage, "BigQueryReadClient", client_mock)
+    client_mock.reset_mock()
+    return client_mock
+
+
+@pytest.fixture
+def bq_query_result_mock():
+    with mock.patch.object(bigquery.job.QueryJob, "result") as query_result_mock:
+        yield query_result_mock
+
+
+@pytest.fixture
+def bq_query_result_mock_fail():
+    with mock.patch.object(bigquery.job.QueryJob, "result") as query_result_mock_fail:
+        query_result_mock_fail.side_effect = exceptions.BadRequest("400 Syntax error")
+        yield query_result_mock_fail
+
+
+@pytest.fixture
+def ray_get_mock():
+    with mock.patch.object(ray, "get") as ray_get:
+        ray_get.return_value = None
+        yield ray_get
+
+
+class TestReadBigQuery:
+    """Tests for BigQuery Read."""
+
+    @pytest.mark.parametrize(
+        "parallelism",
+        [1, 2, 3, 4, 10, 100],
+    )
+    def test_create_read_tasks(self, parallelism):
+        bq_ds = BigQueryDatasource(
+            project_id=_TEST_GCP_PROJECT_ID,
+            dataset=_TEST_BQ_DATASET,
+        )
+        read_tasks_list = bq_ds.get_read_tasks(parallelism)
+        assert len(read_tasks_list) == parallelism
+
+    @pytest.mark.parametrize(
+        "parallelism",
+        [1, 2, 3, 4, 10, 100],
+    )
+    def test_create_reader_query(self, parallelism, bq_query_result_mock):
+        bq_ds = BigQueryDatasource(
+            project_id=_TEST_GCP_PROJECT_ID,
+            query="SELECT * FROM mockdataset.mocktable",
+        )
+        read_tasks_list = bq_ds.get_read_tasks(parallelism)
+        bq_query_result_mock.assert_called_once()
+        assert len(read_tasks_list) == parallelism
+
+    @pytest.mark.parametrize(
+        "parallelism",
+        [1, 2, 3, 4, 10, 100],
+    )
+    def test_create_reader_query_bad_request(
+        self,
+        parallelism,
+        bq_query_result_mock_fail,
+    ):
+        bq_ds = BigQueryDatasource(
+            project_id=_TEST_GCP_PROJECT_ID,
+            query="SELECT * FROM mockdataset.mocktable",
+        )
+        with pytest.raises(exceptions.BadRequest):
+            bq_ds.get_read_tasks(parallelism)
+        bq_query_result_mock_fail.assert_called()
+
+    def test_dataset_query_kwargs_provided(self):
+        with pytest.raises(ValueError) as exception:
+            BigQueryDatasource(
+                project_id=_TEST_GCP_PROJECT_ID,
+                dataset=_TEST_BQ_DATASET,
+                query="SELECT * FROM mockdataset.mocktable",
+            )
+        expected_message = (
+            "Query and dataset kwargs cannot both be provided"
+            + " (must be mutually exclusive)."
+        )
+        assert str(exception.value) == expected_message
+
+    def test_create_reader_dataset_not_found(self):
+        parallelism = 4
+        bq_ds = BigQueryDatasource(
+            project_id=_TEST_GCP_PROJECT_ID,
+            dataset="nonexistentdataset.mocktable",
+        )
+        with pytest.raises(ValueError) as exception:
+            bq_ds.get_read_tasks(parallelism)
+        expected_message = (
+            "Dataset nonexistentdataset is not found. Please ensure that it exists."
+        )
+        assert str(exception.value) == expected_message
+
+    def test_create_reader_table_not_found(self):
+        parallelism = 4
+        bq_ds = BigQueryDatasource(
+            project_id=_TEST_GCP_PROJECT_ID,
+            dataset="mockdataset.nonexistenttable",
+        )
+        with pytest.raises(ValueError) as exception:
+            bq_ds.get_read_tasks(parallelism)
+        expected_message = (
+            "Table mockdataset.nonexistenttable is not found."
+            + " Please ensure that it exists."
+        )
+        assert str(exception.value) == expected_message
+
+
+class TestWriteBigQuery:
+    """Tests for BigQuery Write."""
+
+    def _extract_write_result(self, stats: Iterator[Block]):
+        return dict(next(stats).iloc[0])
+
+    def test_write(self, ray_get_mock):
+        bq_datasink = BigQueryDatasink(
+            project_id=_TEST_GCP_PROJECT_ID,
+            dataset=_TEST_BQ_DATASET,
+        )
+        arr = pa.array([2, 4, 5, 100])
+        block = pa.Table.from_arrays([arr], names=["data"])
+        ctx = TaskContext(1, "")
+        bq_datasink.write(
+            blocks=[block],
+            ctx=ctx,
+        )
+
+        collect_stats_fn = generate_collect_write_stats_fn()
+        stats = collect_stats_fn([block], ctx)
+        pd.testing.assert_frame_equal(
+            next(stats),
+            pd.DataFrame(
+                {
+                    "num_rows": [4],
+                    "size_bytes": [32],
+                    "write_return": [None],
+                }
+            ),
+        )
+
+    def test_write_dataset_exists(self, ray_get_mock):
+        bq_datasink = BigQueryDatasink(
+            project_id=_TEST_GCP_PROJECT_ID,
+            dataset="existingdataset" + "." + _TEST_BQ_TABLE_ID,
+        )
+        arr = pa.array([2, 4, 5, 100])
+        block = pa.Table.from_arrays([arr], names=["data"])
+        ctx = TaskContext(1, "")
+        bq_datasink.write(
+            blocks=[block],
+            ctx=ctx,
+        )
+        collect_stats_fn = generate_collect_write_stats_fn()
+        stats = collect_stats_fn([block], ctx)
+        pd.testing.assert_frame_equal(
+            next(stats),
+            pd.DataFrame(
+                {
+                    "num_rows": [4],
+                    "size_bytes": [32],
+                    "write_return": [None],
+                }
+            ),
+        )
+
+    def test_write_empty_block(self, ray_get_mock):
+        """Test that writing a zero-sized block doesn't crash.
+
+        See https://github.com/ray-project/ray/issues/51892
+        """
+        bq_datasink = BigQueryDatasink(
+            project_id=_TEST_GCP_PROJECT_ID,
+            dataset=_TEST_BQ_DATASET,
+        )
+        # Create an empty block with schema but no rows
+        block = pa.Table.from_arrays([pa.array([], type=pa.int64())], names=["data"])
+        ctx = TaskContext(1, "")
+        # This should not raise an error - empty blocks should be skipped
+        bq_datasink.write(
+            blocks=[block],
+            ctx=ctx,
+        )
+
+        # write() always calls ray.get(), but with an empty list since the
+        # zero-row block is filtered out (no remote write tasks launched).
+        ray_get_mock.assert_called_once_with([])
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))