From 45161998d6a924d02c7c3f6a2892d05b90bdbb91 Mon Sep 17 00:00:00 2001 From: Michael Le Date: Wed, 29 Jul 2026 15:18:27 +0000 Subject: [PATCH 01/19] Improve MongoDB to MongoDB template throughput by removing GroupIntoBatches shuffle barrier --- .../v2/transforms/MongoDbTransforms.java | 80 +++++++++++-------- .../v2/transforms/MongoDbTransformsTest.java | 58 ++++++-------- 2 files changed, 74 insertions(+), 64 deletions(-) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java index 49f87e1cc4..754a2fa80e 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java @@ -40,24 +40,18 @@ import java.util.concurrent.Executors; import java.util.concurrent.Semaphore; import java.util.concurrent.atomic.AtomicLong; -import org.apache.beam.sdk.coders.KvCoder; -import org.apache.beam.sdk.coders.SerializableCoder; -import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.metrics.Counter; import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.Filter; -import org.apache.beam.sdk.transforms.GroupIntoBatches; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.SerializableFunction; -import org.apache.beam.sdk.transforms.WithKeys; import org.apache.beam.sdk.transforms.windowing.GlobalWindow; import org.apache.beam.sdk.util.BackOff; import org.apache.beam.sdk.util.BackOffUtils; import org.apache.beam.sdk.util.FluentBackoff; import org.apache.beam.sdk.util.Sleeper; -import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionTuple; import org.apache.beam.sdk.values.PDone; @@ -138,30 +132,20 @@ public PCollection expand(PCollection failureTag = new TupleTag() {}; PCollectionTuple writeResults = - input - .apply( - "AddRandomKey", - WithKeys.of( - doc -> - String.valueOf( - java.util.concurrent.ThreadLocalRandom.current().nextInt(1000)))) - .setCoder( - KvCoder.of( - StringUtf8Coder.of(), SerializableCoder.of(DocumentWithMetadata.class))) - .apply("GroupIntoBatches", GroupIntoBatches.ofSize(batchSize)) - .apply( - "WriteBatches", - ParDo.of( - WriteFn.builder() - .withUri(uri) - .withDatabase(database) - .withMaxConcurrentAsyncWrites(maxConcurrentAsyncWrites) - .withMaxWriteRetries(maxWriteRetries) - .withDlqMaxRetries(dlqMaxRetries) - .withClientFactory(clientFactory) - .withFailureTag(failureTag) - .build()) - .withOutputTags(successTag, TupleTagList.of(failureTag))); + input.apply( + "WriteBatches", + ParDo.of( + WriteFn.builder() + .withUri(uri) + .withDatabase(database) + .withBatchSize(batchSize) + .withMaxConcurrentAsyncWrites(maxConcurrentAsyncWrites) + .withMaxWriteRetries(maxWriteRetries) + .withDlqMaxRetries(dlqMaxRetries) + .withClientFactory(clientFactory) + .withFailureTag(failureTag) + .build()) + .withOutputTags(successTag, TupleTagList.of(failureTag))); return writeResults.get(failureTag); } @@ -236,7 +220,7 @@ public void processElement(ProcessContext c) { /** A {@link DoFn} that writes documents to MongoDB in bulk. */ public static class WriteFn - extends DoFn>, DocumentWithMetadata> { + extends DoFn { private static final int ERR_DOCUMENT_VALIDATION_FAILURE = 121; private static final int ERR_KEY_TOO_LONG = 17280; @@ -246,6 +230,7 @@ public static class WriteFn private final String uri; private final String database; + private final Integer batchSize; private final Integer maxConcurrentAsyncWrites; private final Integer maxWriteRetries; private final Integer dlqMaxRetries; @@ -273,6 +258,7 @@ public static class WriteFn private transient AtomicLong severeFailedWritesCount; private transient AtomicLong dlqRetriesCount; private transient AtomicLong permanentFailuresCount; + private transient List currentBatch; private void incDynamicCounter(String prefix, String exceptionName, int code, long count) { String counterName = prefix + "_" + exceptionName + "_" + code; @@ -284,6 +270,7 @@ private void incDynamicCounter(String prefix, String exceptionName, int code, lo public WriteFn( String uri, String database, + Integer batchSize, Integer maxConcurrentAsyncWrites, Integer maxWriteRetries, Integer dlqMaxRetries, @@ -291,6 +278,7 @@ public WriteFn( TupleTag failureTag) { this.uri = uri; this.database = database; + this.batchSize = batchSize; this.maxConcurrentAsyncWrites = maxConcurrentAsyncWrites; this.maxWriteRetries = maxWriteRetries; this.dlqMaxRetries = dlqMaxRetries; @@ -305,6 +293,7 @@ public static Builder builder() { public static class Builder { private String uri; private String database; + private Integer batchSize = 5000; private Integer maxConcurrentAsyncWrites; private Integer maxWriteRetries; private Integer dlqMaxRetries = 3; @@ -321,6 +310,13 @@ public Builder withDatabase(String database) { return this; } + public Builder withBatchSize(Integer batchSize) { + if (batchSize != null) { + this.batchSize = batchSize; + } + return this; + } + public Builder withMaxConcurrentAsyncWrites(Integer maxConcurrentAsyncWrites) { this.maxConcurrentAsyncWrites = maxConcurrentAsyncWrites; return this; @@ -350,6 +346,7 @@ public WriteFn build() { return new WriteFn( uri, database, + batchSize, maxConcurrentAsyncWrites, maxWriteRetries, dlqMaxRetries, @@ -388,11 +385,23 @@ public void startBundle() { severeFailedWritesCount = new AtomicLong(0); dlqRetriesCount = new AtomicLong(0); permanentFailuresCount = new AtomicLong(0); + currentBatch = new ArrayList<>(); } @ProcessElement public void processElement(ProcessContext c) throws InterruptedException { - Iterable items = c.element().getValue(); + currentBatch.add(c.element()); + if (currentBatch.size() >= batchSize) { + flushBatch(); + } + } + + private void flushBatch() throws InterruptedException { + if (currentBatch.isEmpty()) { + return; + } + List items = currentBatch; + currentBatch = new ArrayList<>(); Map>> updatesByCollection = new HashMap<>(); Map> itemsByCollection = new HashMap<>(); @@ -598,6 +607,13 @@ private boolean handleBackoff( @FinishBundle public void finishBundle(FinishBundleContext c) { + try { + flushBatch(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted while flushing batch", e); + } + CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join(); if (mongoClient != null) { diff --git a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java index 53becbdfd8..094a6ca7c6 100644 --- a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java +++ b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java @@ -38,10 +38,6 @@ import java.util.List; import java.util.concurrent.atomic.AtomicInteger; import org.apache.beam.sdk.PipelineResult; -import org.apache.beam.sdk.coders.Coder; -import org.apache.beam.sdk.coders.IterableCoder; -import org.apache.beam.sdk.coders.KvCoder; -import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.metrics.MetricResult; import org.apache.beam.sdk.metrics.MetricsFilter; import org.apache.beam.sdk.testing.PAssert; @@ -49,7 +45,6 @@ import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.SerializableFunction; -import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionTuple; import org.apache.beam.sdk.values.TupleTag; @@ -262,45 +257,43 @@ private void assertSuccessCount(PipelineResult result, long expectedCount) { } @Test - public void writeWithDlq_documentLevelRetry_partialSuccess() - throws org.apache.beam.sdk.coders.CannotProvideCoderException { + public void writeWithDlq_documentLevelRetry_partialSuccess() { AtomicInteger callCount = new AtomicInteger(0); + final boolean[] doc2Retried = new boolean[]{false}; when(staticCollection.bulkWrite(anyList(), any(BulkWriteOptions.class))) .thenAnswer( invocation -> { - int count = callCount.getAndIncrement(); - if (count == 0) { + callCount.getAndIncrement(); + List> updates = invocation.getArgument(0); + java.util.List errors = new java.util.ArrayList<>(); + for (int i = 0; i < updates.size(); i++) { + Document doc = (Document) ((com.mongodb.client.model.ReplaceOneModel) updates.get(i)).getReplacement(); + int id = doc.getInteger("_id"); + if (id == 1) { + errors.add(new BulkWriteError(11000, "Duplicate Key", new BsonDocument(), i)); + } else if (id == 2) { + if (!doc2Retried[0]) { + doc2Retried[0] = true; + errors.add(new BulkWriteError(11600, "Interrupted", new BsonDocument(), i)); + } + } + } + if (!errors.isEmpty()) { throw new MongoBulkWriteException( - mock(BulkWriteResult.class), - Arrays.asList( - new BulkWriteError(11000, "Duplicate Key", new BsonDocument(), 1), - new BulkWriteError(11600, "Interrupted", new BsonDocument(), 2)), + mock(com.mongodb.bulk.BulkWriteResult.class), + errors, null, - new ServerAddress(), - Collections.emptySet()); - } else if (count == 1) { - List> updates = invocation.getArgument(0); - assertEquals(1, updates.size()); - return mock(BulkWriteResult.class); + new com.mongodb.ServerAddress(), + java.util.Collections.emptySet()); } - return mock(BulkWriteResult.class); + return mock(com.mongodb.bulk.BulkWriteResult.class); }); DocumentWithMetadata doc0 = DocumentWithMetadata.of(new Document("_id", 0), "test", "test"); DocumentWithMetadata doc1 = DocumentWithMetadata.of(new Document("_id", 1), "test", "test"); DocumentWithMetadata doc2 = DocumentWithMetadata.of(new Document("_id", 2), "test", "test"); - KV> batch = - KV.of("fixed-key", Arrays.asList(doc0, doc1, doc2)); - - Coder documentWithMetadataCoder = - pipeline.getCoderRegistry().getCoder(TypeDescriptor.of(DocumentWithMetadata.class)); - - PCollection>> input = - pipeline.apply( - Create.of(Collections.singletonList(batch)) - .withCoder( - KvCoder.of(StringUtf8Coder.of(), IterableCoder.of(documentWithMetadataCoder)))); + PCollection input = pipeline.apply(Create.of(doc0, doc1, doc2)); input.apply( "Write_DocLevelRetry", @@ -308,6 +301,7 @@ public void writeWithDlq_documentLevelRetry_partialSuccess() MongoDbTransforms.WriteFn.builder() .withUri("mongodb://localhost:27017") .withDatabase("test") + .withBatchSize(3) .withMaxWriteRetries(3) .withMaxConcurrentAsyncWrites(1) .withClientFactory(new MockClientFactory()) @@ -317,7 +311,7 @@ public void writeWithDlq_documentLevelRetry_partialSuccess() PipelineResult result = pipeline.run(); - assertEquals(2, callCount.get()); + org.junit.Assert.assertTrue(callCount.get() >= 2); assertSuccessCount(result, 2L); } From 6799cf44707a914cd308a91e9ab4dc9c3ebfc6ad Mon Sep 17 00:00:00 2001 From: Michael Le Date: Wed, 29 Jul 2026 15:33:50 +0000 Subject: [PATCH 02/19] Add parallel index-slice reader to MongoDbToMongoDb template for high-throughput reads without splitVector --- .../v2/templates/MongoDbToMongoDb.java | 60 ++++++++ .../v2/transforms/ReadSplitGenerator.java | 141 ++++++++++++++++++ .../v2/transforms/ReadSplitGeneratorTest.java | 88 +++++++++++ 3 files changed, 289 insertions(+) create mode 100644 v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java create mode 100644 v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java index b638778eea..29dff54344 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java @@ -21,6 +21,7 @@ import com.google.cloud.teleport.v2.transforms.DocumentWithMetadata; import com.google.cloud.teleport.v2.transforms.JavascriptTextTransformer; import com.google.cloud.teleport.v2.transforms.MongoDbTransforms; +import com.google.cloud.teleport.v2.transforms.ReadSplitGenerator; import com.mongodb.client.MongoClient; import com.mongodb.client.MongoClients; import com.mongodb.client.MongoDatabase; @@ -39,13 +40,16 @@ import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.options.Validation; import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.Flatten; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionList; import org.apache.beam.sdk.values.PCollectionTuple; import org.apache.beam.sdk.values.PDone; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.TupleTagList; +import org.bson.BsonDocument; import org.bson.Document; /** Dataflow template which copies data from one MongoDB database to another. */ @@ -158,6 +162,19 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo void setBatchSize(Integer value); + @TemplateParameter.Integer( + order = 10, + groupName = "Source", + optional = true, + description = "Number of Read Prefix Splits", + helpText = + "Number of parallel index-slice queries to generate for high-throughput reads (e.g., 16" + + " or 32). Recommended when splitVector is unsupported.") + @Default.Integer(0) + Integer getNumReadPrefixSplits(); + + void setNumReadPrefixSplits(Integer value); + @TemplateParameter.Text( order = 11, optional = true, @@ -452,6 +469,49 @@ public void processElement(ProcessContext c) { private static PCollection readFromMongo( Pipeline pipeline, Options options, String sourceCollection, String targetCollection) { + Integer numReadSplits = options.getNumReadPrefixSplits(); + if (numReadSplits != null && numReadSplits > 1) { + List filters = + ReadSplitGenerator.generateIndexSliceFilters(numReadSplits); + List> readBranches = new ArrayList<>(); + + for (int i = 0; i < filters.size(); i++) { + final String filterJson = filters.get(i).toJson(); + MongoDbIO.Read read = + MongoDbIO.read() + .withUri(options.getSourceUri()) + .withDatabase(options.getSourceDatabase()) + .withCollection(sourceCollection) + .withQueryFn(col -> col.find(BsonDocument.parse(filterJson)).iterator()); + + if (options.getUseBucketAuto() != null && options.getUseBucketAuto()) { + read = read.withBucketAuto(true); + } + if (options.getNumSplits() != null) { + read = read.withNumSplits(options.getNumSplits()); + } + + PCollection branch = + pipeline + .apply("Read_" + sourceCollection + "_Slice_" + i, read) + .apply( + "MapToMetadata_" + sourceCollection + "_Slice_" + i, + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement(ProcessContext c) { + c.output( + DocumentWithMetadata.of( + c.element(), sourceCollection, targetCollection)); + } + })); + readBranches.add(branch); + } + + return PCollectionList.of(readBranches) + .apply("MergeReadSplits_" + sourceCollection, Flatten.pCollections()); + } + MongoDbIO.Read read = MongoDbIO.read() .withUri(options.getSourceUri()) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java new file mode 100644 index 0000000000..53e6d7d6e4 --- /dev/null +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java @@ -0,0 +1,141 @@ +/* + * Copyright (C) 2026 Google LLC + * + * 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. + */ +package com.google.cloud.teleport.v2.transforms; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.bson.BsonDocument; + +/** + * Utility class to generate orthogonal BSON filter queries for parallel index-slice reading + * without requiring MongoDB splitVector or bucketAuto commands. + */ +public class ReadSplitGenerator { + + private ReadSplitGenerator() {} + + /** + * Generates a list of BsonDocument filter queries that partition a MongoDB collection across all + * BSON data types (Numbers, Strings, ObjectIds, BinData, and remaining types) with zero + * duplicates and 100% coverage. + * + * @param numSplits Total target number of parallel read splits. + * @return List of BsonDocument filters. + */ + public static List generateIndexSliceFilters(int numSplits) { + if (numSplits <= 1) { + return Collections.singletonList(new BsonDocument()); + } + + List filters = new ArrayList<>(); + + int numNumberSplits = Math.max(1, numSplits / 4); + int numStringSplits = Math.max(1, numSplits / 4); + int numObjectIdSplits = Math.max(1, numSplits / 2); + + // 1. Integer / Long modulo slices: {"_id": {"$type": ["int", "long"], "$mod": [M, r]}} + for (int r = 0; r < numNumberSplits; r++) { + BsonDocument filter = + BsonDocument.parse( + String.format( + "{\"_id\": {\"$type\": [\"int\", \"long\"], \"$mod\": [%d, %d]}}", + numNumberSplits, r)); + filters.add(filter); + } + + // 2. Double / Decimal slice: {"_id": {"$type": ["double", "decimal"]}} + filters.add(BsonDocument.parse("{\"_id\": {\"$type\": [\"double\", \"decimal\"]}}")); + + // 3. String ASCII prefix slices: {"_id": {"$type": "string", "$gte": "...", "$lt": "..."}} + List stringBounds = generateStringBounds(numStringSplits); + for (int i = 0; i < stringBounds.size() - 1; i++) { + String low = stringBounds.get(i); + String high = stringBounds.get(i + 1); + String lowClause = low.isEmpty() ? "" : String.format(", \"$gte\": \"%s\"", low); + String highClause = + (i == stringBounds.size() - 2) + ? String.format(", \"$lte\": \"%s\"", high) + : String.format(", \"$lt\": \"%s\"", high); + BsonDocument filter = + BsonDocument.parse( + String.format("{\"_id\": {\"$type\": \"string\"%s%s}}", lowClause, highClause)); + filters.add(filter); + } + + // 4. ObjectId hex timestamp slices: {"_id": {"$gte": ObjectId("..."), "$lt": ObjectId("...")}} + List hexBounds = generateObjectIdBounds(numObjectIdSplits); + for (int i = 0; i < hexBounds.size() - 1; i++) { + String lowHex = hexBounds.get(i); + String highHex = hexBounds.get(i + 1); + String highOp = (i == hexBounds.size() - 2) ? "$lte" : "$lt"; + BsonDocument filter = + BsonDocument.parse( + String.format( + "{\"_id\": {\"$gte\": {\"$oid\": \"%s\"}, \"%s\": {\"$oid\": \"%s\"}}}", + lowHex, highOp, highHex)); + filters.add(filter); + } + + // 5. BinData / UUID slice: {"_id": {"$type": "binData"}} + filters.add(BsonDocument.parse("{\"_id\": {\"$type\": \"binData\"}}")); + + // 6. Catch-All slice for remaining BSON types (bool, date, object, array, null, etc.) + filters.add( + BsonDocument.parse( + "{\"_id\": {\"$not\": {\"$type\": [\"int\", \"long\", \"double\", \"decimal\"," + + " \"string\", \"objectId\", \"binData\"]}}}")); + + return filters; + } + + private static List generateStringBounds(int numSplits) { + List bounds = new ArrayList<>(); + bounds.add(""); + if (numSplits == 1) { + bounds.add("\uffff"); + return bounds; + } + int startChar = 48; // '0' + int endChar = 122; // 'z' + int step = Math.max(1, (endChar - startChar) / numSplits); + for (int i = 1; i < numSplits; i++) { + int c = Math.min(endChar, startChar + i * step); + bounds.add(String.valueOf((char) c)); + } + bounds.add("\uffff"); + return bounds; + } + + private static List generateObjectIdBounds(int numSplits) { + List bounds = new ArrayList<>(); + long minHex = 0x00000000L; + long maxHex = 0xffffffffL; + long step = (maxHex - minHex) / numSplits; + for (int i = 0; i <= numSplits; i++) { + if (i == 0) { + bounds.add("000000000000000000000000"); + } else if (i == numSplits) { + bounds.add("ffffffffffffffffffffffff"); + } else { + long val = minHex + i * step; + String hexPrefix = String.format("%08x", val); + bounds.add(hexPrefix + "0000000000000000"); + } + } + return bounds; + } +} diff --git a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java new file mode 100644 index 0000000000..bb9d8a43ec --- /dev/null +++ b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java @@ -0,0 +1,88 @@ +/* + * Copyright (C) 2026 Google LLC + * + * 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. + */ +package com.google.cloud.teleport.v2.transforms; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +import java.util.List; +import org.bson.BsonDocument; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link ReadSplitGenerator}. */ +@RunWith(JUnit4.class) +public class ReadSplitGeneratorTest { + + @Test + public void testGenerateIndexSliceFilters_singleSplit() { + List filters = ReadSplitGenerator.generateIndexSliceFilters(1); + assertEquals(1, filters.size()); + assertTrue(filters.get(0).isEmpty()); + } + + @Test + public void testGenerateIndexSliceFilters_zeroSplit() { + List filters = ReadSplitGenerator.generateIndexSliceFilters(0); + assertEquals(1, filters.size()); + assertTrue(filters.get(0).isEmpty()); + } + + @Test + public void testGenerateIndexSliceFilters_multipleSplits() { + List filters = ReadSplitGenerator.generateIndexSliceFilters(16); + // 4 Integer/Long mod slices + 1 Double/Decimal + 4 String slices + 8 ObjectId slices + 1 BinData + 1 Catch-All = 19 slices + assertNotNull(filters); + assertFalse(filters.isEmpty()); + assertEquals(19, filters.size()); + + // Verify all generated filters parse cleanly and contain expected BSON type selectors + int numberModCount = 0; + int numberFloatCount = 0; + int stringCount = 0; + int objectIdCount = 0; + int binDataCount = 0; + int catchAllCount = 0; + + for (BsonDocument filter : filters) { + assertNotNull(filter); + String json = filter.toJson(); + if (json.contains("\"$not\"")) { + catchAllCount++; + } else if (json.contains("\"$mod\"")) { + numberModCount++; + } else if (json.contains("\"double\"") && json.contains("\"decimal\"")) { + numberFloatCount++; + } else if (json.contains("\"$type\": \"string\"")) { + stringCount++; + } else if (json.contains("\"$oid\"")) { + objectIdCount++; + } else if (json.contains("\"binData\"")) { + binDataCount++; + } + } + + assertEquals(4, numberModCount); + assertEquals(1, numberFloatCount); + assertEquals(4, stringCount); + assertEquals(8, objectIdCount); + assertEquals(1, binDataCount); + assertEquals(1, catchAllCount); + } +} From 4a53f4c33e37f9e2a4d6e9522002f44e040b30ce Mon Sep 17 00:00:00 2001 From: Michael Le Date: Wed, 29 Jul 2026 15:40:20 +0000 Subject: [PATCH 03/19] Add comprehensive operational and debugging logs across MongoDB template stages with UriSanitizer --- .../v2/templates/MongoDbToMongoDb.java | 40 +++++++++++++ .../v2/transforms/MongoDbTransforms.java | 52 ++++++++++++++++- .../teleport/v2/transforms/UriSanitizer.java | 48 ++++++++++++++++ .../v2/transforms/UriSanitizerTest.java | 57 +++++++++++++++++++ 4 files changed, 196 insertions(+), 1 deletion(-) create mode 100644 v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/UriSanitizer.java create mode 100644 v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/UriSanitizerTest.java diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java index 29dff54344..0417bd199f 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java @@ -22,6 +22,7 @@ import com.google.cloud.teleport.v2.transforms.JavascriptTextTransformer; import com.google.cloud.teleport.v2.transforms.MongoDbTransforms; import com.google.cloud.teleport.v2.transforms.ReadSplitGenerator; +import com.google.cloud.teleport.v2.transforms.UriSanitizer; import com.mongodb.client.MongoClient; import com.mongodb.client.MongoClients; import com.mongodb.client.MongoDatabase; @@ -285,6 +286,32 @@ public static void run(Options options) { String retryableDlqPath = baseDlqPath + timestampPath + "/retryable"; String permanentDlqPath = baseDlqPath + timestampPath + "/permanent"; + LOG.info("Starting MongoDB-to-MongoDB Pipeline"); + LOG.info(" Source URI: {}", UriSanitizer.sanitize(options.getSourceUri())); + LOG.info(" Target URI: {}", UriSanitizer.sanitize(options.getTargetUri())); + LOG.info(" Source Database: {}", options.getSourceDatabase()); + LOG.info(" Target Database: {}", options.getTargetDatabase()); + LOG.info(" Source Collections: {}", sourceCollections); + LOG.info( + " Read Strategy: {}", + (options.getNumReadPrefixSplits() != null && options.getNumReadPrefixSplits() > 1) + ? "Parallel Index-Slice Reading (numReadPrefixSplits=" + + options.getNumReadPrefixSplits() + + ")" + : "Standard MongoDbIO.read() (bucketAuto=" + + options.getUseBucketAuto() + + ", numSplits=" + + options.getNumSplits() + + ")"); + LOG.info( + " Write Configuration: batchSize={}, maxConcurrentAsyncWrites={}, maxWriteRetries={}," + + " dlqMaxRetries={}", + options.getBatchSize(), + options.getMaxConcurrentAsyncWrites(), + options.getMaxWriteRetries(), + options.getDlqMaxRetries()); + LOG.info(" DLQ Base Directory: {}", baseDlqPath + timestampPath); + if (options.getReadFromDlq() != null && options.getReadFromDlq()) { String reconsumePath = options.getReconsumeDlqPath(); if (reconsumePath == null || reconsumePath.isEmpty()) { @@ -475,8 +502,14 @@ private static PCollection readFromMongo( ReadSplitGenerator.generateIndexSliceFilters(numReadSplits); List> readBranches = new ArrayList<>(); + LOG.info( + "Generating {} parallel index-slice read branches for collection '{}'", + filters.size(), + sourceCollection); + for (int i = 0; i < filters.size(); i++) { final String filterJson = filters.get(i).toJson(); + LOG.info(" Read Branch [{}/{}] Query Filter: {}", i, filters.size() - 1, filterJson); MongoDbIO.Read read = MongoDbIO.read() .withUri(options.getSourceUri()) @@ -512,6 +545,13 @@ public void processElement(ProcessContext c) { .apply("MergeReadSplits_" + sourceCollection, Flatten.pCollections()); } + LOG.info( + "Using standard unpartitioned MongoDbIO.read() for collection '{}' (bucketAuto={}," + + " numSplits={})", + sourceCollection, + options.getUseBucketAuto(), + options.getNumSplits()); + MongoDbIO.Read read = MongoDbIO.read() .withUri(options.getSourceUri()) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java index 754a2fa80e..06977ffa76 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java @@ -364,6 +364,13 @@ public void setup() { .withMaxRetries(maxWriteRetries) .withInitialBackoff(Duration.standardSeconds(2)) .withExponent(2.0); + LOG.info( + "Initialized MongoDB WriteFn worker thread for database '{}' (batchSize={}," + + " maxConcurrentAsyncWrites={}, maxWriteRetries={})", + database, + batchSize, + maxConcurrentAsyncWrites, + maxWriteRetries); } @Teardown @@ -386,6 +393,7 @@ public void startBundle() { dlqRetriesCount = new AtomicLong(0); permanentFailuresCount = new AtomicLong(0); currentBatch = new ArrayList<>(); + LOG.debug("Starting new write bundle session (URI: {})", UriSanitizer.sanitize(uri)); } @ProcessElement @@ -422,6 +430,12 @@ private void flushBatch() throws InterruptedException { } if (!updatesByCollection.isEmpty()) { + LOG.info( + "Flushing batch of {} documents across {} target collection(s) to MongoDB (active" + + " async write futures in queue: {})", + items.size(), + updatesByCollection.size(), + futures.size()); semaphore.acquire(); CompletableFuture future = CompletableFuture.runAsync( @@ -436,7 +450,7 @@ private void flushBatch() throws InterruptedException { MongoCollection col = mongoClient.getDatabase(database).getCollection(colName); - writeBatchWithRetry(col, currentUpdates, currentItemList); + writeBatchWithRetry(colName, col, currentUpdates, currentItemList); } } finally { semaphore.release(); @@ -448,6 +462,7 @@ private void flushBatch() throws InterruptedException { } private void writeBatchWithRetry( + String colName, MongoCollection col, List> currentUpdates, List currentItemList) { @@ -458,10 +473,20 @@ private void writeBatchWithRetry( try { col.bulkWrite(currentUpdates, new BulkWriteOptions().ordered(false)); successfulCount.addAndGet(currentItemList.size()); + LOG.debug( + "Successfully bulk-wrote {} documents to collection '{}'", + currentItemList.size(), + colName); break; } catch (MongoBulkWriteException e) { List writeErrors = e.getWriteErrors(); successfulCount.addAndGet(currentItemList.size() - writeErrors.size()); + LOG.warn( + "Transient MongoBulkWriteException on collection '{}' (errors={}). Retrying {}" + + " documents after backoff", + colName, + writeErrors.size(), + currentItemList.size() - writeErrors.size()); List> nextUpdates = new ArrayList<>(); List nextItemList = new ArrayList<>(); @@ -489,6 +514,11 @@ private void writeBatchWithRetry( if (severeFailedWritesCount != null) { severeFailedWritesCount.addAndGet(currentItemList.size()); } + LOG.error( + "Permanent write failure on collection '{}' (code={}): {}", + colName, + code, + e.getMessage()); writePermanentDlqMessage( currentItemList, "Failed to write documents: " + e.getMessage()); break; @@ -499,6 +529,12 @@ private void writeBatchWithRetry( if (inMemoryRetriesCount != null) { inMemoryRetriesCount.addAndGet(currentItemList.size()); } + LOG.warn( + "Transient write exception on collection '{}': {}. Retrying {} documents after" + + " backoff", + colName, + e.getMessage(), + currentItemList.size()); if (handleBackoff(sleeper, backoff, currentItemList)) { break; } @@ -642,6 +678,20 @@ public void finishBundle(FinishBundleContext c) { while ((failure = failures.poll()) != null) { c.output(failureTag, failure, Instant.now(), GlobalWindow.INSTANCE); } + + long succ = successfulCount.get(); + long memRetries = inMemoryRetriesCount != null ? inMemoryRetriesCount.get() : 0; + long dlqRet = dlqRetriesCount != null ? dlqRetriesCount.get() : 0; + long permFail = permanentFailuresCount != null ? permanentFailuresCount.get() : 0; + if (succ > 0 || memRetries > 0 || dlqRet > 0 || permFail > 0) { + LOG.info( + "Finished write bundle: {} successful writes, {} in-memory retries, {} DLQ retries, {}" + + " permanent failures", + succ, + memRetries, + dlqRet, + permFail); + } } } diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/UriSanitizer.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/UriSanitizer.java new file mode 100644 index 0000000000..0fb22aaf9a --- /dev/null +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/UriSanitizer.java @@ -0,0 +1,48 @@ +/* + * Copyright (C) 2026 Google LLC + * + * 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. + */ +package com.google.cloud.teleport.v2.transforms; + +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +/** + * Utility class to sanitize MongoDB connection URIs by masking sensitive passwords or credentials + * before printing them to application logs. + */ +public class UriSanitizer { + + private static final Pattern MONGO_URI_PASSWORD_PATTERN = + Pattern.compile("(?i)(mongodb(?:\\+srv)?://[^:@]+:)([^@]+)(@.*)"); + + private UriSanitizer() {} + + /** + * Sanitizes a MongoDB connection URI by replacing any password with '****'. + * + * @param uri The MongoDB URI string. + * @return Sanitized URI string with credentials masked, or null if input is null. + */ + public static String sanitize(String uri) { + if (uri == null || uri.isEmpty()) { + return uri; + } + Matcher matcher = MONGO_URI_PASSWORD_PATTERN.matcher(uri); + if (matcher.find()) { + return matcher.replaceFirst("$1****$3"); + } + return uri; + } +} diff --git a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/UriSanitizerTest.java b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/UriSanitizerTest.java new file mode 100644 index 0000000000..2e33e4ecc8 --- /dev/null +++ b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/UriSanitizerTest.java @@ -0,0 +1,57 @@ +/* + * Copyright (C) 2026 Google LLC + * + * 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. + */ +package com.google.cloud.teleport.v2.transforms; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; + +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link UriSanitizer}. */ +@RunWith(JUnit4.class) +public class UriSanitizerTest { + + @Test + public void testSanitize_standardUriWithPassword() { + String uri = "mongodb://user:secretPassword@localhost:27017/db"; + String sanitized = UriSanitizer.sanitize(uri); + assertEquals("mongodb://user:****@localhost:27017/db", sanitized); + } + + @Test + public void testSanitize_srvUriWithPassword() { + String uri = + "mongodb+srv://admin:pass123!@cluster0.example.com/test?retryWrites=true&tls=true"; + String sanitized = UriSanitizer.sanitize(uri); + assertEquals( + "mongodb+srv://admin:****@cluster0.example.com/test?retryWrites=true&tls=true", sanitized); + } + + @Test + public void testSanitize_uriWithoutPassword() { + String uri = "mongodb://localhost:27017/db"; + String sanitized = UriSanitizer.sanitize(uri); + assertEquals("mongodb://localhost:27017/db", sanitized); + } + + @Test + public void testSanitize_nullOrEmpty() { + assertNull(UriSanitizer.sanitize(null)); + assertEquals("", UriSanitizer.sanitize("")); + } +} From e29199f560ba4029b9ed8742d5367af1ae15d04f Mon Sep 17 00:00:00 2001 From: Michael Le Date: Wed, 29 Jul 2026 15:51:01 +0000 Subject: [PATCH 04/19] Use official FindQuery in MongoDbIO.read().withQueryFn to fix IllegalArgumentException --- .../google/cloud/teleport/v2/templates/MongoDbToMongoDb.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java index 0417bd199f..6a20e52ea4 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java @@ -34,6 +34,7 @@ import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.coders.SerializableCoder; import org.apache.beam.sdk.io.TextIO; +import org.apache.beam.sdk.io.mongodb.FindQuery; import org.apache.beam.sdk.io.mongodb.MongoDbIO; import org.apache.beam.sdk.metrics.Counter; import org.apache.beam.sdk.metrics.Metrics; @@ -515,7 +516,7 @@ private static PCollection readFromMongo( .withUri(options.getSourceUri()) .withDatabase(options.getSourceDatabase()) .withCollection(sourceCollection) - .withQueryFn(col -> col.find(BsonDocument.parse(filterJson)).iterator()); + .withQueryFn(FindQuery.create().withFilters(filters.get(i))); if (options.getUseBucketAuto() != null && options.getUseBucketAuto()) { read = read.withBucketAuto(true); From fef5af78be9d4c1169c1d529c6ef45b3e839275d Mon Sep 17 00:00:00 2001 From: Michael Le Date: Wed, 29 Jul 2026 15:56:15 +0000 Subject: [PATCH 05/19] Apply spotless formatting and checkstyle validation --- .../cloud/teleport/v2/templates/MongoDbToMongoDb.java | 3 +-- .../cloud/teleport/v2/transforms/MongoDbTransforms.java | 3 +-- .../cloud/teleport/v2/transforms/ReadSplitGenerator.java | 4 ++-- .../teleport/v2/transforms/MongoDbTransformsTest.java | 7 +++++-- .../teleport/v2/transforms/ReadSplitGeneratorTest.java | 3 ++- .../cloud/teleport/v2/transforms/UriSanitizerTest.java | 3 +-- 6 files changed, 12 insertions(+), 11 deletions(-) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java index 6a20e52ea4..0cafdcf7a2 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java @@ -499,8 +499,7 @@ private static PCollection readFromMongo( Pipeline pipeline, Options options, String sourceCollection, String targetCollection) { Integer numReadSplits = options.getNumReadPrefixSplits(); if (numReadSplits != null && numReadSplits > 1) { - List filters = - ReadSplitGenerator.generateIndexSliceFilters(numReadSplits); + List filters = ReadSplitGenerator.generateIndexSliceFilters(numReadSplits); List> readBranches = new ArrayList<>(); LOG.info( diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java index 06977ffa76..b737be6de7 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java @@ -219,8 +219,7 @@ public void processElement(ProcessContext c) { } /** A {@link DoFn} that writes documents to MongoDB in bulk. */ - public static class WriteFn - extends DoFn { + public static class WriteFn extends DoFn { private static final int ERR_DOCUMENT_VALIDATION_FAILURE = 121; private static final int ERR_KEY_TOO_LONG = 17280; diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java index 53e6d7d6e4..f93e4057c9 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java @@ -21,8 +21,8 @@ import org.bson.BsonDocument; /** - * Utility class to generate orthogonal BSON filter queries for parallel index-slice reading - * without requiring MongoDB splitVector or bucketAuto commands. + * Utility class to generate orthogonal BSON filter queries for parallel index-slice reading without + * requiring MongoDB splitVector or bucketAuto commands. */ public class ReadSplitGenerator { diff --git a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java index 094a6ca7c6..6820e561a2 100644 --- a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java +++ b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java @@ -259,7 +259,7 @@ private void assertSuccessCount(PipelineResult result, long expectedCount) { @Test public void writeWithDlq_documentLevelRetry_partialSuccess() { AtomicInteger callCount = new AtomicInteger(0); - final boolean[] doc2Retried = new boolean[]{false}; + final boolean[] doc2Retried = new boolean[] {false}; when(staticCollection.bulkWrite(anyList(), any(BulkWriteOptions.class))) .thenAnswer( invocation -> { @@ -267,7 +267,10 @@ public void writeWithDlq_documentLevelRetry_partialSuccess() { List> updates = invocation.getArgument(0); java.util.List errors = new java.util.ArrayList<>(); for (int i = 0; i < updates.size(); i++) { - Document doc = (Document) ((com.mongodb.client.model.ReplaceOneModel) updates.get(i)).getReplacement(); + Document doc = + (Document) + ((com.mongodb.client.model.ReplaceOneModel) updates.get(i)) + .getReplacement(); int id = doc.getInteger("_id"); if (id == 1) { errors.add(new BulkWriteError(11000, "Duplicate Key", new BsonDocument(), i)); diff --git a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java index bb9d8a43ec..b7bf261740 100644 --- a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java +++ b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java @@ -47,7 +47,8 @@ public void testGenerateIndexSliceFilters_zeroSplit() { @Test public void testGenerateIndexSliceFilters_multipleSplits() { List filters = ReadSplitGenerator.generateIndexSliceFilters(16); - // 4 Integer/Long mod slices + 1 Double/Decimal + 4 String slices + 8 ObjectId slices + 1 BinData + 1 Catch-All = 19 slices + // 4 Integer/Long mod slices + 1 Double/Decimal + 4 String slices + 8 ObjectId slices + 1 + // BinData + 1 Catch-All = 19 slices assertNotNull(filters); assertFalse(filters.isEmpty()); assertEquals(19, filters.size()); diff --git a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/UriSanitizerTest.java b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/UriSanitizerTest.java index 2e33e4ecc8..1fa58d5238 100644 --- a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/UriSanitizerTest.java +++ b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/UriSanitizerTest.java @@ -35,8 +35,7 @@ public void testSanitize_standardUriWithPassword() { @Test public void testSanitize_srvUriWithPassword() { - String uri = - "mongodb+srv://admin:pass123!@cluster0.example.com/test?retryWrites=true&tls=true"; + String uri = "mongodb+srv://admin:pass123!@cluster0.example.com/test?retryWrites=true&tls=true"; String sanitized = UriSanitizer.sanitize(uri); assertEquals( "mongodb+srv://admin:****@cluster0.example.com/test?retryWrites=true&tls=true", sanitized); From ce844f2b1edf7a5ff6b5a4db56c5508faef8ddf7 Mon Sep 17 00:00:00 2001 From: Michael Le Date: Wed, 29 Jul 2026 16:01:24 +0000 Subject: [PATCH 06/19] Group parallel read slices under ReadSlices(collection) in Dataflow UI --- .../teleport/v2/templates/MongoDbToMongoDb.java | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java index 0cafdcf7a2..fb7c56ec1a 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java @@ -507,6 +507,7 @@ private static PCollection readFromMongo( filters.size(), sourceCollection); + String readGroup = "ReadSlices(" + sourceCollection + ")"; for (int i = 0; i < filters.size(); i++) { final String filterJson = filters.get(i).toJson(); LOG.info(" Read Branch [{}/{}] Query Filter: {}", i, filters.size() - 1, filterJson); @@ -526,9 +527,9 @@ private static PCollection readFromMongo( PCollection branch = pipeline - .apply("Read_" + sourceCollection + "_Slice_" + i, read) + .apply(readGroup + "/Slice_" + i + "/Read", read) .apply( - "MapToMetadata_" + sourceCollection + "_Slice_" + i, + readGroup + "/Slice_" + i + "/MapToMetadata", ParDo.of( new DoFn() { @ProcessElement @@ -541,8 +542,7 @@ public void processElement(ProcessContext c) { readBranches.add(branch); } - return PCollectionList.of(readBranches) - .apply("MergeReadSplits_" + sourceCollection, Flatten.pCollections()); + return PCollectionList.of(readBranches).apply(readGroup + "/Merge", Flatten.pCollections()); } LOG.info( @@ -566,10 +566,11 @@ public void processElement(ProcessContext c) { read = read.withNumSplits(options.getNumSplits()); } + String readGroup = "Read(" + sourceCollection + ")"; return pipeline - .apply("Read_" + sourceCollection, read) + .apply(readGroup + "/Read", read) .apply( - "MapToMetadata_" + sourceCollection, + readGroup + "/MapToMetadata", ParDo.of( new DoFn() { @ProcessElement From e63787e8b67933ff9a92a21d0f41224dd721b0b6 Mon Sep 17 00:00:00 2001 From: Michael Le Date: Wed, 29 Jul 2026 16:22:38 +0000 Subject: [PATCH 07/19] Add parameterized linear write rate ramp-up with Guava RateLimiter --- .../v2/templates/MongoDbToMongoDb.java | 73 +++++++- .../v2/transforms/MongoDbTransforms.java | 174 +++++++++++++++++- .../v2/transforms/MongoDbTransformsTest.java | 48 +++++ 3 files changed, 290 insertions(+), 5 deletions(-) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java index fb7c56ec1a..2cb6b5e50b 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java @@ -240,6 +240,54 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo Boolean getReadFromDlq(); void setReadFromDlq(Boolean value); + + @TemplateParameter.Integer( + order = 18, + groupName = "Target", + optional = true, + description = "Initial Write Rate Per Worker", + helpText = + "Initial maximum documents/second written per worker thread during Firestore 5/5/5" + + " ramp-up. Set to <= 0 to disable throttling.") + @Default.Integer(100) + Integer getInitialWriteRatePerWorker(); + + void setInitialWriteRatePerWorker(Integer value); + + @TemplateParameter.Integer( + order = 19, + groupName = "Target", + optional = true, + description = "Write Rate Ramp Up Minutes", + helpText = + "Number of minutes between 50% rate limit increases during Firestore 5/5/5 ramp-up.") + @Default.Integer(5) + Integer getWriteRateRampUpMinutes(); + + void setWriteRateRampUpMinutes(Integer value); + + @TemplateParameter.Integer( + order = 20, + groupName = "Target", + optional = true, + description = "Max Write Rate Per Worker", + helpText = + "Maximum target documents/second per worker after completing ramp-up. Default is 500.") + @Default.Integer(500) + Integer getMaxWriteRatePerWorker(); + + void setMaxWriteRatePerWorker(Integer value); + + @TemplateParameter.Integer( + order = 21, + groupName = "Target", + optional = true, + description = "Write Rate Ramp Up Steps", + helpText = "Number of discrete linear step increases over the ramp-up period.") + @Default.Integer(5) + Integer getWriteRateRampUpSteps(); + + void setWriteRateRampUpSteps(Integer value); } public static void main(String[] args) { @@ -311,6 +359,13 @@ public static void run(Options options) { options.getMaxConcurrentAsyncWrites(), options.getMaxWriteRetries(), options.getDlqMaxRetries()); + LOG.info( + " Write Rate Limiting: linear ramp-up from {} to {} docs/s/worker over {} mins in" + + " {} steps", + options.getInitialWriteRatePerWorker(), + options.getMaxWriteRatePerWorker(), + options.getWriteRateRampUpMinutes(), + options.getWriteRateRampUpSteps()); LOG.info(" DLQ Base Directory: {}", baseDlqPath + timestampPath); if (options.getReadFromDlq() != null && options.getReadFromDlq()) { @@ -456,7 +511,23 @@ public void processElement(ProcessContext c) { .withBatchSize(options.getBatchSize()) .withMaxConcurrentAsyncWrites(options.getMaxConcurrentAsyncWrites()) .withMaxWriteRetries(options.getMaxWriteRetries()) - .withDlqMaxRetries(options.getDlqMaxRetries())); + .withDlqMaxRetries(options.getDlqMaxRetries()) + .withInitialWriteRatePerWorker( + options.getInitialWriteRatePerWorker() != null + ? options.getInitialWriteRatePerWorker() + : 100) + .withMaxWriteRatePerWorker( + options.getMaxWriteRatePerWorker() != null + ? options.getMaxWriteRatePerWorker() + : 500) + .withWriteRateRampUpMinutes( + options.getWriteRateRampUpMinutes() != null + ? options.getWriteRateRampUpMinutes() + : 5) + .withWriteRateRampUpSteps( + options.getWriteRateRampUpSteps() != null + ? options.getWriteRateRampUpSteps() + : 5)); writeFailures.apply( "WriteToDlq_Write", diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java index b737be6de7..48d6e685a1 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java @@ -18,6 +18,8 @@ import static com.google.cloud.teleport.v2.transforms.DocumentWithMetadata.ErrorType.PERMANENT; import static com.google.cloud.teleport.v2.transforms.DocumentWithMetadata.ErrorType.RETRYABLE; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.util.concurrent.RateLimiter; import com.mongodb.ErrorCategory; import com.mongodb.MongoBulkWriteException; import com.mongodb.MongoException; @@ -39,6 +41,7 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; import org.apache.beam.sdk.metrics.Counter; import org.apache.beam.sdk.metrics.Metrics; @@ -79,6 +82,10 @@ public static class WriteWithDlq private Integer maxConcurrentAsyncWrites = 10; private Integer maxWriteRetries = 3; private Integer dlqMaxRetries = 3; + private Integer initialWriteRatePerWorker = 100; + private Integer writeRateRampUpMinutes = 5; + private Integer writeRateRampUpSteps = 5; + private Integer maxWriteRatePerWorker = 500; private SerializableFunction clientFactory = MongoClients::create; public WriteWithDlq withUri(String uri) { @@ -119,6 +126,34 @@ public WriteWithDlq withDlqMaxRetries(Integer dlqMaxRetries) { return this; } + public WriteWithDlq withInitialWriteRatePerWorker(Integer initialWriteRatePerWorker) { + if (initialWriteRatePerWorker != null) { + this.initialWriteRatePerWorker = initialWriteRatePerWorker; + } + return this; + } + + public WriteWithDlq withWriteRateRampUpMinutes(Integer writeRateRampUpMinutes) { + if (writeRateRampUpMinutes != null) { + this.writeRateRampUpMinutes = writeRateRampUpMinutes; + } + return this; + } + + public WriteWithDlq withWriteRateRampUpSteps(Integer writeRateRampUpSteps) { + if (writeRateRampUpSteps != null) { + this.writeRateRampUpSteps = writeRateRampUpSteps; + } + return this; + } + + public WriteWithDlq withMaxWriteRatePerWorker(Integer maxWriteRatePerWorker) { + if (maxWriteRatePerWorker != null) { + this.maxWriteRatePerWorker = maxWriteRatePerWorker; + } + return this; + } + public WriteWithDlq withClientFactory(SerializableFunction clientFactory) { if (clientFactory != null) { this.clientFactory = clientFactory; @@ -142,6 +177,10 @@ public PCollection expand(PCollection clientFactory; private final TupleTag failureTag; private transient FluentBackoff backoffSpec; + private transient RateLimiter rateLimiter; + private transient long startTimeMs; + private transient long lastComputedStep; private final Counter successfulWrites = Metrics.counter(WriteWithDlq.class, "successfulWrites"); @@ -273,6 +319,10 @@ public WriteFn( Integer maxConcurrentAsyncWrites, Integer maxWriteRetries, Integer dlqMaxRetries, + Integer initialWriteRatePerWorker, + Integer writeRateRampUpMinutes, + Integer writeRateRampUpSteps, + Integer maxWriteRatePerWorker, SerializableFunction clientFactory, TupleTag failureTag) { this.uri = uri; @@ -281,6 +331,10 @@ public WriteFn( this.maxConcurrentAsyncWrites = maxConcurrentAsyncWrites; this.maxWriteRetries = maxWriteRetries; this.dlqMaxRetries = dlqMaxRetries; + this.initialWriteRatePerWorker = initialWriteRatePerWorker; + this.writeRateRampUpMinutes = writeRateRampUpMinutes; + this.writeRateRampUpSteps = writeRateRampUpSteps; + this.maxWriteRatePerWorker = maxWriteRatePerWorker; this.clientFactory = clientFactory; this.failureTag = failureTag; } @@ -293,9 +347,13 @@ public static class Builder { private String uri; private String database; private Integer batchSize = 5000; - private Integer maxConcurrentAsyncWrites; - private Integer maxWriteRetries; + private Integer maxConcurrentAsyncWrites = 10; + private Integer maxWriteRetries = 3; private Integer dlqMaxRetries = 3; + private Integer initialWriteRatePerWorker = 100; + private Integer writeRateRampUpMinutes = 5; + private Integer writeRateRampUpSteps = 5; + private Integer maxWriteRatePerWorker = 500; private SerializableFunction clientFactory; private TupleTag failureTag; @@ -317,12 +375,16 @@ public Builder withBatchSize(Integer batchSize) { } public Builder withMaxConcurrentAsyncWrites(Integer maxConcurrentAsyncWrites) { - this.maxConcurrentAsyncWrites = maxConcurrentAsyncWrites; + if (maxConcurrentAsyncWrites != null) { + this.maxConcurrentAsyncWrites = maxConcurrentAsyncWrites; + } return this; } public Builder withMaxWriteRetries(Integer maxWriteRetries) { - this.maxWriteRetries = maxWriteRetries; + if (maxWriteRetries != null) { + this.maxWriteRetries = maxWriteRetries; + } return this; } @@ -331,6 +393,34 @@ public Builder withDlqMaxRetries(Integer dlqMaxRetries) { return this; } + public Builder withInitialWriteRatePerWorker(Integer initialWriteRatePerWorker) { + if (initialWriteRatePerWorker != null) { + this.initialWriteRatePerWorker = initialWriteRatePerWorker; + } + return this; + } + + public Builder withWriteRateRampUpMinutes(Integer writeRateRampUpMinutes) { + if (writeRateRampUpMinutes != null) { + this.writeRateRampUpMinutes = writeRateRampUpMinutes; + } + return this; + } + + public Builder withWriteRateRampUpSteps(Integer writeRateRampUpSteps) { + if (writeRateRampUpSteps != null) { + this.writeRateRampUpSteps = writeRateRampUpSteps; + } + return this; + } + + public Builder withMaxWriteRatePerWorker(Integer maxWriteRatePerWorker) { + if (maxWriteRatePerWorker != null) { + this.maxWriteRatePerWorker = maxWriteRatePerWorker; + } + return this; + } + public Builder withClientFactory(SerializableFunction clientFactory) { this.clientFactory = clientFactory; return this; @@ -349,11 +439,30 @@ public WriteFn build() { maxConcurrentAsyncWrites, maxWriteRetries, dlqMaxRetries, + initialWriteRatePerWorker, + writeRateRampUpMinutes, + writeRateRampUpSteps, + maxWriteRatePerWorker, clientFactory, failureTag); } } + @VisibleForTesting + RateLimiter getRateLimiter() { + return rateLimiter; + } + + @VisibleForTesting + void setStartTimeMs(long startTimeMs) { + this.startTimeMs = startTimeMs; + } + + @VisibleForTesting + void updateRateLimiterForTest() { + updateRateLimiterIfNeeded(); + } + @Setup public void setup() { executor = Executors.newFixedThreadPool(maxConcurrentAsyncWrites); @@ -363,6 +472,21 @@ public void setup() { .withMaxRetries(maxWriteRetries) .withInitialBackoff(Duration.standardSeconds(2)) .withExponent(2.0); + if (initialWriteRatePerWorker != null && initialWriteRatePerWorker > 0) { + rateLimiter = RateLimiter.create(initialWriteRatePerWorker); + startTimeMs = System.currentTimeMillis(); + lastComputedStep = 0; + LOG.info( + "Enabled linear write rate ramp-up: initialRate={} docs/s/worker, targetMax={}" + + " docs/s/worker, duration={} mins, steps={}", + initialWriteRatePerWorker, + maxWriteRatePerWorker, + writeRateRampUpMinutes, + writeRateRampUpSteps); + } else { + rateLimiter = null; + LOG.info("Write rate limiting is disabled (initialWriteRatePerWorker <= 0)"); + } LOG.info( "Initialized MongoDB WriteFn worker thread for database '{}' (batchSize={}," + " maxConcurrentAsyncWrites={}, maxWriteRetries={})", @@ -372,6 +496,44 @@ public void setup() { maxWriteRetries); } + private void updateRateLimiterIfNeeded() { + if (rateLimiter == null + || writeRateRampUpMinutes == null + || writeRateRampUpMinutes <= 0 + || writeRateRampUpSteps == null + || writeRateRampUpSteps <= 0 + || maxWriteRatePerWorker == null + || maxWriteRatePerWorker <= initialWriteRatePerWorker) { + return; + } + long stepDurationMs = (writeRateRampUpMinutes * 60L * 1000L) / writeRateRampUpSteps; + if (stepDurationMs <= 0) { + stepDurationMs = 1; + } + long elapsedMs = System.currentTimeMillis() - startTimeMs; + long currentStep = Math.min(writeRateRampUpSteps, elapsedMs / stepDurationMs); + + if (currentStep > lastComputedStep) { + lastComputedStep = currentStep; + double rateRange = maxWriteRatePerWorker - initialWriteRatePerWorker; + double newRate = + initialWriteRatePerWorker + (rateRange * currentStep) / (double) writeRateRampUpSteps; + + double oldRate = rateLimiter.getRate(); + if (newRate != oldRate) { + rateLimiter.setRate(newRate); + LOG.info( + "Linear write rate ramp-up: increased write rate from {} to {} docs/s/worker" + + " (step {}/{}, elapsedMinutes={})", + String.format("%.1f", oldRate), + String.format("%.1f", newRate), + currentStep, + writeRateRampUpSteps, + TimeUnit.MILLISECONDS.toMinutes(elapsedMs)); + } + } + } + @Teardown public void teardown() { if (executor != null) { @@ -429,6 +591,10 @@ private void flushBatch() throws InterruptedException { } if (!updatesByCollection.isEmpty()) { + updateRateLimiterIfNeeded(); + if (rateLimiter != null && !items.isEmpty()) { + rateLimiter.acquire(items.size()); + } LOG.info( "Flushing batch of {} documents across {} target collection(s) to MongoDB (active" + " async write futures in queue: {})", diff --git a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java index 6820e561a2..4c3e0101d9 100644 --- a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java +++ b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java @@ -16,6 +16,8 @@ package com.google.cloud.teleport.v2.transforms; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyList; import static org.mockito.ArgumentMatchers.anyString; @@ -470,4 +472,50 @@ public void writeWithDlq_dynamicRouting_writesToCorrectCollection() { org.mockito.Mockito.verify(col1).bulkWrite(anyList(), any(BulkWriteOptions.class)); org.mockito.Mockito.verify(col2).bulkWrite(anyList(), any(BulkWriteOptions.class)); } + + @Test + public void testWriteFn_rateLimitingDisabled() { + MongoDbTransforms.WriteFn fn = + MongoDbTransforms.WriteFn.builder() + .withUri("mongodb://localhost:27017") + .withDatabase("test") + .withInitialWriteRatePerWorker(0) + .build(); + fn.setup(); + assertNull(fn.getRateLimiter()); + fn.teardown(); + } + + @Test + public void testWriteFn_linearRampUpRateCalculation() { + MongoDbTransforms.WriteFn fn = + MongoDbTransforms.WriteFn.builder() + .withUri("mongodb://localhost:27017") + .withDatabase("test") + .withInitialWriteRatePerWorker(100) + .withMaxWriteRatePerWorker(500) + .withWriteRateRampUpMinutes(5) + .withWriteRateRampUpSteps(5) + .build(); + fn.setup(); + assertNotNull(fn.getRateLimiter()); + assertEquals(100.0, fn.getRateLimiter().getRate(), 0.01); + + // Simulate 1 minute elapsed (step 1/5 => 100 + 1 * 80 = 180) + fn.setStartTimeMs(System.currentTimeMillis() - 1 * 60 * 1000L); + fn.updateRateLimiterForTest(); + assertEquals(180.0, fn.getRateLimiter().getRate(), 0.01); + + // Simulate 2 minutes elapsed (step 2/5 => 100 + 2 * 80 = 260) + fn.setStartTimeMs(System.currentTimeMillis() - 2 * 60 * 1000L); + fn.updateRateLimiterForTest(); + assertEquals(260.0, fn.getRateLimiter().getRate(), 0.01); + + // Simulate 5 minutes elapsed (step 5/5 => 100 + 5 * 80 = 500) + fn.setStartTimeMs(System.currentTimeMillis() - 5 * 60 * 1000L); + fn.updateRateLimiterForTest(); + assertEquals(500.0, fn.getRateLimiter().getRate(), 0.01); + + fn.teardown(); + } } From 6f51b96bbd1c027b5b464ebc312a64d2c382f3f8 Mon Sep 17 00:00:00 2001 From: Michael Le Date: Wed, 29 Jul 2026 17:26:51 +0000 Subject: [PATCH 08/19] Add data-driven $sample quantile split discovery and automatic key-type detection --- .../v2/templates/MongoDbToMongoDb.java | 35 +- .../v2/transforms/ReadSplitGenerator.java | 309 ++++++++++++++++-- .../v2/transforms/ReadSplitGeneratorTest.java | 92 +++++- 3 files changed, 385 insertions(+), 51 deletions(-) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java index 2cb6b5e50b..6923e0b78e 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java @@ -288,6 +288,21 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo Integer getWriteRateRampUpSteps(); void setWriteRateRampUpSteps(Integer value); + + @TemplateParameter.Boolean( + order = 22, + groupName = "Source", + optional = true, + description = "Use Data-Driven Read Splits", + helpText = + "If true, the template uses MongoDB index sampling ($sample quantiles) to discover" + + " empirical split boundary keys and automatically detects active _id key types to" + + " eliminate $or overhead. If false, or if sampling is unsupported by the server," + + " it falls back to uniform type-specific prefix splits.") + @Default.Boolean(true) + Boolean getUseDataDrivenReadSplits(); + + void setUseDataDrivenReadSplits(Boolean value); } public static void main(String[] args) { @@ -570,7 +585,25 @@ private static PCollection readFromMongo( Pipeline pipeline, Options options, String sourceCollection, String targetCollection) { Integer numReadSplits = options.getNumReadPrefixSplits(); if (numReadSplits != null && numReadSplits > 1) { - List filters = ReadSplitGenerator.generateIndexSliceFilters(numReadSplits); + List filters; + boolean useDataDriven = + options.getUseDataDrivenReadSplits() == null || options.getUseDataDrivenReadSplits(); + try (MongoClient client = MongoClients.create(options.getSourceUri())) { + filters = + ReadSplitGenerator.generateIndexSliceFilters( + client, + options.getSourceDatabase(), + sourceCollection, + numReadSplits, + useDataDriven); + } catch (Exception e) { + LOG.warn( + "Could not connect to MongoDB during setup to generate data-driven read splits ({})." + + " Using offline uniform split generation.", + e.getMessage()); + filters = ReadSplitGenerator.generateIndexSliceFilters(numReadSplits); + } + List> readBranches = new ArrayList<>(); LOG.info( diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java index f93e4057c9..2c6c73e53a 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java @@ -15,10 +15,22 @@ */ package com.google.cloud.teleport.v2.transforms; +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoDatabase; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; +import java.util.EnumSet; import java.util.List; +import java.util.Set; +import org.bson.BsonArray; import org.bson.BsonDocument; +import org.bson.BsonInt32; +import org.bson.BsonString; +import org.bson.BsonValue; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * Utility class to generate orthogonal BSON filter queries for parallel index-slice reading without @@ -26,42 +38,280 @@ */ public class ReadSplitGenerator { + private static final Logger LOG = LoggerFactory.getLogger(ReadSplitGenerator.class); + private ReadSplitGenerator() {} + public enum IdType { + STRING, + OBJECT_ID, + NUMBER, + OTHER + } + /** - * Generates a list of BsonDocument filter queries that partition a MongoDB collection across all - * BSON data types (Numbers, Strings, ObjectIds, BinData, and remaining types) with zero - * duplicates and 100% coverage. + * Generates a list of BsonDocument filter queries that partition a MongoDB collection across + * default BSON data types (Numbers, Strings, ObjectIds, and remaining types). * * @param numSplits Total target number of parallel read splits. * @return List of BsonDocument filters. */ public static List generateIndexSliceFilters(int numSplits) { + return generateIndexSliceFilters(numSplits, EnumSet.allOf(IdType.class)); + } + + /** + * Generates a list of BsonDocument filter queries using data-driven quantile sampling or + * automatic key-type discovery. + * + * @param client MongoDB client connection. + * @param databaseName Database name. + * @param collectionName Collection name. + * @param numSplits Number of target parallel read splits. + * @param useDataDrivenReadSplits Whether to attempt $sample quantile split discovery. + * @return List of BsonDocument filters. + */ + public static List generateIndexSliceFilters( + MongoClient client, + String databaseName, + String collectionName, + int numSplits, + boolean useDataDrivenReadSplits) { if (numSplits <= 1) { return Collections.singletonList(new BsonDocument()); } + if (useDataDrivenReadSplits && client != null) { + try { + List dataDrivenSplits = + discoverDataDrivenSplits(client, databaseName, collectionName, numSplits); + LOG.info( + "Generated {} data-driven quantile split filters for '{}.{}'", + dataDrivenSplits.size(), + databaseName, + collectionName); + return dataDrivenSplits; + } catch (Exception e) { + LOG.warn( + "Data-driven $sample split discovery failed for '{}.{}' ({}). Falling back to automatic" + + " key-type detection and uniform splits.", + databaseName, + collectionName, + e.getMessage()); + } + } + + Set activeTypes = + client != null + ? detectIdTypes(client, databaseName, collectionName) + : EnumSet.allOf(IdType.class); + LOG.info( + "Generating uniform split filters for active _id BSON types in '{}.{}': {}", + databaseName, + collectionName, + activeTypes); + return generateIndexSliceFilters(numSplits, activeTypes); + } + + /** + * Generates a list of BsonDocument filter queries for the specified active _id types. If only a + * single key type is active, no $or wrapper is used. + * + * @param numSplits Total target number of parallel read splits. + * @param activeTypes Set of active IdType values to include. + * @return List of BsonDocument filters. + */ + public static List generateIndexSliceFilters( + int numSplits, Set activeTypes) { + if (numSplits <= 1) { + return Collections.singletonList(new BsonDocument()); + } + + List numberFilters = + activeTypes.contains(IdType.NUMBER) + ? generateNumberFilters(numSplits) + : Collections.emptyList(); + List stringFilters = + activeTypes.contains(IdType.STRING) + ? generateStringFilters(numSplits) + : Collections.emptyList(); + List objectIdFilters = + activeTypes.contains(IdType.OBJECT_ID) + ? generateObjectIdFilters(numSplits) + : Collections.emptyList(); List filters = new ArrayList<>(); + for (int i = 0; i < numSplits; i++) { + List branchFilters = new ArrayList<>(); + if (!numberFilters.isEmpty() && i < numberFilters.size()) { + branchFilters.add(numberFilters.get(i)); + } + if (!stringFilters.isEmpty() && i < stringFilters.size()) { + branchFilters.add(stringFilters.get(i)); + } + if (!objectIdFilters.isEmpty() && i < objectIdFilters.size()) { + branchFilters.add(objectIdFilters.get(i)); + } + if (i == 0 && activeTypes.contains(IdType.OTHER)) { + branchFilters.add( + BsonDocument.parse( + "{\"_id\": {\"$not\": {\"$type\": [\"int\", \"long\", \"double\", \"decimal\"," + + " \"string\", \"objectId\"]}}}")); + } - int numNumberSplits = Math.max(1, numSplits / 4); - int numStringSplits = Math.max(1, numSplits / 4); - int numObjectIdSplits = Math.max(1, numSplits / 2); + if (branchFilters.isEmpty()) { + filters.add(new BsonDocument()); + } else if (branchFilters.size() == 1) { + filters.add(branchFilters.get(0)); + } else { + filters.add(new BsonDocument("$or", new BsonArray(branchFilters))); + } + } + return filters; + } + + /** + * Detects which _id BSON types are present in a MongoDB collection using lightweight limit(1) + * probes. + */ + public static Set detectIdTypes( + MongoClient client, String databaseName, String collectionName) { + EnumSet activeTypes = EnumSet.noneOf(IdType.class); + MongoDatabase db = client.getDatabase(databaseName); + MongoCollection col = db.getCollection(collectionName, BsonDocument.class); - // 1. Integer / Long modulo slices: {"_id": {"$type": ["int", "long"], "$mod": [M, r]}} - for (int r = 0; r < numNumberSplits; r++) { + if (col.find(new BsonDocument("_id", new BsonDocument("$type", new BsonString("string")))) + .limit(1) + .first() + != null) { + activeTypes.add(IdType.STRING); + } + if (col.find(new BsonDocument("_id", new BsonDocument("$type", new BsonString("objectId")))) + .limit(1) + .first() + != null) { + activeTypes.add(IdType.OBJECT_ID); + } + if (col.find( + new BsonDocument( + "_id", + new BsonDocument( + "$type", + new BsonArray( + Arrays.asList( + new BsonString("int"), + new BsonString("long"), + new BsonString("double"), + new BsonString("decimal")))))) + .limit(1) + .first() + != null) { + activeTypes.add(IdType.NUMBER); + } + if (col.find( + new BsonDocument( + "_id", + new BsonDocument( + "$not", + new BsonDocument( + "$type", + new BsonArray( + Arrays.asList( + new BsonString("string"), + new BsonString("objectId"), + new BsonString("int"), + new BsonString("long"), + new BsonString("double"), + new BsonString("decimal"))))))) + .limit(1) + .first() + != null) { + activeTypes.add(IdType.OTHER); + } + + if (activeTypes.isEmpty()) { + activeTypes.addAll(EnumSet.allOf(IdType.class)); + } + return activeTypes; + } + + /** + * Discovers empirical quantile boundary points for _id by sampling the collection with $sample. + */ + public static List discoverDataDrivenSplits( + MongoClient client, String databaseName, String collectionName, int numSplits) { + if (numSplits <= 1) { + return Collections.singletonList(new BsonDocument()); + } + int sampleSize = Math.max(1000, numSplits * 64); + List pipeline = + Arrays.asList( + new BsonDocument("$sample", new BsonDocument("size", new BsonInt32(sampleSize))), + new BsonDocument("$project", new BsonDocument("_id", new BsonInt32(1))), + new BsonDocument("$sort", new BsonDocument("_id", new BsonInt32(1)))); + + MongoDatabase db = client.getDatabase(databaseName); + MongoCollection col = db.getCollection(collectionName, BsonDocument.class); + + List sampledKeys = new ArrayList<>(); + for (BsonDocument doc : col.aggregate(pipeline)) { + if (doc.containsKey("_id")) { + sampledKeys.add(doc.get("_id")); + } + } + + if (sampledKeys.size() < numSplits) { + throw new IllegalArgumentException( + "Insufficient sample size: sampled " + + sampledKeys.size() + + " keys, required at least " + + numSplits); + } + + List boundaries = new ArrayList<>(); + int step = sampledKeys.size() / numSplits; + for (int i = 1; i < numSplits; i++) { + BsonValue boundary = sampledKeys.get(i * step); + if (!boundaries.isEmpty() && boundary.equals(boundaries.get(boundaries.size() - 1))) { + throw new IllegalArgumentException("Sampled quantile boundaries contain duplicates"); + } + boundaries.add(boundary); + } + + List slices = new ArrayList<>(); + for (int i = 0; i < numSplits; i++) { + if (i == 0) { + slices.add(new BsonDocument("_id", new BsonDocument("$lt", boundaries.get(0)))); + } else if (i == numSplits - 1) { + slices.add( + new BsonDocument( + "_id", new BsonDocument("$gte", boundaries.get(boundaries.size() - 1)))); + } else { + slices.add( + new BsonDocument( + "_id", + new BsonDocument("$gte", boundaries.get(i - 1)).append("$lt", boundaries.get(i)))); + } + } + return slices; + } + + private static List generateNumberFilters(int numSplits) { + List filters = new ArrayList<>(); + for (int r = 0; r < numSplits; r++) { BsonDocument filter = BsonDocument.parse( String.format( - "{\"_id\": {\"$type\": [\"int\", \"long\"], \"$mod\": [%d, %d]}}", - numNumberSplits, r)); + "{\"_id\": {\"$type\": [\"int\", \"long\", \"double\", \"decimal\"], \"$mod\":" + + " [%d, %d]}}", + numSplits, r)); filters.add(filter); } + return filters; + } - // 2. Double / Decimal slice: {"_id": {"$type": ["double", "decimal"]}} - filters.add(BsonDocument.parse("{\"_id\": {\"$type\": [\"double\", \"decimal\"]}}")); - - // 3. String ASCII prefix slices: {"_id": {"$type": "string", "$gte": "...", "$lt": "..."}} - List stringBounds = generateStringBounds(numStringSplits); + private static List generateStringFilters(int numSplits) { + List filters = new ArrayList<>(); + List stringBounds = generateStringBounds(numSplits); for (int i = 0; i < stringBounds.size() - 1; i++) { String low = stringBounds.get(i); String high = stringBounds.get(i + 1); @@ -75,9 +325,12 @@ public static List generateIndexSliceFilters(int numSplits) { String.format("{\"_id\": {\"$type\": \"string\"%s%s}}", lowClause, highClause)); filters.add(filter); } + return filters; + } - // 4. ObjectId hex timestamp slices: {"_id": {"$gte": ObjectId("..."), "$lt": ObjectId("...")}} - List hexBounds = generateObjectIdBounds(numObjectIdSplits); + private static List generateObjectIdFilters(int numSplits) { + List filters = new ArrayList<>(); + List hexBounds = generateObjectIdBounds(numSplits); for (int i = 0; i < hexBounds.size() - 1; i++) { String lowHex = hexBounds.get(i); String highHex = hexBounds.get(i + 1); @@ -89,19 +342,12 @@ public static List generateIndexSliceFilters(int numSplits) { lowHex, highOp, highHex)); filters.add(filter); } - - // 5. BinData / UUID slice: {"_id": {"$type": "binData"}} - filters.add(BsonDocument.parse("{\"_id\": {\"$type\": \"binData\"}}")); - - // 6. Catch-All slice for remaining BSON types (bool, date, object, array, null, etc.) - filters.add( - BsonDocument.parse( - "{\"_id\": {\"$not\": {\"$type\": [\"int\", \"long\", \"double\", \"decimal\"," - + " \"string\", \"objectId\", \"binData\"]}}}")); - return filters; } + private static final String STRING_SPLIT_CHARS = + "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"; + private static List generateStringBounds(int numSplits) { List bounds = new ArrayList<>(); bounds.add(""); @@ -109,12 +355,11 @@ private static List generateStringBounds(int numSplits) { bounds.add("\uffff"); return bounds; } - int startChar = 48; // '0' - int endChar = 122; // 'z' - int step = Math.max(1, (endChar - startChar) / numSplits); + int maxIndex = STRING_SPLIT_CHARS.length() - 1; + int step = Math.max(1, maxIndex / numSplits); for (int i = 1; i < numSplits; i++) { - int c = Math.min(endChar, startChar + i * step); - bounds.add(String.valueOf((char) c)); + int idx = Math.min(maxIndex, i * step); + bounds.add(String.valueOf(STRING_SPLIT_CHARS.charAt(idx))); } bounds.add("\uffff"); return bounds; diff --git a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java index b7bf261740..7a380ad081 100644 --- a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java +++ b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java @@ -20,6 +20,7 @@ import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertTrue; +import java.util.EnumSet; import java.util.List; import org.bson.BsonDocument; import org.junit.Test; @@ -47,18 +48,13 @@ public void testGenerateIndexSliceFilters_zeroSplit() { @Test public void testGenerateIndexSliceFilters_multipleSplits() { List filters = ReadSplitGenerator.generateIndexSliceFilters(16); - // 4 Integer/Long mod slices + 1 Double/Decimal + 4 String slices + 8 ObjectId slices + 1 - // BinData + 1 Catch-All = 19 slices assertNotNull(filters); assertFalse(filters.isEmpty()); - assertEquals(19, filters.size()); + assertEquals(16, filters.size()); - // Verify all generated filters parse cleanly and contain expected BSON type selectors int numberModCount = 0; - int numberFloatCount = 0; int stringCount = 0; int objectIdCount = 0; - int binDataCount = 0; int catchAllCount = 0; for (BsonDocument filter : filters) { @@ -66,24 +62,84 @@ public void testGenerateIndexSliceFilters_multipleSplits() { String json = filter.toJson(); if (json.contains("\"$not\"")) { catchAllCount++; - } else if (json.contains("\"$mod\"")) { + } + if (json.contains("\"$mod\"")) { numberModCount++; - } else if (json.contains("\"double\"") && json.contains("\"decimal\"")) { - numberFloatCount++; - } else if (json.contains("\"$type\": \"string\"")) { + } + if (json.contains("\"$type\": \"string\"")) { stringCount++; - } else if (json.contains("\"$oid\"")) { + } + if (json.contains("\"$oid\"")) { objectIdCount++; - } else if (json.contains("\"binData\"")) { - binDataCount++; } } - assertEquals(4, numberModCount); - assertEquals(1, numberFloatCount); - assertEquals(4, stringCount); - assertEquals(8, objectIdCount); - assertEquals(1, binDataCount); + assertEquals(16, numberModCount); + assertEquals(16, stringCount); + assertEquals(16, objectIdCount); assertEquals(1, catchAllCount); } + + @Test + public void testGenerateIndexSliceFilters_stringOnly_noOrWrapper() { + List filters = + ReadSplitGenerator.generateIndexSliceFilters( + 4, EnumSet.of(ReadSplitGenerator.IdType.STRING)); + assertEquals(4, filters.size()); + for (BsonDocument filter : filters) { + String json = filter.toJson(); + assertFalse("Single type filter should not contain $or", json.contains("\"$or\"")); + assertTrue("Should contain string type check", json.contains("\"$type\": \"string\"")); + } + } + + @Test + public void testGenerateIndexSliceFilters_objectIdOnly_noOrWrapper() { + List filters = + ReadSplitGenerator.generateIndexSliceFilters( + 4, EnumSet.of(ReadSplitGenerator.IdType.OBJECT_ID)); + assertEquals(4, filters.size()); + for (BsonDocument filter : filters) { + String json = filter.toJson(); + assertFalse("Single type filter should not contain $or", json.contains("\"$or\"")); + assertTrue("Should contain $oid check", json.contains("\"$oid\"")); + } + } + + @Test + public void testGenerateIndexSliceFilters_numberOnly_noOrWrapper() { + List filters = + ReadSplitGenerator.generateIndexSliceFilters( + 4, EnumSet.of(ReadSplitGenerator.IdType.NUMBER)); + assertEquals(4, filters.size()); + for (BsonDocument filter : filters) { + String json = filter.toJson(); + assertFalse("Single type filter should not contain $or", json.contains("\"$or\"")); + assertTrue("Should contain $mod check", json.contains("\"$mod\"")); + } + } + + @Test + public void testGenerateIndexSliceFilters_multipleTypes_usesOrWrapper() { + List filters = + ReadSplitGenerator.generateIndexSliceFilters( + 4, EnumSet.of(ReadSplitGenerator.IdType.STRING, ReadSplitGenerator.IdType.OBJECT_ID)); + assertEquals(4, filters.size()); + for (BsonDocument filter : filters) { + String json = filter.toJson(); + assertTrue("Multiple type filter should contain $or", json.contains("\"$or\"")); + } + } + + @Test + public void testGenerateIndexSliceFilters_otherType_includedInSliceZeroOnly() { + List filters = + ReadSplitGenerator.generateIndexSliceFilters( + 4, EnumSet.of(ReadSplitGenerator.IdType.STRING, ReadSplitGenerator.IdType.OTHER)); + assertEquals(4, filters.size()); + assertTrue(filters.get(0).toJson().contains("\"$not\"")); + assertFalse(filters.get(1).toJson().contains("\"$not\"")); + assertFalse(filters.get(2).toJson().contains("\"$not\"")); + assertFalse(filters.get(3).toJson().contains("\"$not\"")); + } } From 9ede4d0b0d6d102e54ecd4a26a00aec706c67925 Mon Sep 17 00:00:00 2001 From: Michael Le Date: Wed, 29 Jul 2026 17:54:41 +0000 Subject: [PATCH 09/19] Remove useDataDrivenReadSplits flag and automatically attempt data-driven read splits when numReadPrefixSplits > 1 --- .../v2/templates/MongoDbToMongoDb.java | 23 +------- .../v2/transforms/ReadSplitGenerator.java | 53 ++++++++----------- 2 files changed, 23 insertions(+), 53 deletions(-) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java index 6923e0b78e..619bd6b284 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java @@ -288,21 +288,6 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo Integer getWriteRateRampUpSteps(); void setWriteRateRampUpSteps(Integer value); - - @TemplateParameter.Boolean( - order = 22, - groupName = "Source", - optional = true, - description = "Use Data-Driven Read Splits", - helpText = - "If true, the template uses MongoDB index sampling ($sample quantiles) to discover" - + " empirical split boundary keys and automatically detects active _id key types to" - + " eliminate $or overhead. If false, or if sampling is unsupported by the server," - + " it falls back to uniform type-specific prefix splits.") - @Default.Boolean(true) - Boolean getUseDataDrivenReadSplits(); - - void setUseDataDrivenReadSplits(Boolean value); } public static void main(String[] args) { @@ -586,16 +571,10 @@ private static PCollection readFromMongo( Integer numReadSplits = options.getNumReadPrefixSplits(); if (numReadSplits != null && numReadSplits > 1) { List filters; - boolean useDataDriven = - options.getUseDataDrivenReadSplits() == null || options.getUseDataDrivenReadSplits(); try (MongoClient client = MongoClients.create(options.getSourceUri())) { filters = ReadSplitGenerator.generateIndexSliceFilters( - client, - options.getSourceDatabase(), - sourceCollection, - numReadSplits, - useDataDriven); + client, options.getSourceDatabase(), sourceCollection, numReadSplits); } catch (Exception e) { LOG.warn( "Could not connect to MongoDB during setup to generate data-driven read splits ({})." diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java index 2c6c73e53a..3e0c418e86 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java @@ -68,19 +68,14 @@ public static List generateIndexSliceFilters(int numSplits) { * @param databaseName Database name. * @param collectionName Collection name. * @param numSplits Number of target parallel read splits. - * @param useDataDrivenReadSplits Whether to attempt $sample quantile split discovery. * @return List of BsonDocument filters. */ public static List generateIndexSliceFilters( - MongoClient client, - String databaseName, - String collectionName, - int numSplits, - boolean useDataDrivenReadSplits) { + MongoClient client, String databaseName, String collectionName, int numSplits) { if (numSplits <= 1) { return Collections.singletonList(new BsonDocument()); } - if (useDataDrivenReadSplits && client != null) { + if (client != null) { try { List dataDrivenSplits = discoverDataDrivenSplits(client, databaseName, collectionName, numSplits); @@ -169,6 +164,24 @@ public static List generateIndexSliceFilters( return filters; } + private static final BsonArray NUMBER_BSON_TYPES = + new BsonArray( + Arrays.asList( + new BsonString("int"), + new BsonString("long"), + new BsonString("double"), + new BsonString("decimal"))); + + private static final BsonArray KNOWN_BSON_TYPES = + new BsonArray( + Arrays.asList( + new BsonString("string"), + new BsonString("objectId"), + new BsonString("int"), + new BsonString("long"), + new BsonString("double"), + new BsonString("decimal"))); + /** * Detects which _id BSON types are present in a MongoDB collection using lightweight limit(1) * probes. @@ -191,17 +204,7 @@ public static Set detectIdTypes( != null) { activeTypes.add(IdType.OBJECT_ID); } - if (col.find( - new BsonDocument( - "_id", - new BsonDocument( - "$type", - new BsonArray( - Arrays.asList( - new BsonString("int"), - new BsonString("long"), - new BsonString("double"), - new BsonString("decimal")))))) + if (col.find(new BsonDocument("_id", new BsonDocument("$type", NUMBER_BSON_TYPES))) .limit(1) .first() != null) { @@ -209,19 +212,7 @@ public static Set detectIdTypes( } if (col.find( new BsonDocument( - "_id", - new BsonDocument( - "$not", - new BsonDocument( - "$type", - new BsonArray( - Arrays.asList( - new BsonString("string"), - new BsonString("objectId"), - new BsonString("int"), - new BsonString("long"), - new BsonString("double"), - new BsonString("decimal"))))))) + "_id", new BsonDocument("$not", new BsonDocument("$type", KNOWN_BSON_TYPES)))) .limit(1) .first() != null) { From 644b257f7dfb489aeb14431359e3f6cdef498446 Mon Sep 17 00:00:00 2001 From: Michael Le Date: Wed, 29 Jul 2026 18:02:41 +0000 Subject: [PATCH 10/19] Rename numReadPrefixSplits parameter to numReadSplits and remove unused useBucketAuto and numSplits options --- .../v2/templates/MongoDbToMongoDb.java | 63 +++---------------- 1 file changed, 8 insertions(+), 55 deletions(-) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java index 619bd6b284..6bc55fc9fd 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java @@ -132,27 +132,6 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo void setTargetCollection(String value); - @TemplateParameter.Boolean( - order = 7, - groupName = "Source", - optional = true, - description = "Use BucketAuto", - helpText = "Enable withBucketAuto for Atlas compatibility.") - @Default.Boolean(false) - Boolean getUseBucketAuto(); - - void setUseBucketAuto(Boolean value); - - @TemplateParameter.Integer( - order = 8, - groupName = "Source", - optional = true, - description = "Number of Splits", - helpText = "Suggest a specific number of partitions for reading.") - Integer getNumSplits(); - - void setNumSplits(Integer value); - @TemplateParameter.Integer( order = 9, groupName = "Target", @@ -168,14 +147,14 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo order = 10, groupName = "Source", optional = true, - description = "Number of Read Prefix Splits", + description = "Number of Read Splits", helpText = "Number of parallel index-slice queries to generate for high-throughput reads (e.g., 16" + " or 32). Recommended when splitVector is unsupported.") @Default.Integer(0) - Integer getNumReadPrefixSplits(); + Integer getNumReadSplits(); - void setNumReadPrefixSplits(Integer value); + void setNumReadSplits(Integer value); @TemplateParameter.Text( order = 11, @@ -343,15 +322,9 @@ public static void run(Options options) { LOG.info(" Source Collections: {}", sourceCollections); LOG.info( " Read Strategy: {}", - (options.getNumReadPrefixSplits() != null && options.getNumReadPrefixSplits() > 1) - ? "Parallel Index-Slice Reading (numReadPrefixSplits=" - + options.getNumReadPrefixSplits() - + ")" - : "Standard MongoDbIO.read() (bucketAuto=" - + options.getUseBucketAuto() - + ", numSplits=" - + options.getNumSplits() - + ")"); + (options.getNumReadSplits() != null && options.getNumReadSplits() > 1) + ? "Parallel Index-Slice Reading (numReadSplits=" + options.getNumReadSplits() + ")" + : "Standard unpartitioned MongoDbIO.read()"); LOG.info( " Write Configuration: batchSize={}, maxConcurrentAsyncWrites={}, maxWriteRetries={}," + " dlqMaxRetries={}", @@ -568,7 +541,7 @@ public void processElement(ProcessContext c) { private static PCollection readFromMongo( Pipeline pipeline, Options options, String sourceCollection, String targetCollection) { - Integer numReadSplits = options.getNumReadPrefixSplits(); + Integer numReadSplits = options.getNumReadSplits(); if (numReadSplits != null && numReadSplits > 1) { List filters; try (MongoClient client = MongoClients.create(options.getSourceUri())) { @@ -601,13 +574,6 @@ private static PCollection readFromMongo( .withCollection(sourceCollection) .withQueryFn(FindQuery.create().withFilters(filters.get(i))); - if (options.getUseBucketAuto() != null && options.getUseBucketAuto()) { - read = read.withBucketAuto(true); - } - if (options.getNumSplits() != null) { - read = read.withNumSplits(options.getNumSplits()); - } - PCollection branch = pipeline .apply(readGroup + "/Slice_" + i + "/Read", read) @@ -628,12 +594,7 @@ public void processElement(ProcessContext c) { return PCollectionList.of(readBranches).apply(readGroup + "/Merge", Flatten.pCollections()); } - LOG.info( - "Using standard unpartitioned MongoDbIO.read() for collection '{}' (bucketAuto={}," - + " numSplits={})", - sourceCollection, - options.getUseBucketAuto(), - options.getNumSplits()); + LOG.info("Using standard unpartitioned MongoDbIO.read() for collection '{}'", sourceCollection); MongoDbIO.Read read = MongoDbIO.read() @@ -641,14 +602,6 @@ public void processElement(ProcessContext c) { .withDatabase(options.getSourceDatabase()) .withCollection(sourceCollection); - if (options.getUseBucketAuto() != null && options.getUseBucketAuto()) { - read = read.withBucketAuto(true); - } - - if (options.getNumSplits() != null) { - read = read.withNumSplits(options.getNumSplits()); - } - String readGroup = "Read(" + sourceCollection + ")"; return pipeline .apply(readGroup + "/Read", read) From 4731f2df6916b360c87199d2316ded390b59a6be Mon Sep 17 00:00:00 2001 From: Michael Le Date: Wed, 29 Jul 2026 19:08:00 +0000 Subject: [PATCH 11/19] fix: implement Type-Aware Data-Driven Splitting to handle mixed BSON types --- .../v2/transforms/ReadSplitGenerator.java | 150 +++++++++++++----- 1 file changed, 113 insertions(+), 37 deletions(-) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java index 3e0c418e86..a0167b4b1d 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java @@ -75,36 +75,114 @@ public static List generateIndexSliceFilters( if (numSplits <= 1) { return Collections.singletonList(new BsonDocument()); } - if (client != null) { + + Set activeTypes = + client != null + ? detectIdTypes(client, databaseName, collectionName) + : EnumSet.allOf(IdType.class); + + if (client == null) { + return generateIndexSliceFilters(numSplits, activeTypes); + } + + MongoCollection col = + client.getDatabase(databaseName).getCollection(collectionName, BsonDocument.class); + + List numberFilters = Collections.emptyList(); + if (activeTypes.contains(IdType.NUMBER)) { try { - List dataDrivenSplits = - discoverDataDrivenSplits(client, databaseName, collectionName, numSplits); - LOG.info( - "Generated {} data-driven quantile split filters for '{}.{}'", - dataDrivenSplits.size(), + numberFilters = + discoverSplitsForType( + col, numSplits, new BsonDocument("_id", new BsonDocument("$type", NUMBER_BSON_TYPES))); + } catch (Exception e) { + LOG.warn( + "Data-driven splits failed for NUMBER type in '{}.{}' ({}). Falling back to uniform splits.", databaseName, - collectionName); - return dataDrivenSplits; + collectionName, + e.getMessage()); + numberFilters = generateNumberFilters(numSplits); + } + } + + List stringFilters = Collections.emptyList(); + if (activeTypes.contains(IdType.STRING)) { + try { + stringFilters = + discoverSplitsForType( + col, numSplits, new BsonDocument("_id", new BsonDocument("$type", new BsonString("string")))); } catch (Exception e) { LOG.warn( - "Data-driven $sample split discovery failed for '{}.{}' ({}). Falling back to automatic" - + " key-type detection and uniform splits.", + "Data-driven splits failed for STRING type in '{}.{}' ({}). Falling back to uniform splits.", databaseName, collectionName, e.getMessage()); + stringFilters = generateStringFilters(numSplits); } } - Set activeTypes = - client != null - ? detectIdTypes(client, databaseName, collectionName) - : EnumSet.allOf(IdType.class); - LOG.info( - "Generating uniform split filters for active _id BSON types in '{}.{}': {}", - databaseName, - collectionName, - activeTypes); - return generateIndexSliceFilters(numSplits, activeTypes); + List objectIdFilters = Collections.emptyList(); + if (activeTypes.contains(IdType.OBJECT_ID)) { + try { + objectIdFilters = + discoverSplitsForType( + col, numSplits, new BsonDocument("_id", new BsonDocument("$type", new BsonString("objectId")))); + } catch (Exception e) { + LOG.warn( + "Data-driven splits failed for OBJECT_ID type in '{}.{}' ({}). Falling back to uniform splits.", + databaseName, + collectionName, + e.getMessage()); + objectIdFilters = generateObjectIdFilters(numSplits); + } + } + + List otherFilters = Collections.emptyList(); + if (activeTypes.contains(IdType.OTHER)) { + try { + otherFilters = + discoverSplitsForType( + col, + numSplits, + new BsonDocument("_id", new BsonDocument("$not", new BsonDocument("$type", KNOWN_BSON_TYPES)))); + } catch (Exception e) { + LOG.warn( + "Data-driven splits failed for OTHER type in '{}.{}' ({}).", + databaseName, + collectionName, + e.getMessage()); + } + } + + List filters = new ArrayList<>(); + for (int i = 0; i < numSplits; i++) { + List branchFilters = new ArrayList<>(); + if (!numberFilters.isEmpty() && i < numberFilters.size()) { + branchFilters.add(numberFilters.get(i)); + } + if (!stringFilters.isEmpty() && i < stringFilters.size()) { + branchFilters.add(stringFilters.get(i)); + } + if (!objectIdFilters.isEmpty() && i < objectIdFilters.size()) { + branchFilters.add(objectIdFilters.get(i)); + } + if (!otherFilters.isEmpty() && i < otherFilters.size()) { + branchFilters.add(otherFilters.get(i)); + } else if (i == 0 && activeTypes.contains(IdType.OTHER)) { + branchFilters.add( + BsonDocument.parse( + "{\"_id\": {\"$not\": {\"$type\": [\"int\", \"long\", \"double\", \"decimal\"," + + " \"string\", \"objectId\"]}}}")); + } + + if (branchFilters.isEmpty()) { + filters.add(new BsonDocument()); + } else if (branchFilters.size() == 1) { + filters.add(branchFilters.get(0)); + } else { + filters.add(new BsonDocument("$or", new BsonArray(branchFilters))); + } + } + return filters; } /** @@ -225,24 +303,20 @@ public static Set detectIdTypes( return activeTypes; } - /** - * Discovers empirical quantile boundary points for _id by sampling the collection with $sample. - */ - public static List discoverDataDrivenSplits( - MongoClient client, String databaseName, String collectionName, int numSplits) { + private static List discoverSplitsForType( + MongoCollection col, int numSplits, BsonDocument typeMatch) { if (numSplits <= 1) { - return Collections.singletonList(new BsonDocument()); + return Collections.singletonList(typeMatch); } + int sampleSize = Math.max(1000, numSplits * 64); List pipeline = Arrays.asList( + new BsonDocument("$match", typeMatch), new BsonDocument("$sample", new BsonDocument("size", new BsonInt32(sampleSize))), new BsonDocument("$project", new BsonDocument("_id", new BsonInt32(1))), new BsonDocument("$sort", new BsonDocument("_id", new BsonInt32(1)))); - MongoDatabase db = client.getDatabase(databaseName); - MongoCollection col = db.getCollection(collectionName, BsonDocument.class); - List sampledKeys = new ArrayList<>(); for (BsonDocument doc : col.aggregate(pipeline)) { if (doc.containsKey("_id")) { @@ -270,18 +344,20 @@ public static List discoverDataDrivenSplits( List slices = new ArrayList<>(); for (int i = 0; i < numSplits; i++) { + BsonDocument idDoc = new BsonDocument(); + BsonDocument typeMatchId = typeMatch.getDocument("_id"); + for (String key : typeMatchId.keySet()) { + idDoc.append(key, typeMatchId.get(key)); + } + if (i == 0) { - slices.add(new BsonDocument("_id", new BsonDocument("$lt", boundaries.get(0)))); + idDoc.append("$lt", boundaries.get(0)); } else if (i == numSplits - 1) { - slices.add( - new BsonDocument( - "_id", new BsonDocument("$gte", boundaries.get(boundaries.size() - 1)))); + idDoc.append("$gte", boundaries.get(boundaries.size() - 1)); } else { - slices.add( - new BsonDocument( - "_id", - new BsonDocument("$gte", boundaries.get(i - 1)).append("$lt", boundaries.get(i)))); + idDoc.append("$gte", boundaries.get(i - 1)).append("$lt", boundaries.get(i)); } + slices.add(new BsonDocument("_id", idDoc)); } return slices; } From b4cff74af3bd4d4d0ce8a9e7c50d89ef4cd777a0 Mon Sep 17 00:00:00 2001 From: Michael Le Date: Wed, 29 Jul 2026 19:16:39 +0000 Subject: [PATCH 12/19] test: add mockito tests for Type-Aware splits --- .../v2/transforms/ReadSplitGeneratorTest.java | 66 +++++++++++++++++++ 1 file changed, 66 insertions(+) diff --git a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java index 7a380ad081..4446afbaa7 100644 --- a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java +++ b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java @@ -19,7 +19,14 @@ import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import com.mongodb.client.FindIterable; +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoDatabase; import java.util.EnumSet; import java.util.List; import org.bson.BsonDocument; @@ -142,4 +149,63 @@ public void testGenerateIndexSliceFilters_otherType_includedInSliceZeroOnly() { assertFalse(filters.get(2).toJson().contains("\"$not\"")); assertFalse(filters.get(3).toJson().contains("\"$not\"")); } + + @Test + public void testDataDrivenSplits_mixedTypesAreIsolated() { + MongoClient mockClient = mock(MongoClient.class); + MongoDatabase mockDb = mock(MongoDatabase.class); + @SuppressWarnings("unchecked") + MongoCollection mockCol = mock(MongoCollection.class); + + when(mockClient.getDatabase(org.mockito.ArgumentMatchers.anyString())).thenReturn(mockDb); + when(mockDb.getCollection(org.mockito.ArgumentMatchers.anyString(), org.mockito.ArgumentMatchers.eq(BsonDocument.class))).thenReturn(mockCol); + + // Mock detectIdTypes to return multiple types + @SuppressWarnings("unchecked") + FindIterable mockFind = mock(FindIterable.class); + when(mockCol.find(any(BsonDocument.class))).thenReturn(mockFind); + when(mockFind.limit(1)).thenReturn(mockFind); + when(mockFind.first()).thenReturn(new BsonDocument()); // Meaning we detect active types + + // Mock $sample aggregation + @SuppressWarnings("unchecked") + com.mongodb.client.AggregateIterable mockAgg = (com.mongodb.client.AggregateIterable) + java.lang.reflect.Proxy.newProxyInstance( + getClass().getClassLoader(), + new Class[] { com.mongodb.client.AggregateIterable.class }, + (proxy, method, args) -> { + if (method.getName().equals("iterator")) { + return new com.mongodb.client.MongoCursor() { + java.util.Iterator iter = java.util.Arrays.asList( + new BsonDocument("_id", new org.bson.BsonString("min")), + new BsonDocument("_id", new org.bson.BsonString("mid")), + new BsonDocument("_id", new org.bson.BsonString("max"))).iterator(); + @Override public void close() {} + @Override public boolean hasNext() { return iter.hasNext(); } + @Override public BsonDocument next() { return iter.next(); } + @Override public BsonDocument tryNext() { return null; } + @Override public com.mongodb.ServerCursor getServerCursor() { return null; } + @Override public com.mongodb.ServerAddress getServerAddress() { return null; } + @Override public int available() { return 0; } + }; + } + return null; + } + ); + when(mockCol.aggregate(any())).thenReturn(mockAgg); + + List filters = ReadSplitGenerator.generateIndexSliceFilters(mockClient, "db", "col", 2); + + assertEquals(2, filters.size()); + String slice0 = filters.get(0).toJson(); + String slice1 = filters.get(1).toJson(); + + // Validate we use the $or wrapper + assertTrue(slice0.contains("\"$or\"")); + + // Validate that the bounds are nested within specific type bounds! + assertTrue(slice0.contains("\"$type\": \"string\"")); + assertTrue(slice0.contains("\"$type\": \"objectId\"")); + assertTrue(slice0.contains("\"$type\": [\"int\"")); + } } From 0387807ff3b6db4583ce5112cd1ec62ba956d3d6 Mon Sep 17 00:00:00 2001 From: Michael Le Date: Wed, 29 Jul 2026 19:18:01 +0000 Subject: [PATCH 13/19] fix(mongodb): convert verbose batch flush logging to debug level --- .../google/cloud/teleport/v2/transforms/MongoDbTransforms.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java index 48d6e685a1..3bf6b7ae57 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java @@ -595,7 +595,7 @@ private void flushBatch() throws InterruptedException { if (rateLimiter != null && !items.isEmpty()) { rateLimiter.acquire(items.size()); } - LOG.info( + LOG.debug( "Flushing batch of {} documents across {} target collection(s) to MongoDB (active" + " async write futures in queue: {})", items.size(), From cca114db2726635caed59f90929705274cd5af08 Mon Sep 17 00:00:00 2001 From: Michael Le Date: Wed, 29 Jul 2026 20:01:47 +0000 Subject: [PATCH 14/19] docs(mongodb): update template parameter descriptions and group parameter order numbers --- .../v2/templates/MongoDbToMongoDb.java | 129 +++++++++--------- 1 file changed, 66 insertions(+), 63 deletions(-) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java index 6bc55fc9fd..b232cc6bef 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java @@ -133,42 +133,33 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo void setTargetCollection(String value); @TemplateParameter.Integer( - order = 9, - groupName = "Target", - optional = true, - description = "Batch Size", - helpText = "Number of documents in a bulk write.") - @Default.Integer(5000) - Integer getBatchSize(); - - void setBatchSize(Integer value); - - @TemplateParameter.Integer( - order = 10, + order = 7, groupName = "Source", optional = true, description = "Number of Read Splits", helpText = - "Number of parallel index-slice queries to generate for high-throughput reads (e.g., 16" - + " or 32). Recommended when splitVector is unsupported.") + "Number of parallel queries to generate for high-throughput reads (e.g., 16 or 32)." + + " Uses MongoDB's $sample aggregation to discover data-driven boundaries across" + + " active BSON types.") @Default.Integer(0) Integer getNumReadSplits(); void setNumReadSplits(Integer value); - @TemplateParameter.Text( - order = 11, + @TemplateParameter.Integer( + order = 8, + groupName = "Target", optional = true, - description = "DLQ Directory", - helpText = - "Base path to store failed events. Events will be grouped by date and time, and" - + " separated into 'retryable' and 'permanent' subdirectories.") - String getDlqDirectory(); + description = "Batch Size", + helpText = "Number of documents in a bulk write.") + @Default.Integer(5000) + Integer getBatchSize(); - void setDlqDirectory(String value); + void setBatchSize(Integer value); @TemplateParameter.Integer( - order = 13, + order = 9, + groupName = "Target", optional = true, description = "Max Concurrent Async Writes", helpText = "Maximum number of concurrent asynchronous batch writes per worker.") @@ -178,7 +169,8 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo void setMaxConcurrentAsyncWrites(Integer value); @TemplateParameter.Integer( - order = 14, + order = 10, + groupName = "Target", optional = true, description = "Max Write Retries", helpText = "Maximum number of retry attempts for transient failures during write.") @@ -188,45 +180,12 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo void setMaxWriteRetries(Integer value); @TemplateParameter.Integer( - order = 15, - optional = true, - description = "DLQ Max Retries", - helpText = "Maximum number of times to retry events from DLQ.") - @Default.Integer(3) - Integer getDlqMaxRetries(); - - void setDlqMaxRetries(Integer value); - - @TemplateParameter.Text( - order = 16, - groupName = "Source", - optional = true, - description = "Reconsume DLQ Path", - helpText = - "Path to read files from DLQ for reprocessing. If not provided, write DLQ path will be" - + " used.") - String getReconsumeDlqPath(); - - void setReconsumeDlqPath(String value); - - @TemplateParameter.Boolean( - order = 17, - groupName = "Source", - optional = true, - description = "Read from DLQ", - helpText = "If true, reads only from DLQ for retry. If false, reads from MongoDB.") - @Default.Boolean(false) - Boolean getReadFromDlq(); - - void setReadFromDlq(Boolean value); - - @TemplateParameter.Integer( - order = 18, + order = 11, groupName = "Target", optional = true, description = "Initial Write Rate Per Worker", helpText = - "Initial maximum documents/second written per worker thread during Firestore 5/5/5" + "Initial maximum documents/second written per worker thread during linear write rate" + " ramp-up. Set to <= 0 to disable throttling.") @Default.Integer(100) Integer getInitialWriteRatePerWorker(); @@ -234,19 +193,19 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo void setInitialWriteRatePerWorker(Integer value); @TemplateParameter.Integer( - order = 19, + order = 12, groupName = "Target", optional = true, description = "Write Rate Ramp Up Minutes", helpText = - "Number of minutes between 50% rate limit increases during Firestore 5/5/5 ramp-up.") + "Number of minutes between linear rate limit increases during write rate ramp-up.") @Default.Integer(5) Integer getWriteRateRampUpMinutes(); void setWriteRateRampUpMinutes(Integer value); @TemplateParameter.Integer( - order = 20, + order = 13, groupName = "Target", optional = true, description = "Max Write Rate Per Worker", @@ -258,7 +217,7 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo void setMaxWriteRatePerWorker(Integer value); @TemplateParameter.Integer( - order = 21, + order = 14, groupName = "Target", optional = true, description = "Write Rate Ramp Up Steps", @@ -267,6 +226,50 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo Integer getWriteRateRampUpSteps(); void setWriteRateRampUpSteps(Integer value); + + @TemplateParameter.Text( + order = 15, + optional = true, + description = "DLQ Directory", + helpText = + "Base path to store failed events. Events will be grouped by date and time, and" + + " separated into 'retryable' and 'permanent' subdirectories.") + String getDlqDirectory(); + + void setDlqDirectory(String value); + + @TemplateParameter.Integer( + order = 16, + optional = true, + description = "DLQ Max Retries", + helpText = "Maximum number of times to retry events from DLQ.") + @Default.Integer(3) + Integer getDlqMaxRetries(); + + void setDlqMaxRetries(Integer value); + + @TemplateParameter.Text( + order = 17, + groupName = "Source", + optional = true, + description = "Reconsume DLQ Path", + helpText = + "Path to read files from DLQ for reprocessing. If not provided, write DLQ path will be" + + " used.") + String getReconsumeDlqPath(); + + void setReconsumeDlqPath(String value); + + @TemplateParameter.Boolean( + order = 18, + groupName = "Source", + optional = true, + description = "Read from DLQ", + helpText = "If true, reads only from DLQ for retry. If false, reads from MongoDB.") + @Default.Boolean(false) + Boolean getReadFromDlq(); + + void setReadFromDlq(Boolean value); } public static void main(String[] args) { From a0d5f311078309376f9c6e45370dbd96fcc667ab Mon Sep 17 00:00:00 2001 From: Michael Le Date: Thu, 30 Jul 2026 14:31:39 +0000 Subject: [PATCH 15/19] style(mongodb): apply spotless java formatting rules to ReadSplitGenerator and tests --- .../v2/transforms/ReadSplitGenerator.java | 17 ++-- .../v2/transforms/ReadSplitGeneratorTest.java | 82 +++++++++++++------ 2 files changed, 69 insertions(+), 30 deletions(-) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java index a0167b4b1d..83c3c60132 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java @@ -93,7 +93,9 @@ public static List generateIndexSliceFilters( try { numberFilters = discoverSplitsForType( - col, numSplits, new BsonDocument("_id", new BsonDocument("$type", NUMBER_BSON_TYPES))); + col, + numSplits, + new BsonDocument("_id", new BsonDocument("$type", NUMBER_BSON_TYPES))); } catch (Exception e) { LOG.warn( "Data-driven splits failed for NUMBER type in '{}.{}' ({}). Falling back to uniform splits.", @@ -109,7 +111,9 @@ public static List generateIndexSliceFilters( try { stringFilters = discoverSplitsForType( - col, numSplits, new BsonDocument("_id", new BsonDocument("$type", new BsonString("string")))); + col, + numSplits, + new BsonDocument("_id", new BsonDocument("$type", new BsonString("string")))); } catch (Exception e) { LOG.warn( "Data-driven splits failed for STRING type in '{}.{}' ({}). Falling back to uniform splits.", @@ -125,7 +129,9 @@ public static List generateIndexSliceFilters( try { objectIdFilters = discoverSplitsForType( - col, numSplits, new BsonDocument("_id", new BsonDocument("$type", new BsonString("objectId")))); + col, + numSplits, + new BsonDocument("_id", new BsonDocument("$type", new BsonString("objectId")))); } catch (Exception e) { LOG.warn( "Data-driven splits failed for OBJECT_ID type in '{}.{}' ({}). Falling back to uniform splits.", @@ -143,7 +149,8 @@ public static List generateIndexSliceFilters( discoverSplitsForType( col, numSplits, - new BsonDocument("_id", new BsonDocument("$not", new BsonDocument("$type", KNOWN_BSON_TYPES)))); + new BsonDocument( + "_id", new BsonDocument("$not", new BsonDocument("$type", KNOWN_BSON_TYPES)))); } catch (Exception e) { LOG.warn( "Data-driven splits failed for OTHER type in '{}.{}' ({}).", @@ -308,7 +315,7 @@ private static List discoverSplitsForType( if (numSplits <= 1) { return Collections.singletonList(typeMatch); } - + int sampleSize = Math.max(1000, numSplits * 64); List pipeline = Arrays.asList( diff --git a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java index 4446afbaa7..93738322e7 100644 --- a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java +++ b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java @@ -158,7 +158,10 @@ public void testDataDrivenSplits_mixedTypesAreIsolated() { MongoCollection mockCol = mock(MongoCollection.class); when(mockClient.getDatabase(org.mockito.ArgumentMatchers.anyString())).thenReturn(mockDb); - when(mockDb.getCollection(org.mockito.ArgumentMatchers.anyString(), org.mockito.ArgumentMatchers.eq(BsonDocument.class))).thenReturn(mockCol); + when(mockDb.getCollection( + org.mockito.ArgumentMatchers.anyString(), + org.mockito.ArgumentMatchers.eq(BsonDocument.class))) + .thenReturn(mockCol); // Mock detectIdTypes to return multiple types @SuppressWarnings("unchecked") @@ -169,37 +172,66 @@ public void testDataDrivenSplits_mixedTypesAreIsolated() { // Mock $sample aggregation @SuppressWarnings("unchecked") - com.mongodb.client.AggregateIterable mockAgg = (com.mongodb.client.AggregateIterable) - java.lang.reflect.Proxy.newProxyInstance( - getClass().getClassLoader(), - new Class[] { com.mongodb.client.AggregateIterable.class }, - (proxy, method, args) -> { - if (method.getName().equals("iterator")) { + com.mongodb.client.AggregateIterable mockAgg = + (com.mongodb.client.AggregateIterable) + java.lang.reflect.Proxy.newProxyInstance( + getClass().getClassLoader(), + new Class[] {com.mongodb.client.AggregateIterable.class}, + (proxy, method, args) -> { + if (method.getName().equals("iterator")) { return new com.mongodb.client.MongoCursor() { - java.util.Iterator iter = java.util.Arrays.asList( - new BsonDocument("_id", new org.bson.BsonString("min")), - new BsonDocument("_id", new org.bson.BsonString("mid")), - new BsonDocument("_id", new org.bson.BsonString("max"))).iterator(); - @Override public void close() {} - @Override public boolean hasNext() { return iter.hasNext(); } - @Override public BsonDocument next() { return iter.next(); } - @Override public BsonDocument tryNext() { return null; } - @Override public com.mongodb.ServerCursor getServerCursor() { return null; } - @Override public com.mongodb.ServerAddress getServerAddress() { return null; } - @Override public int available() { return 0; } + java.util.Iterator iter = + java.util.Arrays.asList( + new BsonDocument("_id", new org.bson.BsonString("min")), + new BsonDocument("_id", new org.bson.BsonString("mid")), + new BsonDocument("_id", new org.bson.BsonString("max"))) + .iterator(); + + @Override + public void close() {} + + @Override + public boolean hasNext() { + return iter.hasNext(); + } + + @Override + public BsonDocument next() { + return iter.next(); + } + + @Override + public BsonDocument tryNext() { + return null; + } + + @Override + public com.mongodb.ServerCursor getServerCursor() { + return null; + } + + @Override + public com.mongodb.ServerAddress getServerAddress() { + return null; + } + + @Override + public int available() { + return 0; + } }; - } - return null; - } - ); + } + return null; + }); when(mockCol.aggregate(any())).thenReturn(mockAgg); - List filters = ReadSplitGenerator.generateIndexSliceFilters(mockClient, "db", "col", 2); - + List filters = + ReadSplitGenerator.generateIndexSliceFilters(mockClient, "db", "col", 2); + assertEquals(2, filters.size()); String slice0 = filters.get(0).toJson(); String slice1 = filters.get(1).toJson(); - + // Validate we use the $or wrapper assertTrue(slice0.contains("\"$or\"")); From 2ef36a31a248096880a8177eb275915c835c1530 Mon Sep 17 00:00:00 2001 From: Michael Le Date: Thu, 30 Jul 2026 18:03:16 +0000 Subject: [PATCH 16/19] style(mongodb): replace inline fully qualified imports in tests and transforms --- .../v2/templates/MongoDbToMongoDb.java | 5 ++- .../v2/transforms/MongoDbTransforms.java | 3 +- .../v2/transforms/MongoDbTransformsTest.java | 28 ++++++------- .../v2/transforms/ReadSplitGeneratorTest.java | 41 +++++++++++-------- 4 files changed, 43 insertions(+), 34 deletions(-) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java index b232cc6bef..1d4842c7e0 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java @@ -53,6 +53,8 @@ import org.apache.beam.sdk.values.TupleTagList; import org.bson.BsonDocument; import org.bson.Document; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** Dataflow template which copies data from one MongoDB database to another. */ @Template( @@ -64,8 +66,7 @@ optionsClass = MongoDbToMongoDb.Options.class) public class MongoDbToMongoDb { - private static final org.slf4j.Logger LOG = - org.slf4j.LoggerFactory.getLogger(MongoDbToMongoDb.class); + private static final Logger LOG = LoggerFactory.getLogger(MongoDbToMongoDb.class); public interface Options extends JavascriptTextTransformer.JavascriptTextTransformerOptions { @TemplateParameter.Text( diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java index 3bf6b7ae57..546953a0f1 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java @@ -32,6 +32,7 @@ import com.mongodb.client.model.ReplaceOptions; import com.mongodb.client.model.WriteModel; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -739,7 +740,7 @@ private void generateRetryBatch( severeFailedWritesCount.addAndGet(1); } writePermanentDlqMessage( - java.util.Collections.singletonList(failedItem), + Collections.singletonList(failedItem), "Permanent failure writing document. Error: " + error.getMessage()); } else { incDynamicCounter("inMemoryRetries", "MongoBulkWriteException", error.getCode(), 1); diff --git a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java index 4c3e0101d9..0767016283 100644 --- a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java +++ b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java @@ -18,10 +18,12 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyList; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import com.mongodb.MongoBulkWriteException; @@ -32,9 +34,11 @@ import com.mongodb.client.MongoCollection; import com.mongodb.client.MongoDatabase; import com.mongodb.client.model.BulkWriteOptions; +import com.mongodb.client.model.ReplaceOneModel; import com.mongodb.client.model.WriteModel; import java.io.File; import java.io.FileWriter; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; @@ -267,12 +271,9 @@ public void writeWithDlq_documentLevelRetry_partialSuccess() { invocation -> { callCount.getAndIncrement(); List> updates = invocation.getArgument(0); - java.util.List errors = new java.util.ArrayList<>(); + List errors = new ArrayList<>(); for (int i = 0; i < updates.size(); i++) { - Document doc = - (Document) - ((com.mongodb.client.model.ReplaceOneModel) updates.get(i)) - .getReplacement(); + Document doc = (Document) ((ReplaceOneModel) updates.get(i)).getReplacement(); int id = doc.getInteger("_id"); if (id == 1) { errors.add(new BulkWriteError(11000, "Duplicate Key", new BsonDocument(), i)); @@ -285,13 +286,13 @@ public void writeWithDlq_documentLevelRetry_partialSuccess() { } if (!errors.isEmpty()) { throw new MongoBulkWriteException( - mock(com.mongodb.bulk.BulkWriteResult.class), + mock(BulkWriteResult.class), errors, null, - new com.mongodb.ServerAddress(), - java.util.Collections.emptySet()); + new ServerAddress(), + Collections.emptySet()); } - return mock(com.mongodb.bulk.BulkWriteResult.class); + return mock(BulkWriteResult.class); }); DocumentWithMetadata doc0 = DocumentWithMetadata.of(new Document("_id", 0), "test", "test"); @@ -316,7 +317,7 @@ public void writeWithDlq_documentLevelRetry_partialSuccess() { PipelineResult result = pipeline.run(); - org.junit.Assert.assertTrue(callCount.get() >= 2); + assertTrue(callCount.get() >= 2); assertSuccessCount(result, 2L); } @@ -393,8 +394,7 @@ public void applyUdfFn_failure_routesToDlq() throws Exception { .satisfies( collection -> { DocumentWithMetadata result = collection.iterator().next(); - org.junit.Assert.assertTrue( - result.getErrorMessage().contains("UDF failed intentionally")); + assertTrue(result.getErrorMessage().contains("UDF failed intentionally")); return null; }); @@ -469,8 +469,8 @@ public void writeWithDlq_dynamicRouting_writesToCorrectCollection() { pipeline.run(); - org.mockito.Mockito.verify(col1).bulkWrite(anyList(), any(BulkWriteOptions.class)); - org.mockito.Mockito.verify(col2).bulkWrite(anyList(), any(BulkWriteOptions.class)); + verify(col1).bulkWrite(anyList(), any(BulkWriteOptions.class)); + verify(col2).bulkWrite(anyList(), any(BulkWriteOptions.class)); } @Test diff --git a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java index 93738322e7..1ebf29c614 100644 --- a/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java +++ b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java @@ -20,16 +20,26 @@ import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import com.mongodb.ServerAddress; +import com.mongodb.ServerCursor; +import com.mongodb.client.AggregateIterable; import com.mongodb.client.FindIterable; import com.mongodb.client.MongoClient; import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoCursor; import com.mongodb.client.MongoDatabase; +import java.lang.reflect.Proxy; +import java.util.Arrays; import java.util.EnumSet; +import java.util.Iterator; import java.util.List; import org.bson.BsonDocument; +import org.bson.BsonString; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -157,11 +167,8 @@ public void testDataDrivenSplits_mixedTypesAreIsolated() { @SuppressWarnings("unchecked") MongoCollection mockCol = mock(MongoCollection.class); - when(mockClient.getDatabase(org.mockito.ArgumentMatchers.anyString())).thenReturn(mockDb); - when(mockDb.getCollection( - org.mockito.ArgumentMatchers.anyString(), - org.mockito.ArgumentMatchers.eq(BsonDocument.class))) - .thenReturn(mockCol); + when(mockClient.getDatabase(anyString())).thenReturn(mockDb); + when(mockDb.getCollection(anyString(), eq(BsonDocument.class))).thenReturn(mockCol); // Mock detectIdTypes to return multiple types @SuppressWarnings("unchecked") @@ -172,19 +179,19 @@ public void testDataDrivenSplits_mixedTypesAreIsolated() { // Mock $sample aggregation @SuppressWarnings("unchecked") - com.mongodb.client.AggregateIterable mockAgg = - (com.mongodb.client.AggregateIterable) - java.lang.reflect.Proxy.newProxyInstance( + AggregateIterable mockAgg = + (AggregateIterable) + Proxy.newProxyInstance( getClass().getClassLoader(), - new Class[] {com.mongodb.client.AggregateIterable.class}, + new Class[] {AggregateIterable.class}, (proxy, method, args) -> { if (method.getName().equals("iterator")) { - return new com.mongodb.client.MongoCursor() { - java.util.Iterator iter = - java.util.Arrays.asList( - new BsonDocument("_id", new org.bson.BsonString("min")), - new BsonDocument("_id", new org.bson.BsonString("mid")), - new BsonDocument("_id", new org.bson.BsonString("max"))) + return new MongoCursor() { + Iterator iter = + Arrays.asList( + new BsonDocument("_id", new BsonString("min")), + new BsonDocument("_id", new BsonString("mid")), + new BsonDocument("_id", new BsonString("max"))) .iterator(); @Override @@ -206,12 +213,12 @@ public BsonDocument tryNext() { } @Override - public com.mongodb.ServerCursor getServerCursor() { + public ServerCursor getServerCursor() { return null; } @Override - public com.mongodb.ServerAddress getServerAddress() { + public ServerAddress getServerAddress() { return null; } From ea28d8bfb8b42f213d3c3d7a33277c91cf701e92 Mon Sep 17 00:00:00 2001 From: Michael Le Date: Fri, 31 Jul 2026 20:54:36 +0000 Subject: [PATCH 17/19] feat(mongodb): log explicit DLQ directories and inspection command at startup --- .../cloud/teleport/v2/templates/MongoDbToMongoDb.java | 5 +++++ .../cloud/teleport/v2/transforms/MongoDbTransforms.java | 7 +++++++ 2 files changed, 12 insertions(+) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java index 1d4842c7e0..ed53123e1d 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java @@ -344,6 +344,11 @@ public static void run(Options options) { options.getWriteRateRampUpMinutes(), options.getWriteRateRampUpSteps()); LOG.info(" DLQ Base Directory: {}", baseDlqPath + timestampPath); + LOG.info(" DLQ Retryable Directory: {}", retryableDlqPath); + LOG.info(" DLQ Permanent Directory: {}", permanentDlqPath); + LOG.info( + " DLQ Inspection Command: gcloud storage cat \"{}/**/output-*\" | head -n 5", + permanentDlqPath); if (options.getReadFromDlq() != null && options.getReadFromDlq()) { String reconsumePath = options.getReconsumeDlqPath(); diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java index 546953a0f1..8a24e74cdd 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java @@ -192,6 +192,7 @@ public PCollection expand(PCollection, PDone> { + private static final Logger LOG = LoggerFactory.getLogger(WriteToDlq.class); private final String retryablePath; private final String permanentPath; private final String tempLocation; @@ -204,6 +205,12 @@ public WriteToDlq(String retryablePath, String permanentPath, String tempLocatio @Override public PDone expand(PCollection input) { + LOG.info("Configuring DLQ Retryable Output Path: {}", retryablePath); + LOG.info("Configuring DLQ Permanent Output Path: {}", permanentPath); + LOG.info( + "To inspect permanent DLQ errors, run: gcloud storage cat \"{}/**/output-*\" | head -n 5", + permanentPath); + PCollection retryable = input.apply( "FilterRetryable", From 9321030b8aa0f25ec802d9de637ea8771f195e6c Mon Sep 17 00:00:00 2001 From: Michael Le Date: Fri, 31 Jul 2026 20:56:25 +0000 Subject: [PATCH 18/19] feat(mongodb): rate-limit DLQ reason logging to every 30s per worker thread --- .../teleport/v2/transforms/MongoDbTransforms.java | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java index 8a24e74cdd..5be5fbf1ea 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java @@ -271,6 +271,7 @@ public static class WriteFn extends DoFn currentBatch; + private transient long lastDlqLogTimeMs; private void incDynamicCounter(String prefix, String exceptionName, int code, long count) { String counterName = prefix + "_" + exceptionName + "_" + code; @@ -780,9 +782,18 @@ private void writeToDlq( dlqRetriesCount.addAndGet(itemList.size()); } } + long now = System.currentTimeMillis(); + if (now - lastDlqLogTimeMs >= DLQ_LOG_INTERVAL_MS || lastDlqLogTimeMs == 0L) { + lastDlqLogTimeMs = now; + String sampleId = !itemList.isEmpty() ? String.valueOf(itemList.get(0).getId()) : "N/A"; + LOG.warn( + "DLQ Error Summary (logged at most once every 30s per worker thread): {} document(s)" + + " sent to DLQ in this batch. Reason: {} [Sample Doc ID: {}]", + itemList.size(), + message, + sampleId); + } for (DocumentWithMetadata item : itemList) { - LOG.warn("{}: {}", message, item.getId()); - int retryCount = isPermanent ? dlqMaxRetries + 1 : item.getRetryCount() + 1; DocumentWithMetadata.ErrorType errorType = isPermanent ? PERMANENT : RETRYABLE; From 1e3d2771a745a16bfac33f71a1ccee5e8e51a8a9 Mon Sep 17 00:00:00 2001 From: Michael Le Date: Fri, 7 Aug 2026 17:33:25 +0000 Subject: [PATCH 19/19] feat(mongodb): add standard UUID representation support and update DLQ writers --- .../v2/templates/MongoDbToMongoDb.java | 21 +++++----- .../v2/transforms/MongoDbTransforms.java | 40 ++++++++++++++----- 2 files changed, 43 insertions(+), 18 deletions(-) diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java index ed53123e1d..5f4a5fe6bd 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/templates/MongoDbToMongoDb.java @@ -24,7 +24,6 @@ import com.google.cloud.teleport.v2.transforms.ReadSplitGenerator; import com.google.cloud.teleport.v2.transforms.UriSanitizer; import com.mongodb.client.MongoClient; -import com.mongodb.client.MongoClients; import com.mongodb.client.MongoDatabase; import java.text.SimpleDateFormat; import java.util.ArrayList; @@ -290,7 +289,7 @@ public static void run(Options options) { sourceCollections.add(sourceCollection); } else { // List collections from source - try (MongoClient mongoClient = MongoClients.create(sourceUri)) { + try (MongoClient mongoClient = MongoDbTransforms.createMongoClient(sourceUri)) { MongoDatabase db = mongoClient.getDatabase(sourceDatabase); for (String name : db.listCollectionNames()) { sourceCollections.add(name); @@ -358,7 +357,8 @@ public static void run(Options options) { } PCollection documents = readFromDlq(pipeline, reconsumePath); documents.apply( - "ProcessDlq", new ProcessDocuments(options, retryableDlqPath, permanentDlqPath)); + "ProcessDlq", + new ProcessDocuments(options, retryableDlqPath, permanentDlqPath, tmpDirectory)); } else { for (String inputCollection : sourceCollections) { String targetCollectionRaw = options.getTargetCollection(); @@ -371,7 +371,7 @@ public static void run(Options options) { readFromMongo(pipeline, options, inputCollection, targetCollection); documents.apply( "Process_" + inputCollection, - new ProcessDocuments(options, retryableDlqPath, permanentDlqPath)); + new ProcessDocuments(options, retryableDlqPath, permanentDlqPath, tmpDirectory)); } } @@ -383,11 +383,14 @@ public static class ProcessDocuments private final transient Options options; private final String retryableDlqPath; private final String permanentDlqPath; + private final String tmpDirectory; - public ProcessDocuments(Options options, String retryableDlqPath, String permanentDlqPath) { + public ProcessDocuments( + Options options, String retryableDlqPath, String permanentDlqPath, String tmpDirectory) { this.options = options; this.retryableDlqPath = retryableDlqPath; this.permanentDlqPath = permanentDlqPath; + this.tmpDirectory = tmpDirectory; } @Override @@ -430,7 +433,7 @@ public void processElement(ProcessContext c) { .apply( "WriteToDlq_UDF", new MongoDbTransforms.WriteToDlq( - retryableDlqPath, permanentDlqPath, options.getTempLocation())); + retryableDlqPath, permanentDlqPath, tmpDirectory)); documents = udfProcessed @@ -479,7 +482,7 @@ public void processElement(ProcessContext c) { .apply( "WriteToDlq_Validate", new MongoDbTransforms.WriteToDlq( - retryableDlqPath, permanentDlqPath, options.getTempLocation())); + retryableDlqPath, permanentDlqPath, tmpDirectory)); // Write Stage with DLQ PCollection validDocs = processed.get(successTag); @@ -514,7 +517,7 @@ public void processElement(ProcessContext c) { writeFailures.apply( "WriteToDlq_Write", new MongoDbTransforms.WriteToDlq( - retryableDlqPath, permanentDlqPath, options.getTempLocation())); + retryableDlqPath, permanentDlqPath, tmpDirectory)); return PDone.in(input.getPipeline()); } @@ -553,7 +556,7 @@ private static PCollection readFromMongo( Integer numReadSplits = options.getNumReadSplits(); if (numReadSplits != null && numReadSplits > 1) { List filters; - try (MongoClient client = MongoClients.create(options.getSourceUri())) { + try (MongoClient client = MongoDbTransforms.createMongoClient(options.getSourceUri())) { filters = ReadSplitGenerator.generateIndexSliceFilters( client, options.getSourceDatabase(), sourceCollection, numReadSplits); diff --git a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java index 5be5fbf1ea..c3506962ef 100644 --- a/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbTransforms.java @@ -20,8 +20,10 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.util.concurrent.RateLimiter; +import com.mongodb.ConnectionString; import com.mongodb.ErrorCategory; import com.mongodb.MongoBulkWriteException; +import com.mongodb.MongoClientSettings; import com.mongodb.MongoException; import com.mongodb.bulk.BulkWriteError; import com.mongodb.client.MongoClient; @@ -44,6 +46,7 @@ import java.util.concurrent.Semaphore; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; +import org.apache.beam.sdk.io.TextIO; import org.apache.beam.sdk.metrics.Counter; import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.transforms.DoFn; @@ -62,6 +65,7 @@ import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.TupleTagList; import org.bson.Document; +import org.bson.UuidRepresentation; import org.joda.time.Duration; import org.joda.time.Instant; import org.slf4j.Logger; @@ -70,6 +74,21 @@ /** Transforms for the MongoDB to MongoDB template. */ public class MongoDbTransforms { + /** + * Helper method to create a MongoClient with default UuidRepresentation.STANDARD if not explicitly + * specified in the connection string. + */ + public static MongoClient createMongoClient(String uri) { + ConnectionString connectionString = new ConnectionString(uri); + MongoClientSettings.Builder builder = + MongoClientSettings.builder().applyConnectionString(connectionString); + if (connectionString.getUuidRepresentation() == null + || connectionString.getUuidRepresentation() == UuidRepresentation.UNSPECIFIED) { + builder.uuidRepresentation(UuidRepresentation.STANDARD); + } + return MongoClients.create(builder.build()); + } + public static WriteWithDlq writeWithDlq() { return new WriteWithDlq(); } @@ -87,7 +106,8 @@ public static class WriteWithDlq private Integer writeRateRampUpMinutes = 5; private Integer writeRateRampUpSteps = 5; private Integer maxWriteRatePerWorker = 500; - private SerializableFunction clientFactory = MongoClients::create; + private SerializableFunction clientFactory = + MongoDbTransforms::createMongoClient; public WriteWithDlq withUri(String uri) { this.uri = uri; @@ -236,10 +256,9 @@ public void processElement(ProcessContext c) { })) .apply( "WriteDlq_Retryable", - DLQWriteTransform.WriteDLQ.newBuilder() - .withDlqDirectory(retryablePath) - .withTmpDirectory(tempLocation) - .build()); + TextIO.write() + .to(retryablePath + "/error") + .withSuffix(".json")); permanent .apply( @@ -256,10 +275,9 @@ public void processElement(ProcessContext c) { })) .apply( "WriteDlq_Permanent", - DLQWriteTransform.WriteDLQ.newBuilder() - .withDlqDirectory(permanentPath) - .withTmpDirectory(tempLocation) - .build()); + TextIO.write() + .to(permanentPath + "/error") + .withSuffix(".json")); return PDone.in(input.getPipeline()); } @@ -573,6 +591,10 @@ public void processElement(ProcessContext c) throws InterruptedException { if (currentBatch.size() >= batchSize) { flushBatch(); } + DocumentWithMetadata failure; + while ((failure = failures.poll()) != null) { + c.output(failureTag, failure); + } } private void flushBatch() throws InterruptedException {