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..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 @@ -21,8 +21,9 @@ 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.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; @@ -32,6 +33,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; @@ -39,14 +41,19 @@ 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; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** Dataflow template which copies data from one MongoDB database to another. */ @Template( @@ -58,8 +65,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( @@ -126,29 +132,22 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo void setTargetCollection(String value); - @TemplateParameter.Boolean( + @TemplateParameter.Integer( order = 7, groupName = "Source", optional = true, - description = "Use BucketAuto", - helpText = "Enable withBucketAuto for Atlas compatibility.") - @Default.Boolean(false) - Boolean getUseBucketAuto(); + description = "Number of Read Splits", + helpText = + "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 setUseBucketAuto(Boolean value); + void setNumReadSplits(Integer 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", optional = true, description = "Batch Size", @@ -158,19 +157,9 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo void setBatchSize(Integer value); - @TemplateParameter.Text( - order = 11, - 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 = 13, + order = 9, + groupName = "Target", optional = true, description = "Max Concurrent Async Writes", helpText = "Maximum number of concurrent asynchronous batch writes per worker.") @@ -180,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.") @@ -190,8 +180,67 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo void setMaxWriteRetries(Integer value); @TemplateParameter.Integer( + order = 11, + groupName = "Target", + optional = true, + description = "Initial Write Rate Per Worker", + helpText = + "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(); + + void setInitialWriteRatePerWorker(Integer value); + + @TemplateParameter.Integer( + order = 12, + groupName = "Target", + optional = true, + description = "Write Rate Ramp Up Minutes", + helpText = + "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 = 13, + 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 = 14, + 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); + + @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) @@ -200,7 +249,7 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo void setDlqMaxRetries(Integer value); @TemplateParameter.Text( - order = 16, + order = 17, groupName = "Source", optional = true, description = "Reconsume DLQ Path", @@ -212,7 +261,7 @@ public interface Options extends JavascriptTextTransformer.JavascriptTextTransfo void setReconsumeDlqPath(String value); @TemplateParameter.Boolean( - order = 17, + order = 18, groupName = "Source", optional = true, description = "Read from DLQ", @@ -240,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); @@ -268,6 +317,38 @@ 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.getNumReadSplits() != null && options.getNumReadSplits() > 1) + ? "Parallel Index-Slice Reading (numReadSplits=" + options.getNumReadSplits() + ")" + : "Standard unpartitioned MongoDbIO.read()"); + LOG.info( + " Write Configuration: batchSize={}, maxConcurrentAsyncWrites={}, maxWriteRetries={}," + + " dlqMaxRetries={}", + options.getBatchSize(), + 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); + 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(); if (reconsumePath == null || reconsumePath.isEmpty()) { @@ -276,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(); @@ -289,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)); } } @@ -301,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 @@ -348,7 +433,7 @@ public void processElement(ProcessContext c) { .apply( "WriteToDlq_UDF", new MongoDbTransforms.WriteToDlq( - retryableDlqPath, permanentDlqPath, options.getTempLocation())); + retryableDlqPath, permanentDlqPath, tmpDirectory)); documents = udfProcessed @@ -397,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); @@ -411,12 +496,28 @@ 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", new MongoDbTransforms.WriteToDlq( - retryableDlqPath, permanentDlqPath, options.getTempLocation())); + retryableDlqPath, permanentDlqPath, tmpDirectory)); return PDone.in(input.getPipeline()); } @@ -452,24 +553,72 @@ public void processElement(ProcessContext c) { private static PCollection readFromMongo( Pipeline pipeline, Options options, String sourceCollection, String targetCollection) { + Integer numReadSplits = options.getNumReadSplits(); + if (numReadSplits != null && numReadSplits > 1) { + List filters; + try (MongoClient client = MongoDbTransforms.createMongoClient(options.getSourceUri())) { + filters = + ReadSplitGenerator.generateIndexSliceFilters( + client, options.getSourceDatabase(), sourceCollection, numReadSplits); + } 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( + "Generating {} parallel index-slice read branches for collection '{}'", + 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); + MongoDbIO.Read read = + MongoDbIO.read() + .withUri(options.getSourceUri()) + .withDatabase(options.getSourceDatabase()) + .withCollection(sourceCollection) + .withQueryFn(FindQuery.create().withFilters(filters.get(i))); + + PCollection branch = + pipeline + .apply(readGroup + "/Slice_" + i + "/Read", read) + .apply( + readGroup + "/Slice_" + i + "/MapToMetadata", + 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(readGroup + "/Merge", Flatten.pCollections()); + } + + LOG.info("Using standard unpartitioned MongoDbIO.read() for collection '{}'", sourceCollection); + MongoDbIO.Read read = MongoDbIO.read() .withUri(options.getSourceUri()) .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("Read_" + sourceCollection, read) + .apply(readGroup + "/Read", read) .apply( - "MapToMetadata_" + sourceCollection, + readGroup + "/MapToMetadata", ParDo.of( new DoFn() { @ProcessElement 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..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 @@ -18,8 +18,12 @@ 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.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; @@ -30,6 +34,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; @@ -39,31 +44,28 @@ 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.coders.KvCoder; -import org.apache.beam.sdk.coders.SerializableCoder; -import org.apache.beam.sdk.coders.StringUtf8Coder; +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; 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; 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; @@ -72,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(); } @@ -85,7 +102,12 @@ public static class WriteWithDlq private Integer maxConcurrentAsyncWrites = 10; private Integer maxWriteRetries = 3; private Integer dlqMaxRetries = 3; - private SerializableFunction clientFactory = MongoClients::create; + private Integer initialWriteRatePerWorker = 100; + private Integer writeRateRampUpMinutes = 5; + private Integer writeRateRampUpSteps = 5; + private Integer maxWriteRatePerWorker = 500; + private SerializableFunction clientFactory = + MongoDbTransforms::createMongoClient; public WriteWithDlq withUri(String uri) { this.uri = uri; @@ -125,6 +147,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; @@ -138,36 +188,31 @@ 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) + .withInitialWriteRatePerWorker(initialWriteRatePerWorker) + .withWriteRateRampUpMinutes(writeRateRampUpMinutes) + .withWriteRateRampUpSteps(writeRateRampUpSteps) + .withMaxWriteRatePerWorker(maxWriteRatePerWorker) + .withClientFactory(clientFactory) + .withFailureTag(failureTag) + .build()) + .withOutputTags(successTag, TupleTagList.of(failureTag))); return writeResults.get(failureTag); } } public static class WriteToDlq extends PTransform, PDone> { + private static final Logger LOG = LoggerFactory.getLogger(WriteToDlq.class); private final String retryablePath; private final String permanentPath; private final String tempLocation; @@ -180,6 +225,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", @@ -205,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( @@ -225,33 +275,40 @@ 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()); } } /** A {@link DoFn} that writes documents to MongoDB in bulk. */ - public static class WriteFn - extends DoFn>, DocumentWithMetadata> { + public static class WriteFn extends DoFn { private static final int ERR_DOCUMENT_VALIDATION_FAILURE = 121; private static final int ERR_KEY_TOO_LONG = 17280; private static final int ERR_BAD_VALUE = 2; + private static final long DLQ_LOG_INTERVAL_MS = 30_000L; private static final Logger LOG = LoggerFactory.getLogger(WriteFn.class); private final String uri; private final String database; + private final Integer batchSize; private final Integer maxConcurrentAsyncWrites; private final Integer maxWriteRetries; private final Integer dlqMaxRetries; + private final Integer initialWriteRatePerWorker; + private final Integer writeRateRampUpMinutes; + private final Integer writeRateRampUpSteps; + private final Integer maxWriteRatePerWorker; private final SerializableFunction 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 +330,8 @@ public static class WriteFn private transient AtomicLong severeFailedWritesCount; private transient AtomicLong dlqRetriesCount; private transient AtomicLong permanentFailuresCount; + private transient List currentBatch; + private transient long lastDlqLogTimeMs; private void incDynamicCounter(String prefix, String exceptionName, int code, long count) { String counterName = prefix + "_" + exceptionName + "_" + code; @@ -284,16 +343,26 @@ private void incDynamicCounter(String prefix, String exceptionName, int code, lo public WriteFn( String uri, String database, + Integer batchSize, Integer maxConcurrentAsyncWrites, Integer maxWriteRetries, Integer dlqMaxRetries, + Integer initialWriteRatePerWorker, + Integer writeRateRampUpMinutes, + Integer writeRateRampUpSteps, + Integer maxWriteRatePerWorker, SerializableFunction clientFactory, TupleTag failureTag) { this.uri = uri; this.database = database; + this.batchSize = batchSize; 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; } @@ -305,9 +374,14 @@ public static Builder builder() { public static class Builder { private String uri; private String database; - private Integer maxConcurrentAsyncWrites; - private Integer maxWriteRetries; + private Integer batchSize = 5000; + 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; @@ -321,13 +395,24 @@ 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; + if (maxConcurrentAsyncWrites != null) { + this.maxConcurrentAsyncWrites = maxConcurrentAsyncWrites; + } return this; } public Builder withMaxWriteRetries(Integer maxWriteRetries) { - this.maxWriteRetries = maxWriteRetries; + if (maxWriteRetries != null) { + this.maxWriteRetries = maxWriteRetries; + } return this; } @@ -336,6 +421,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; @@ -350,14 +463,34 @@ public WriteFn build() { return new WriteFn( uri, database, + batchSize, 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); @@ -367,6 +500,66 @@ 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={})", + database, + batchSize, + maxConcurrentAsyncWrites, + 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 @@ -388,11 +581,28 @@ public void startBundle() { severeFailedWritesCount = new AtomicLong(0); dlqRetriesCount = new AtomicLong(0); permanentFailuresCount = new AtomicLong(0); + currentBatch = new ArrayList<>(); + LOG.debug("Starting new write bundle session (URI: {})", UriSanitizer.sanitize(uri)); } @ProcessElement public void processElement(ProcessContext c) throws InterruptedException { - Iterable items = c.element().getValue(); + currentBatch.add(c.element()); + if (currentBatch.size() >= batchSize) { + flushBatch(); + } + DocumentWithMetadata failure; + while ((failure = failures.poll()) != null) { + c.output(failureTag, failure); + } + } + + private void flushBatch() throws InterruptedException { + if (currentBatch.isEmpty()) { + return; + } + List items = currentBatch; + currentBatch = new ArrayList<>(); Map>> updatesByCollection = new HashMap<>(); Map> itemsByCollection = new HashMap<>(); @@ -413,6 +623,16 @@ public void processElement(ProcessContext c) throws InterruptedException { } if (!updatesByCollection.isEmpty()) { + updateRateLimiterIfNeeded(); + if (rateLimiter != null && !items.isEmpty()) { + rateLimiter.acquire(items.size()); + } + LOG.debug( + "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( @@ -427,7 +647,7 @@ public void processElement(ProcessContext c) throws InterruptedException { MongoCollection col = mongoClient.getDatabase(database).getCollection(colName); - writeBatchWithRetry(col, currentUpdates, currentItemList); + writeBatchWithRetry(colName, col, currentUpdates, currentItemList); } } finally { semaphore.release(); @@ -439,6 +659,7 @@ public void processElement(ProcessContext c) throws InterruptedException { } private void writeBatchWithRetry( + String colName, MongoCollection col, List> currentUpdates, List currentItemList) { @@ -449,10 +670,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<>(); @@ -480,6 +711,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; @@ -490,6 +726,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; } @@ -529,7 +771,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); @@ -562,9 +804,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; @@ -598,6 +849,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) { @@ -626,6 +884,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/ReadSplitGenerator.java b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java new file mode 100644 index 0000000000..83c3c60132 --- /dev/null +++ b/v2/mongodb-to-mongodb/src/main/java/com/google/cloud/teleport/v2/transforms/ReadSplitGenerator.java @@ -0,0 +1,460 @@ +/* + * 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 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 + * requiring MongoDB splitVector or bucketAuto commands. + */ +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 + * 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. + * @return List of BsonDocument filters. + */ + public static List generateIndexSliceFilters( + MongoClient client, String databaseName, String collectionName, int numSplits) { + if (numSplits <= 1) { + return Collections.singletonList(new BsonDocument()); + } + + 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 { + 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, + 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 splits failed for STRING type in '{}.{}' ({}). Falling back to uniform splits.", + databaseName, + collectionName, + e.getMessage()); + stringFilters = generateStringFilters(numSplits); + } + } + + 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; + } + + /** + * 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\"]}}}")); + } + + 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; + } + + 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. + */ + 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); + + 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", NUMBER_BSON_TYPES))) + .limit(1) + .first() + != null) { + activeTypes.add(IdType.NUMBER); + } + if (col.find( + new BsonDocument( + "_id", new BsonDocument("$not", new BsonDocument("$type", KNOWN_BSON_TYPES)))) + .limit(1) + .first() + != null) { + activeTypes.add(IdType.OTHER); + } + + if (activeTypes.isEmpty()) { + activeTypes.addAll(EnumSet.allOf(IdType.class)); + } + return activeTypes; + } + + private static List discoverSplitsForType( + MongoCollection col, int numSplits, BsonDocument typeMatch) { + if (numSplits <= 1) { + 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)))); + + 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++) { + BsonDocument idDoc = new BsonDocument(); + BsonDocument typeMatchId = typeMatch.getDocument("_id"); + for (String key : typeMatchId.keySet()) { + idDoc.append(key, typeMatchId.get(key)); + } + + if (i == 0) { + idDoc.append("$lt", boundaries.get(0)); + } else if (i == numSplits - 1) { + idDoc.append("$gte", boundaries.get(boundaries.size() - 1)); + } else { + idDoc.append("$gte", boundaries.get(i - 1)).append("$lt", boundaries.get(i)); + } + slices.add(new BsonDocument("_id", idDoc)); + } + 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\", \"double\", \"decimal\"], \"$mod\":" + + " [%d, %d]}}", + numSplits, r)); + filters.add(filter); + } + return filters; + } + + 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); + 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); + } + return filters; + } + + 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); + 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); + } + return filters; + } + + private static final String STRING_SPLIT_CHARS = + "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"; + + private static List generateStringBounds(int numSplits) { + List bounds = new ArrayList<>(); + bounds.add(""); + if (numSplits == 1) { + bounds.add("\uffff"); + return bounds; + } + int maxIndex = STRING_SPLIT_CHARS.length() - 1; + int step = Math.max(1, maxIndex / numSplits); + for (int i = 1; i < numSplits; i++) { + int idx = Math.min(maxIndex, i * step); + bounds.add(String.valueOf(STRING_SPLIT_CHARS.charAt(idx))); + } + 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/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/MongoDbTransformsTest.java b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbTransformsTest.java index 53becbdfd8..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 @@ -16,10 +16,14 @@ 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.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; @@ -30,18 +34,16 @@ 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; 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 +51,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,26 +263,34 @@ 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); + List errors = new ArrayList<>(); + for (int i = 0; i < updates.size(); i++) { + 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)); + } 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)), + errors, null, new ServerAddress(), Collections.emptySet()); - } else if (count == 1) { - List> updates = invocation.getArgument(0); - assertEquals(1, updates.size()); - return mock(BulkWriteResult.class); } return mock(BulkWriteResult.class); }); @@ -290,17 +299,7 @@ public void writeWithDlq_documentLevelRetry_partialSuccess() 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 +307,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 +317,7 @@ public void writeWithDlq_documentLevelRetry_partialSuccess() PipelineResult result = pipeline.run(); - assertEquals(2, callCount.get()); + assertTrue(callCount.get() >= 2); assertSuccessCount(result, 2L); } @@ -394,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; }); @@ -470,7 +469,53 @@ 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 + 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(); } } 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..1ebf29c614 --- /dev/null +++ b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/ReadSplitGeneratorTest.java @@ -0,0 +1,250 @@ +/* + * 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 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; + +/** 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); + assertNotNull(filters); + assertFalse(filters.isEmpty()); + assertEquals(16, filters.size()); + + int numberModCount = 0; + int stringCount = 0; + int objectIdCount = 0; + int catchAllCount = 0; + + for (BsonDocument filter : filters) { + assertNotNull(filter); + String json = filter.toJson(); + if (json.contains("\"$not\"")) { + catchAllCount++; + } + if (json.contains("\"$mod\"")) { + numberModCount++; + } + if (json.contains("\"$type\": \"string\"")) { + stringCount++; + } + if (json.contains("\"$oid\"")) { + objectIdCount++; + } + } + + 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\"")); + } + + @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(anyString())).thenReturn(mockDb); + when(mockDb.getCollection(anyString(), 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") + AggregateIterable mockAgg = + (AggregateIterable) + Proxy.newProxyInstance( + getClass().getClassLoader(), + new Class[] {AggregateIterable.class}, + (proxy, method, args) -> { + if (method.getName().equals("iterator")) { + 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 + 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 ServerCursor getServerCursor() { + return null; + } + + @Override + public 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\"")); + } +} 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..1fa58d5238 --- /dev/null +++ b/v2/mongodb-to-mongodb/src/test/java/com/google/cloud/teleport/v2/transforms/UriSanitizerTest.java @@ -0,0 +1,56 @@ +/* + * 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("")); + } +}