This is an automated email from the ASF dual-hosted git repository.
derrickaw pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 2b5c1131af3 [yaml] - mongodb write normalization (#38376)
2b5c1131af3 is described below
commit 2b5c1131af3cc0d64a97a6b2d82b082274761635
Author: Derrick Williams <[email protected]>
AuthorDate: Fri May 29 13:43:10 2026 -0400
[yaml] - mongodb write normalization (#38376)
* MongoDB IO write connector for beam yaml
* revert bigtable change
* add yaml and more transformer write support
* more edits to support parity between java and python
* remove read
* remove updateConfiguration
* remove updateField
* update external transforms
* remove unnecessary comments
* add clarifying comments
* fix gemini comments
* address gemini comments
* address reviewer comments
* Updated to 1024 batch size
* remove duplicate function due to rebasing
---------
Co-authored-by: Arnav Arora <[email protected]>
---
sdks/java/io/expansion-service/build.gradle | 1 +
.../apache/beam/sdk/io/mongodb/MongoDbUtils.java | 74 +++++++++
.../MongoDbWriteSchemaTransformConfiguration.java | 85 +++++++++++
.../MongoDbWriteSchemaTransformProvider.java | 168 +++++++++++++++++++++
.../beam/sdk/io/mongodb/MongoDbUtilsTest.java | 135 +++++++++++++++++
.../MongoDbWriteSchemaTransformProviderTest.java | 166 ++++++++++++++++++++
sdks/python/apache_beam/yaml/integration_tests.py | 50 ++++++
sdks/python/apache_beam/yaml/standard_io.yaml | 19 +++
sdks/python/apache_beam/yaml/tests/mongodb.yaml | 44 ++++++
sdks/python/apache_beam/yaml/yaml_io.py | 48 ++++++
sdks/standard_external_transforms.yaml | 32 +++-
11 files changed, 821 insertions(+), 1 deletion(-)
diff --git a/sdks/java/io/expansion-service/build.gradle
b/sdks/java/io/expansion-service/build.gradle
index 1045ad4aeed..32894b97809 100644
--- a/sdks/java/io/expansion-service/build.gradle
+++ b/sdks/java/io/expansion-service/build.gradle
@@ -96,6 +96,7 @@ dependencies {
runtimeOnly project(path: ":sdks:java:io:iceberg:bqms", configuration:
"shadow")
runtimeOnly library.java.bigdataoss_util_hadoop
+ runtimeOnly project(":sdks:java:io:mongodb")
runtimeOnly library.java.kafka_clients
runtimeOnly library.java.slf4j_jdk14
diff --git
a/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbUtils.java
b/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbUtils.java
new file mode 100644
index 00000000000..a5acfb1d19f
--- /dev/null
+++
b/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbUtils.java
@@ -0,0 +1,74 @@
+/*
+ * 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.
+ */
+package org.apache.beam.sdk.io.mongodb;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.schemas.Schema.Field;
+import org.apache.beam.sdk.values.Row;
+import org.bson.BsonNull;
+import org.bson.Document;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Utility methods for MongoDB IO. */
+public class MongoDbUtils {
+
+ /** Converts a Beam {@link Row} to a BSON {@link Document}. */
+ public static Document toDocument(Row row) {
+ Object converted = convertToBsonValue(row);
+ if (converted instanceof Document) {
+ return (Document) converted;
+ }
+ throw new IllegalArgumentException(
+ "Expected Document but got "
+ + (converted != null ? converted.getClass().getName() : "null"));
+ }
+
+ private static @Nullable Object convertToBsonValue(@Nullable Object value) {
+ if (value == null) {
+ return new BsonNull();
+ }
+ if (value instanceof Row) {
+ Row row = (Row) value;
+ Document doc = new Document();
+ for (Field field : row.getSchema().getFields()) {
+ Object fieldValue = row.getValue(field.getName());
+ Object converted = convertToBsonValue(fieldValue);
+ doc.append(field.getName(), converted != null ? converted : new
BsonNull());
+ }
+ return doc;
+ } else if (value instanceof Iterable) {
+ List<Object> bsonList = new ArrayList<>();
+ for (Object item : (Iterable<?>) value) {
+ Object converted = convertToBsonValue(item);
+ bsonList.add(converted != null ? converted : new BsonNull());
+ }
+ return bsonList;
+ } else if (value instanceof Map) {
+ Map<?, ?> map = (Map<?, ?>) value;
+ Document doc = new Document();
+ for (Map.Entry<?, ?> entry : map.entrySet()) {
+ Object converted = convertToBsonValue(entry.getValue());
+ doc.append(String.valueOf(entry.getKey()), converted != null ?
converted : new BsonNull());
+ }
+ return doc;
+ }
+ return value;
+ }
+}
diff --git
a/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbWriteSchemaTransformConfiguration.java
b/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbWriteSchemaTransformConfiguration.java
new file mode 100644
index 00000000000..275cfe88001
--- /dev/null
+++
b/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbWriteSchemaTransformConfiguration.java
@@ -0,0 +1,85 @@
+/*
+ * 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.
+ */
+package org.apache.beam.sdk.io.mongodb;
+
+import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription;
+import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Configuration class for the MongoDB Write transform. */
+@DefaultSchema(AutoValueSchema.class)
+@AutoValue
+public abstract class MongoDbWriteSchemaTransformConfiguration implements
Serializable {
+
+ @SchemaFieldDescription("The connection URI for the MongoDB server.")
+ public abstract String getUri();
+
+ @SchemaFieldDescription("The MongoDB database to write to.")
+ public abstract String getDatabase();
+
+ @SchemaFieldDescription("The MongoDB collection to write to.")
+ public abstract String getCollection();
+
+ @SchemaFieldDescription("The number of documents to include in each batch
write.")
+ @Nullable
+ public abstract Long getBatchSize();
+
+ @SchemaFieldDescription(
+ "This option specifies whether and where to output unwritable rows.
Note: Error handling is currently limited to data conversion failures before
sending to the MongoDB driver, as the underlying MongoDbIO does not yet support
dead-letter queues for write failures.")
+ @Nullable
+ public abstract ErrorHandling getErrorHandling();
+
+ public void validate() {
+ checkArgument(getUri() != null && !getUri().isEmpty(), "MongoDB URI must
be specified.");
+ checkArgument(
+ getDatabase() != null && !getDatabase().isEmpty(), "MongoDB database
must be specified.");
+ checkArgument(
+ getCollection() != null && !getCollection().isEmpty(),
+ "MongoDB collection must be specified.");
+
+ Long batchSize = getBatchSize();
+ if (batchSize != null) {
+ checkArgument(batchSize > 0, "Batch size must be positive.");
+ }
+ }
+
+ public static Builder builder() {
+ return new AutoValue_MongoDbWriteSchemaTransformConfiguration.Builder();
+ }
+
+ @AutoValue.Builder
+ public abstract static class Builder {
+ public abstract Builder setUri(String uri);
+
+ public abstract Builder setDatabase(String database);
+
+ public abstract Builder setCollection(String collection);
+
+ public abstract Builder setBatchSize(Long batchSize);
+
+ public abstract Builder setErrorHandling(ErrorHandling errorHandling);
+
+ public abstract MongoDbWriteSchemaTransformConfiguration build();
+ }
+}
diff --git
a/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbWriteSchemaTransformProvider.java
b/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbWriteSchemaTransformProvider.java
new file mode 100644
index 00000000000..dde265ce2b8
--- /dev/null
+++
b/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbWriteSchemaTransformProvider.java
@@ -0,0 +1,168 @@
+/*
+ * 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.
+ */
+package org.apache.beam.sdk.io.mongodb;
+
+import com.google.auto.service.AutoService;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.io.Serializable;
+import java.util.Collections;
+import java.util.List;
+import org.apache.beam.sdk.coders.AtomicCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
+import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
+import org.apache.beam.sdk.schemas.transforms.TypedSchemaTransformProvider;
+import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionRowTuple;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.bson.Document;
+
+/** An implementation of {@link TypedSchemaTransformProvider} for writing to
MongoDB. */
+@AutoService(SchemaTransformProvider.class)
+public class MongoDbWriteSchemaTransformProvider
+ extends
TypedSchemaTransformProvider<MongoDbWriteSchemaTransformConfiguration> {
+
+ public static class DocumentCoder extends AtomicCoder<Document> implements
Serializable {
+ private static final DocumentCoder INSTANCE = new DocumentCoder();
+
+ private DocumentCoder() {}
+
+ public static DocumentCoder of() {
+ return INSTANCE;
+ }
+
+ @Override
+ public void encode(Document value, OutputStream outStream) throws
java.io.IOException {
+ StringUtf8Coder.of().encode(value.toJson(), outStream);
+ }
+
+ @Override
+ public Document decode(InputStream inStream) throws java.io.IOException {
+ String json = StringUtf8Coder.of().decode(inStream);
+ return Document.parse(json);
+ }
+ }
+
+ private static final String INPUT_TAG = "input";
+ public static final TupleTag<Document> OUTPUT_TAG = new TupleTag<Document>()
{};
+ public static final TupleTag<Row> ERROR_TAG = new TupleTag<Row>() {};
+
+ private static final org.apache.beam.sdk.metrics.Counter errorCounter =
+ org.apache.beam.sdk.metrics.Metrics.counter(
+ MongoDbWriteSchemaTransformProvider.class,
"MongoDB-write-error-counter");
+
+ @Override
+ protected SchemaTransform from(MongoDbWriteSchemaTransformConfiguration
configuration) {
+ return new MongoDbWriteSchemaTransform(configuration);
+ }
+
+ @Override
+ public String identifier() {
+ return "beam:schematransform:org.apache.beam:mongodb_write:v1";
+ }
+
+ @Override
+ public List<String> inputCollectionNames() {
+ return Collections.singletonList(INPUT_TAG);
+ }
+
+ /** The {@link SchemaTransform} that performs the write operation. */
+ private static class MongoDbWriteSchemaTransform extends SchemaTransform {
+ private final MongoDbWriteSchemaTransformConfiguration configuration;
+
+ MongoDbWriteSchemaTransform(MongoDbWriteSchemaTransformConfiguration
configuration) {
+ configuration.validate();
+ this.configuration = configuration;
+ }
+
+ @Override
+ public PCollectionRowTuple expand(PCollectionRowTuple input) {
+ // Retrieve the input PCollection of Rows and its schema.
+ PCollection<Row> rows = input.get(INPUT_TAG);
+ org.apache.beam.sdk.schemas.Schema inputSchema = rows.getSchema();
+
+ // Determine if error handling is enabled and set up the error schema.
+ boolean handleErrors =
ErrorHandling.hasOutput(configuration.getErrorHandling());
+ org.apache.beam.sdk.schemas.Schema errorSchema =
ErrorHandling.errorSchema(inputSchema);
+
+ // Convert Beam Rows to BSON Documents, emitting errors to a separate
tag if enabled.
+ PCollectionTuple outputTuple =
+ rows.apply(
+ "ConvertToDocument",
+ ParDo.of(new RowToBsonDocumentFn(handleErrors, errorSchema))
+ .withOutputTags(OUTPUT_TAG, TupleTagList.of(ERROR_TAG)));
+
+ PCollection<Document> documents =
outputTuple.get(OUTPUT_TAG).setCoder(DocumentCoder.of());
+
+ // Configure the MongoDB write operation.
+ MongoDbIO.Write write =
+ MongoDbIO.write()
+ .withUri(configuration.getUri())
+ .withDatabase(configuration.getDatabase())
+ .withCollection(configuration.getCollection());
+
+ Long batchSize = configuration.getBatchSize();
+ if (batchSize != null) {
+ write = write.withBatchSize(batchSize);
+ }
+
+ // Apply the MongoDB write transform.
+ documents.apply("WriteToMongo", write);
+
+ // Extract and format the error collection.
+ PCollection<Row> errorOutput =
outputTuple.get(ERROR_TAG).setRowSchema(errorSchema);
+
+ // Return the error collection as specified by the configuration.
+ ErrorHandling errorHandling = configuration.getErrorHandling();
+ return PCollectionRowTuple.of(
+ (handleErrors && errorHandling != null) ? errorHandling.getOutput()
: "errors",
+ errorOutput);
+ }
+ }
+
+ /** Converts a Beam {@link Row} to a BSON {@link Document}. */
+ static class RowToBsonDocumentFn extends DoFn<Row, Document> {
+ private final boolean handleErrors;
+ private final org.apache.beam.sdk.schemas.Schema errorSchema;
+
+ RowToBsonDocumentFn(boolean handleErrors,
org.apache.beam.sdk.schemas.Schema errorSchema) {
+ this.handleErrors = handleErrors;
+ this.errorSchema = errorSchema;
+ }
+
+ @ProcessElement
+ public void processElement(@Element Row row, MultiOutputReceiver receiver)
{
+ try {
+ receiver.get(OUTPUT_TAG).output(MongoDbUtils.toDocument(row));
+ } catch (Exception e) {
+ if (!handleErrors) {
+ throw new RuntimeException(e);
+ }
+ errorCounter.inc();
+ receiver.get(ERROR_TAG).output(ErrorHandling.errorRecord(errorSchema,
row, e));
+ }
+ }
+ }
+}
diff --git
a/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbUtilsTest.java
b/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbUtilsTest.java
new file mode 100644
index 00000000000..ec11bb86591
--- /dev/null
+++
b/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbUtilsTest.java
@@ -0,0 +1,135 @@
+/*
+ * 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.
+ */
+package org.apache.beam.sdk.io.mongodb;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.values.Row;
+import org.bson.BsonNull;
+import org.bson.Document;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link MongoDbUtils}. */
+@RunWith(JUnit4.class)
+public class MongoDbUtilsTest {
+
+ @Test
+ public void testToDocumentWithSimplePrimitives() {
+ Schema schema =
+ Schema.builder()
+ .addStringField("stringField")
+ .addInt32Field("intField")
+ .addBooleanField("booleanField")
+ .addDoubleField("doubleField")
+ .build();
+
+ Row row = Row.withSchema(schema).addValues("hello", 42, true,
3.14).build();
+
+ Document doc = MongoDbUtils.toDocument(row);
+
+ assertNotNull(doc);
+ assertEquals("hello", doc.get("stringField"));
+ assertEquals(42, doc.get("intField"));
+ assertEquals(true, doc.get("booleanField"));
+ assertEquals(3.14, doc.get("doubleField"));
+ }
+
+ @Test
+ public void testToDocumentWithNestedRow() {
+ Schema nestedSchema =
+
Schema.builder().addStringField("nestedString").addInt32Field("nestedInt").build();
+
+ Schema parentSchema =
+ Schema.builder()
+ .addStringField("parentString")
+ .addRowField("nestedRow", nestedSchema)
+ .build();
+
+ Row nestedRow = Row.withSchema(nestedSchema).addValues("nestedValue",
100).build();
+ Row parentRow = Row.withSchema(parentSchema).addValues("parentValue",
nestedRow).build();
+
+ Document doc = MongoDbUtils.toDocument(parentRow);
+
+ assertNotNull(doc);
+ assertEquals("parentValue", doc.get("parentString"));
+
+ Object nestedObj = doc.get("nestedRow");
+ assertTrue(nestedObj instanceof Document);
+ Document nestedDoc = (Document) nestedObj;
+ assertEquals("nestedValue", nestedDoc.get("nestedString"));
+ assertEquals(100, nestedDoc.get("nestedInt"));
+ }
+
+ @Test
+ public void testToDocumentWithIterable() {
+ Schema schema = Schema.builder().addArrayField("listField",
FieldType.STRING).build();
+
+ Row row = Row.withSchema(schema).addValue(Arrays.asList("a", "b",
"c")).build();
+
+ Document doc = MongoDbUtils.toDocument(row);
+
+ assertNotNull(doc);
+ Object listObj = doc.get("listField");
+ assertTrue(listObj instanceof List);
+ List<?> list = (List<?>) listObj;
+ assertEquals(3, list.size());
+ assertEquals("a", list.get(0));
+ assertEquals("b", list.get(1));
+ assertEquals("c", list.get(2));
+ }
+
+ @Test
+ public void testToDocumentWithMap() {
+ Schema schema =
+ Schema.builder().addMapField("mapField", FieldType.STRING,
FieldType.INT32).build();
+
+ Map<String, Integer> map = Collections.singletonMap("key", 42);
+ Row row = Row.withSchema(schema).addValue(map).build();
+
+ Document doc = MongoDbUtils.toDocument(row);
+
+ assertNotNull(doc);
+ Object mapObj = doc.get("mapField");
+ assertTrue(mapObj instanceof Document);
+ Document mapDoc = (Document) mapObj;
+ assertEquals(42, mapDoc.get("key"));
+ }
+
+ @Test
+ public void testToDocumentWithNullValues() {
+ Schema schema = Schema.builder().addNullableField("nullableString",
FieldType.STRING).build();
+
+ Row row = Row.withSchema(schema).addValue(null).build();
+
+ Document doc = MongoDbUtils.toDocument(row);
+
+ assertNotNull(doc);
+ Object val = doc.get("nullableString");
+ assertTrue(val instanceof BsonNull);
+ }
+}
diff --git
a/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbWriteSchemaTransformProviderTest.java
b/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbWriteSchemaTransformProviderTest.java
new file mode 100644
index 00000000000..864c516617e
--- /dev/null
+++
b/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbWriteSchemaTransformProviderTest.java
@@ -0,0 +1,166 @@
+/*
+ * 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.
+ */
+package org.apache.beam.sdk.io.mongodb;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThrows;
+
+import java.util.Collections;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.SchemaRegistry;
+import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling;
+import org.apache.beam.sdk.testing.PAssert;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.bson.Document;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link MongoDbWriteSchemaTransformProvider}. */
+@RunWith(JUnit4.class)
+public class MongoDbWriteSchemaTransformProviderTest {
+
+ @Rule public transient TestPipeline p = TestPipeline.create();
+
+ @Test
+ public void testInvalidConfigMissingUri() {
+ assertThrows(
+ IllegalStateException.class,
+ () -> {
+ MongoDbWriteSchemaTransformConfiguration.builder()
+ .setDatabase("db")
+ .setCollection("col")
+ .build()
+ .validate();
+ });
+ }
+
+ @Test
+ public void testInvalidConfigMissingDatabase() {
+ assertThrows(
+ IllegalStateException.class,
+ () -> {
+ MongoDbWriteSchemaTransformConfiguration.builder()
+ .setUri("mongodb://localhost:27017")
+ .setCollection("col")
+ .build()
+ .validate();
+ });
+ }
+
+ @Test
+ public void testInvalidConfigMissingCollection() {
+ assertThrows(
+ IllegalStateException.class,
+ () -> {
+ MongoDbWriteSchemaTransformConfiguration.builder()
+ .setUri("mongodb://localhost:27017")
+ .setDatabase("db")
+ .build()
+ .validate();
+ });
+ }
+
+ @Test
+ public void testInvalidConfigNegativeBatchSize() {
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> {
+ MongoDbWriteSchemaTransformConfiguration.builder()
+ .setUri("mongodb://localhost:27017")
+ .setDatabase("db")
+ .setCollection("col")
+ .setBatchSize(-1L)
+ .build()
+ .validate();
+ });
+ }
+
+ @Test
+ public void testConfigurationSchema() throws Exception {
+ Schema schema =
+
SchemaRegistry.createDefault().getSchema(MongoDbWriteSchemaTransformConfiguration.class);
+
+ // We expect 5 fields now (uri, database, collection, batchSize,
errorHandling)
+ assertEquals(5, schema.getFieldCount());
+ assertNotNull(schema.getField("uri"));
+ assertNotNull(schema.getField("database"));
+ assertNotNull(schema.getField("collection"));
+ assertNotNull(schema.getField("batchSize"));
+ assertNotNull(schema.getField("errorHandling"));
+ }
+
+ @Test
+ public void testRowToBsonDocumentFn() {
+ Schema beamSchema =
+ Schema.builder()
+ .addStringField("name")
+ .addInt32Field("age")
+ .addNullableStringField("country")
+ .build();
+
+ Row row =
+ Row.withSchema(beamSchema)
+ .withFieldValue("name", "John")
+ .withFieldValue("age", 30)
+ .withFieldValue("country", null)
+ .build();
+
+ PCollection<Row> inputRows =
+
p.apply(Create.of(Collections.singletonList(row))).setRowSchema(beamSchema);
+
+ Schema errorSchema = ErrorHandling.errorSchema(beamSchema);
+ PCollectionTuple outputTuple =
+ inputRows.apply(
+ "ConvertToDocument",
+ ParDo.of(
+ new
MongoDbWriteSchemaTransformProvider.RowToBsonDocumentFn(false, errorSchema))
+ .withOutputTags(
+ MongoDbWriteSchemaTransformProvider.OUTPUT_TAG,
+
TupleTagList.of(MongoDbWriteSchemaTransformProvider.ERROR_TAG)));
+
+ PCollection<Document> bsonDocuments =
+ outputTuple
+ .get(MongoDbWriteSchemaTransformProvider.OUTPUT_TAG)
+ .setCoder(MongoDbWriteSchemaTransformProvider.DocumentCoder.of());
+
+
outputTuple.get(MongoDbWriteSchemaTransformProvider.ERROR_TAG).setRowSchema(errorSchema);
+
+ PAssert.that(bsonDocuments)
+ .satisfies(
+ documents -> {
+ Document doc = documents.iterator().next();
+ assertEquals("John", doc.get("name"));
+ assertEquals(30, doc.get("age"));
+ // The RowToBsonDocumentFn retains nulls explicitly in the BSON
document
+ assertEquals(true, doc.containsKey("country"));
+ assertEquals(null, doc.get("country"));
+ return null;
+ });
+
+ p.run().waitUntilFinish();
+ }
+}
diff --git a/sdks/python/apache_beam/yaml/integration_tests.py
b/sdks/python/apache_beam/yaml/integration_tests.py
index 7aad41e0607..2d0b2787fc9 100644
--- a/sdks/python/apache_beam/yaml/integration_tests.py
+++ b/sdks/python/apache_beam/yaml/integration_tests.py
@@ -45,6 +45,7 @@ from testcontainers.core.container import DockerContainer
from testcontainers.core.waiting_utils import wait_for_logs
from testcontainers.google import PubSubContainer
from testcontainers.kafka import KafkaContainer
+from testcontainers.mongodb import MongoDbContainer
from testcontainers.mssql import SqlServerContainer
from testcontainers.mysql import MySqlContainer
from testcontainers.postgres import PostgresContainer
@@ -62,6 +63,8 @@ from apache_beam.yaml.conftest import yaml_test_files_dir
_LOGGER = logging.getLogger(__name__)
+_MONGO_CONTAINER_IMAGE = 'mongo:7.0.7'
+
@contextlib.contextmanager
def gcs_temp_dir(bucket):
@@ -200,6 +203,53 @@ def temp_bigtable_table(project, prefix='yaml_bt_it_'):
_LOGGER.warning("Failed to clean up instance")
[email protected]
+def temp_mongodb_table():
+ """
+ provides a temporary MongoDB instance.
+
+ starts a MongoDB container, creates a unique database
+ and collection name for test isolation, and yields them as a dictionary.
+
+ This allows YAML test files to get connection details without hardcoding
them.
+ Example usage in a YAML test file's fixture section:
+
+ fixtures:
+ - name: mongo_vars
+ type: path.to.this.file.mongodb_fixture
+
+ Then, in the pipeline definition, you can use placeholders like:
+ - uri: ${mongo_vars.URI}
+ - database: ${mongo_vars.DATABASE}
+ - collection: ${mongo_vars.COLLECTION}
+ """
+ _LOGGER.info("Setting up MongoDB fixture...")
+ mongo_container = MongoDbContainer(_MONGO_CONTAINER_IMAGE)
+ try:
+ mongo_container.start()
+ mongo_uri = mongo_container.get_connection_url()
+
+ db_name = f'db_{uuid.uuid4().hex}'
+ collection_name = f'collection_{uuid.uuid4().hex}'
+
+ _LOGGER.info(
+ "MongoDB container started. URI: [%s], DB: [%s], Collection: [%s]",
+ mongo_uri,
+ db_name,
+ collection_name)
+
+ yield {
+ 'URI': mongo_uri,
+ 'DATABASE': db_name,
+ 'COLLECTION': collection_name,
+ }
+
+ finally:
+ _LOGGER.info("Tearing down MongoDB fixture...")
+ mongo_container.stop()
+ _LOGGER.info("MongoDB container stopped.")
+
+
@contextlib.contextmanager
def temp_sqlite_database(prefix='yaml_jdbc_it_'):
"""Context manager to provide a temporary SQLite database via JDBC for
diff --git a/sdks/python/apache_beam/yaml/standard_io.yaml
b/sdks/python/apache_beam/yaml/standard_io.yaml
index a48c3accff4..781d3de193e 100644
--- a/sdks/python/apache_beam/yaml/standard_io.yaml
+++ b/sdks/python/apache_beam/yaml/standard_io.yaml
@@ -115,6 +115,7 @@
'WriteToIceberg': 'apache_beam.yaml.yaml_io.write_to_iceberg'
'ReadFromTFRecord': 'apache_beam.yaml.yaml_io.read_from_tfrecord'
'WriteToTFRecord': 'apache_beam.yaml.yaml_io.write_to_tfrecord'
+ 'WriteToMongoDB': 'apache_beam.yaml.yaml_io.write_to_mongodb'
# General File Formats
@@ -429,3 +430,21 @@
config:
gradle_target: 'sdks:java:io:expansion-service:shadowJar'
+#MongoDB
+- type: renaming
+ transforms:
+ 'WriteToMongoDB': 'WriteToMongoDB'
+ config:
+ mappings:
+ 'WriteToMongoDB':
+ connection_uri: "uri"
+ database: "database"
+ collection: "collection"
+ batch_size: "batch_size"
+ error_handling: "error_handling"
+ underlying_provider:
+ type: beamJar
+ transforms:
+ 'WriteToMongoDB':
'beam:schematransform:org.apache.beam:mongodb_write:v1'
+ config:
+ gradle_target: 'sdks:java:io:expansion-service:shadowJar'
diff --git a/sdks/python/apache_beam/yaml/tests/mongodb.yaml
b/sdks/python/apache_beam/yaml/tests/mongodb.yaml
new file mode 100644
index 00000000000..efce73c8987
--- /dev/null
+++ b/sdks/python/apache_beam/yaml/tests/mongodb.yaml
@@ -0,0 +1,44 @@
+#
+# 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"
+
+pipelines:
+ - pipeline:
+ type: composite
+ transforms:
+ - type: Create
+ name: CreateData
+ config:
+ elements:
+ - { id: 1, name: "John" }
+ - { id: 2, name: "Jane" }
+ - type: WriteToMongoDB
+ name: WriteData
+ input: CreateData
+ config:
+ connection_uri: '{mongo_vars[URI]}'
+ database: '{mongo_vars[DATABASE]}'
+ collection: '{mongo_vars[COLLECTION]}'
+ error_handling:
+ output: my_error_output
+ - type: AssertEqual
+ input: WriteData.my_error_output
+ config:
+ elements: []
diff --git a/sdks/python/apache_beam/yaml/yaml_io.py
b/sdks/python/apache_beam/yaml/yaml_io.py
index f59e730ac81..eaa0d431750 100644
--- a/sdks/python/apache_beam/yaml/yaml_io.py
+++ b/sdks/python/apache_beam/yaml/yaml_io.py
@@ -29,6 +29,7 @@ from collections.abc import Iterable
from collections.abc import Mapping
from typing import Any
from typing import Optional
+from typing import Union
import fastavro
@@ -725,6 +726,53 @@ def write_to_tfrecord(
compression_type=getattr(CompressionTypes, compression_type))
[email protected]_fn
+@yaml_errors.maybe_with_exception_handling_transform_fn
+def write_to_mongodb(
+ pcoll,
+ *,
+ database: str,
+ collection: str,
+ connection_uri: Optional[str] = None,
+ batch_size: int = 1024,
+ extra_client_params: Optional[Mapping[str, Any]] = None):
+ """Writes data to MongoDB.
+
+ Args:
+ pcoll: The input PCollection of Beam Rows.
+ database: The MongoDB database name.
+ collection: The MongoDB collection name.
+ connection_uri: The MongoDB connection string. e.g.
"mongodb://localhost:27017"
+ batch_size: Number of documents per bulk_write to MongoDB.
+ extra_client_params: Optional MongoClient parameters.
+ """
+ from apache_beam.io import mongodbio
+
+ def row_to_dict(value):
+ if value is None:
+ return None
+ if hasattr(value, '_asdict'):
+ return {k: row_to_dict(v) for k, v in value._asdict().items()}
+ elif hasattr(value, 'as_dict'):
+ return {k: row_to_dict(v) for k, v in value.as_dict().items()}
+ elif isinstance(value, (list, tuple)):
+ return [row_to_dict(v) for v in value]
+ elif isinstance(value, Mapping):
+ return {k: row_to_dict(v) for k, v in value.items()}
+ else:
+ return value
+
+ return (
+ pcoll
+ | beam.Map(row_to_dict)
+ | mongodbio.WriteToMongoDB(
+ uri=connection_uri,
+ db=database,
+ coll=collection,
+ batch_size=batch_size,
+ extra_client_params=extra_client_params))
+
+
@beam.ptransform_fn
def match_all(
pcoll,
diff --git a/sdks/standard_external_transforms.yaml
b/sdks/standard_external_transforms.yaml
index 1c536ce319d..057c4e3f47d 100644
--- a/sdks/standard_external_transforms.yaml
+++ b/sdks/standard_external_transforms.yaml
@@ -19,7 +19,7 @@
# configuration in /sdks/standard_expansion_services.yaml.
# Refer to gen_xlang_wrappers.py for more info.
#
-# Last updated on: 2025-06-05
+# Last updated on: 2026-05-06
- default_service: sdks:java:io:expansion-service:shadowJar
description: 'Outputs a PCollection of Beam Rows, each containing a single
INT64
@@ -50,6 +50,36 @@
type: int64
identifier: beam:schematransform:org.apache.beam:generate_sequence:v1
name: GenerateSequence
+- default_service: sdks:java:io:expansion-service:shadowJar
+ description: ''
+ destinations:
+ python: apache_beam/io
+ fields:
+ - description: The number of documents to include in each batch write.
+ name: batch_size
+ nullable: true
+ type: int64
+ - description: The MongoDB collection to write to.
+ name: collection
+ nullable: false
+ type: str
+ - description: The MongoDB database to write to.
+ name: database
+ nullable: false
+ type: str
+ - description: 'This option specifies whether and where to output unwritable
rows.
+ Note: Error handling is currently limited to data conversion failures
before
+ sending to the MongoDB driver, as the underlying MongoDbIO does not yet
support
+ dead-letter queues for write failures.'
+ name: error_handling
+ nullable: true
+ type: Row(output=<class 'str'>)
+ - description: The connection URI for the MongoDB server.
+ name: uri
+ nullable: false
+ type: str
+ identifier: beam:schematransform:org.apache.beam:mongodb_write:v1
+ name: MongodbWrite
- default_service: sdks:java:io:expansion-service:shadowJar
description: ''
destinations:
