From e370a2ac1648bf7c8e5d07a885a277878bafed13 Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Thu, 8 Oct 2026 20:31:36 +0000 Subject: [PATCH] Add mongodb_to_bigquery yaml blueprint example and e2e test --- .../yaml/examples/testing/examples_test.py | 45 +++++++++- .../blueprint/mongodb_to_bigquery.yaml | 70 +++++++++++++++ .../e2e/mongodb_to_bigquery.yaml | 88 +++++++++++++++++++ 3 files changed, 201 insertions(+), 2 deletions(-) create mode 100644 sdks/python/apache_beam/yaml/examples/transforms/blueprint/mongodb_to_bigquery.yaml create mode 100644 sdks/python/apache_beam/yaml/extended_tests/e2e/mongodb_to_bigquery.yaml diff --git a/sdks/python/apache_beam/yaml/examples/testing/examples_test.py b/sdks/python/apache_beam/yaml/examples/testing/examples_test.py index 3ea7e0f07d80..daf72a5f8d59 100644 --- a/sdks/python/apache_beam/yaml/examples/testing/examples_test.py +++ b/sdks/python/apache_beam/yaml/examples/testing/examples_test.py @@ -683,7 +683,8 @@ def _kafka_test_preprocessor( 'test_wordCountImport_yaml', 'test_wordCountInheritance_yaml', 'test_iceberg_to_alloydb_yaml', - 'test_bigquery_write_yaml' + 'test_bigquery_write_yaml', + 'test_mongodb_to_bigquery_yaml' ]) def _io_write_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): @@ -972,6 +973,45 @@ def __mysql_io_read_test_preprocessor( test_spec, lambda url: url.split('/')[3].split('?')[0], 'MySql') +@YamlExamplesTestSuite.register_test_preprocessor([ + 'test_mongodb_to_bigquery_yaml', +]) +def _mongodb_io_read_test_preprocessor( + test_spec: dict, expected: list[str], env: TestEnvironment): + """ + Preprocessor for tests that involve reading from MongoDB. + + This preprocessor replaces any ReadFromMongoDB transform with a Create + transform that reads from a predefined in-memory dictionary. This allows + the test to verify the pipeline's correctness without relying on an active + MongoDB instance. + + Args: + test_spec: The dictionary representation of the YAML pipeline specification. + expected: A list of strings representing the expected output of the + pipeline. + env: The TestEnvironment object providing utilities for creating temporary + files. + + Returns: + The modified test_spec dictionary with ReadFromMongoDB transforms replaced. + """ + if pipeline := test_spec.get('pipeline', None): + for transform in pipeline.get('transforms', []): + if transform.get('type', '') == 'ReadFromMongoDB': + config = transform['config'] + database, collection = config['database'], config['collection'] + transform['type'] = 'Create' + transform['config'] = { + k: v + for k, v in config.items() if k.startswith('__') + } + transform['config']['elements'] = INPUT_TABLES[( + 'MongoDB', str(database), str(collection))] + + return test_spec + + def _db_io_read_test_processor( test_spec: dict, database_url_fn: Callable, database_type: str): """ @@ -1333,7 +1373,8 @@ def _jinja_preprocessor(raw_spec_string: str, test_name: str): ('SqlServer', 'shipment', 'shipments'): input_data.shipments_data(), ('Postgres', 'shipment', 'shipments'): input_data.shipments_data(), ('Oracle', 'shipment', 'shipments'): input_data.shipments_data(), - ('MySql', 'shipment', 'shipments'): input_data.shipments_data() + ('MySql', 'shipment', 'shipments'): input_data.shipments_data(), + ('MongoDB', 'shipment', 'shipments'): input_data.shipments_data() } YAML_DOCS_DIR = os.path.join(os.path.dirname(__file__)) diff --git a/sdks/python/apache_beam/yaml/examples/transforms/blueprint/mongodb_to_bigquery.yaml b/sdks/python/apache_beam/yaml/examples/transforms/blueprint/mongodb_to_bigquery.yaml new file mode 100644 index 000000000000..62a2572d6bf1 --- /dev/null +++ b/sdks/python/apache_beam/yaml/examples/transforms/blueprint/mongodb_to_bigquery.yaml @@ -0,0 +1,70 @@ +# coding=utf-8 +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You 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. + +# This is an example of a Beam YAML pipeline that reads from a MongoDB database +# and writes to BigQuery. This matches the Dataflow Template located +# here - https://cloud.google.com/dataflow/docs/guides/templates/provided/mongodb-to-bigquery + +pipeline: + type: chain + transforms: + # Step 1: Reading data from MongoDB + - type: ReadFromMongoDB + name: ReadFromMongoDB + config: + uri: "mongodb://localhost:27017" + database: "shipment" + collection: "shipments" + schema: | + { + "type": "object", + "properties": { + "shipment_id": {"type": "string"}, + "customer_id": {"type": "string"}, + "shipment_date": {"type": "string"}, + "shipment_cost": {"type": "number"}, + "customer_name": {"type": "string"}, + "customer_email": {"type": "string"} + }, + "required": [ + "shipment_id", + "customer_id", + "shipment_date", + "shipment_cost", + "customer_name", + "customer_email" + ] + } + # Step 2: Write records out to BigQuery + - type: WriteToBigQuery + name: WriteShipments + config: + table: "apache-beam-testing.yaml_test.shipments" + create_disposition: "CREATE_NEVER" + write_disposition: "WRITE_APPEND" + num_streams: 1 + +options: + temp_location: "gs://apache-beam-testing/temp" + +# Expected: +# Row(shipment_id='S1', customer_id='C1', shipment_date='2023-05-01', shipment_cost=150.0, customer_name='Alice', customer_email='alice@example.com') +# Row(shipment_id='S2', customer_id='C2', shipment_date='2023-06-12', shipment_cost=300.0, customer_name='Bob', customer_email='bob@example.com') +# Row(shipment_id='S3', customer_id='C1', shipment_date='2023-05-10', shipment_cost=20.0, customer_name='Alice', customer_email='alice@example.com') +# Row(shipment_id='S4', customer_id='C4', shipment_date='2024-07-01', shipment_cost=150.0, customer_name='Derek', customer_email='derek@example.com') +# Row(shipment_id='S5', customer_id='C5', shipment_date='2023-05-09', shipment_cost=300.0, customer_name='Erin', customer_email='erin@example.com') +# Row(shipment_id='S6', customer_id='C4', shipment_date='2024-07-02', shipment_cost=150.0, customer_name='Derek', customer_email='derek@example.com') diff --git a/sdks/python/apache_beam/yaml/extended_tests/e2e/mongodb_to_bigquery.yaml b/sdks/python/apache_beam/yaml/extended_tests/e2e/mongodb_to_bigquery.yaml new file mode 100644 index 000000000000..c1db35bd5f06 --- /dev/null +++ b/sdks/python/apache_beam/yaml/extended_tests/e2e/mongodb_to_bigquery.yaml @@ -0,0 +1,88 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You 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. +# + +fixtures: + - name: MONGO_VARS + type: "apache_beam.yaml.integration_tests.temp_mongodb_table" + - name: BQ_TABLE + type: "apache_beam.yaml.integration_tests.temp_bigquery_table" + config: + project: "apache-beam-testing" + - name: TEMP_DIR + type: "apache_beam.yaml.integration_tests.gcs_temp_dir" + config: + bucket: "gs://temp-storage-for-end-to-end-tests/temp-it" + +pipelines: + # Pipeline 1: Write records to MongoDB + - pipeline: + type: chain + transforms: + - type: Create + config: + elements: + - {id: 1, name: "Alice", score: 95.5} + - {id: 2, name: "Bob", score: 88.0} + - {id: 3, name: "Charlie", score: 92.3} + - type: WriteToMongoDB + config: + uri: "{MONGO_VARS[URI]}" + database: "{MONGO_VARS[DATABASE]}" + collection: "{MONGO_VARS[COLLECTION]}" + + # Pipeline 2: Read from MongoDB and write to BigQuery + - pipeline: + type: chain + transforms: + - type: ReadFromMongoDB + config: + uri: "{MONGO_VARS[URI]}" + database: "{MONGO_VARS[DATABASE]}" + collection: "{MONGO_VARS[COLLECTION]}" + schema: | + { + "type": "object", + "properties": { + "id": {"type": "integer"}, + "name": {"type": "string"}, + "score": {"type": "number"} + }, + "required": ["id", "name", "score"] + } + - type: WriteToBigQuery + config: + table: "{BQ_TABLE}" + options: + project: "apache-beam-testing" + temp_location: "{TEMP_DIR}" + + # Pipeline 3: Read from BigQuery and verify records + - pipeline: + type: chain + transforms: + - type: ReadFromBigQuery + config: + table: "{BQ_TABLE}" + - type: AssertEqual + config: + elements: + - {id: 1, name: "Alice", score: 95.5} + - {id: 2, name: "Bob", score: 88.0} + - {id: 3, name: "Charlie", score: 92.3} + options: + project: "apache-beam-testing" + temp_location: "{TEMP_DIR}"