From d6d9dfaa9ef43005078ecfaa9f903851ea837a17 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Esteves?= Date: Wed, 15 Nov 2023 15:35:28 +0100 Subject: [PATCH 1/2] Add BigQueryContainer --- google/setup.py | 1 + google/testcontainers/google/__init__.py | 1 + google/testcontainers/google/bigquery.py | 64 +++++++++++++++++++ google/tests/test_bigquery.py | 36 +++++++++++ .../tests/{test_google.py => test_pubsub.py} | 0 5 files changed, 102 insertions(+) create mode 100644 google/testcontainers/google/bigquery.py create mode 100644 google/tests/test_bigquery.py rename google/tests/{test_google.py => test_pubsub.py} (100%) diff --git a/google/setup.py b/google/setup.py index 10a1247d..20f2493a 100644 --- a/google/setup.py +++ b/google/setup.py @@ -12,6 +12,7 @@ url="https://github.com/testcontainers/testcontainers-python", install_requires=[ "testcontainers-core", + "google-cloud-bigquery>=2" "google-cloud-pubsub>=2", ], python_requires=">=3.7", diff --git a/google/testcontainers/google/__init__.py b/google/testcontainers/google/__init__.py index 71665bea..d840276c 100644 --- a/google/testcontainers/google/__init__.py +++ b/google/testcontainers/google/__init__.py @@ -1 +1,2 @@ from .pubsub import PubSubContainer # noqa +from .bigquery import BigQueryContainer # noqa \ No newline at end of file diff --git a/google/testcontainers/google/bigquery.py b/google/testcontainers/google/bigquery.py new file mode 100644 index 00000000..5c9f3477 --- /dev/null +++ b/google/testcontainers/google/bigquery.py @@ -0,0 +1,64 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may +# not use this file except in compliance with the License. You may obtain +# a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations +# under the License. +from testcontainers.core.container import DockerContainer +from google.cloud.bigquery import Client as BigQueryClient +from google.api_core.client_options import ClientOptions +from google.auth.credentials import AnonymousCredentials + + +class BigQueryContainer(DockerContainer): + """ + Docker container to emulate BigQuery, based on https://github.com/testcontainers/testcontainers-java/blob/main/modules/gcloud/src/main/java/org/testcontainers/containers/BigQueryEmulatorContainer.java. + Uses ghcr.io/goccy/bigquery-emulator image by default. + Example: + + The example will spin up a Google Cloud BigQuery emulator that you can use for integration + tests. + + .. doctest:: + + >>> from testcontainers.google import BigQueryContainer + >>> from testcontainers.core.waiting_utils import wait_for_logs + >>> from google.cloud.bigquery import QueryJobConfig + + >>> with BigQueryContainer() as bigquery: + ... wait_for_logs(bigquery, "gRPC server listening", timeout=60) + ... client = bigquery.get_client() + ... result = client.query("SELECT 1", job_config=QueryJobConfig()).result() + ... print(result.total_rows) + """ + def __init__(self, image: str = "ghcr.io/goccy/bigquery-emulator:latest", project: str = "test-project", + http_port: int = 9050, grpc_port: int = 9060, **kwargs) -> None: + super(BigQueryContainer, self).__init__(image=image, **kwargs) + self.project = project + self.http_port = http_port + self.grpc_port = grpc_port + self.with_exposed_ports(http_port, grpc_port) + command = [ + "--project", project, + "--port", str(http_port), + "--grpc-port", str(grpc_port), + ] + self.with_command(' '.join(command)) + + def get_emulator_http_endpoint(self) -> str: + return f"http://{self.get_container_host_ip()}:{self.get_exposed_port(self.http_port)}" + + def get_client(self) -> BigQueryClient: + client_options = ClientOptions(api_endpoint=self.get_emulator_http_endpoint()) + return BigQueryClient( + project=self.project, + client_options=client_options, + credentials=AnonymousCredentials(), + ) + diff --git a/google/tests/test_bigquery.py b/google/tests/test_bigquery.py new file mode 100644 index 00000000..3e646da8 --- /dev/null +++ b/google/tests/test_bigquery.py @@ -0,0 +1,36 @@ +from testcontainers.google import BigQueryContainer +from testcontainers.core.waiting_utils import wait_for_logs +from google.cloud.bigquery import QueryJobConfig, Client as BigQueryClient + + +def test_pubsub_container(): + with BigQueryContainer() as bigquery: + wait_for_logs(bigquery, "gRPC server listening", timeout=60) + + client: BigQueryClient = bigquery.get_client() + + # Function DDL + fn_stmt = ''' + CREATE FUNCTION testr(arr ARRAY>) AS ( + ( + SELECT SUM(IF(elem.name = "foo",elem.val,null)) + FROM UNNEST(arr) AS elem + ) + ) + ''' + + client.query(fn_stmt, job_config=QueryJobConfig()).result() + + select_stmt = ''' + SELECT + testr([ + STRUCT("foo", 10), + STRUCT("bar", 40), + STRUCT("foo", 20) + ]) + ''' + + result = client.query(select_stmt, job_config=QueryJobConfig()).result() + result = [ column for row in result for column in row ] + + assert result == [ 30 ] \ No newline at end of file diff --git a/google/tests/test_google.py b/google/tests/test_pubsub.py similarity index 100% rename from google/tests/test_google.py rename to google/tests/test_pubsub.py From 8ea36ff4cc645620ee23f28e9298f9b87ce3f3f0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Esteves?= Date: Wed, 15 Nov 2023 16:16:38 +0100 Subject: [PATCH 2/2] Add missing comma to setup.py --- google/setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/setup.py b/google/setup.py index 20f2493a..ef3611bc 100644 --- a/google/setup.py +++ b/google/setup.py @@ -12,7 +12,7 @@ url="https://github.com/testcontainers/testcontainers-python", install_requires=[ "testcontainers-core", - "google-cloud-bigquery>=2" + "google-cloud-bigquery>=2", "google-cloud-pubsub>=2", ], python_requires=">=3.7",