diff --git a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/templates/DataStreamMongoDBToFirestore.java b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/templates/DataStreamMongoDBToFirestore.java index 62d05a0c35..4de6095d45 100644 --- a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/templates/DataStreamMongoDBToFirestore.java +++ b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/templates/DataStreamMongoDBToFirestore.java @@ -37,18 +37,21 @@ import com.google.cloud.teleport.v2.transforms.DLQWriteTransform; import com.google.cloud.teleport.v2.transforms.JavascriptTextTransformer.FailsafeJavascriptUdf; import com.google.cloud.teleport.v2.transforms.JavascriptTextTransformer.JavascriptTextTransformerOptions; +import com.google.cloud.teleport.v2.transforms.MongoDbBulkTransforms; +import com.google.cloud.teleport.v2.transforms.MongoDbChangeEventContextCoder; import com.google.cloud.teleport.v2.transforms.MongoDbEventDeadLetterQueueSanitizer; import com.google.cloud.teleport.v2.transforms.ProcessChangeEventFn; +import com.google.cloud.teleport.v2.transforms.StatefulDeduplicationFn; +import com.google.cloud.teleport.v2.transforms.TimestampSortKey; +import com.google.cloud.teleport.v2.transforms.TimestampSortKeyCoder; import com.google.cloud.teleport.v2.transforms.Utils; import com.google.cloud.teleport.v2.values.FailsafeElement; import com.google.common.base.Strings; import com.google.common.base.Throwables; import com.mongodb.MongoBulkWriteException; -import com.mongodb.MongoClientSettings; import com.mongodb.bulk.BulkWriteError; import com.mongodb.bulk.BulkWriteResult; import com.mongodb.client.MongoClient; -import com.mongodb.client.MongoClients; import com.mongodb.client.MongoCollection; import com.mongodb.client.MongoDatabase; import com.mongodb.client.model.BulkWriteOptions; @@ -68,7 +71,6 @@ import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; -import org.apache.beam.sdk.coders.SerializableCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.io.FileSystems; import org.apache.beam.sdk.metrics.Counter; @@ -82,14 +84,19 @@ import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.Reshuffle; +import org.apache.beam.sdk.transforms.WithKeys; import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindows; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PBegin; 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.TupleTag; import org.apache.beam.sdk.values.TupleTagList; +import org.apache.beam.sdk.values.TypeDescriptors; import org.bson.Document; -import org.bson.UuidRepresentation; import org.bson.conversions.Bson; import org.joda.time.Duration; import org.joda.time.Instant; @@ -114,20 +121,10 @@ + " Storage bucket and writes them to a Firestore with MongoDB compatibility database. It" + " is intended for data migration from Datastream sources to Firestore with MongoDB" + " compatibility.\n", - "Data consistency is guaranteed only at the end of migration when all data has been written" - + " to the destination database. To store ordering information for each record written to" - + " the destination database, this template creates an additional collection (called a" - + " shadow collection) for each collection in the source database. This is used to ensure" - + " consistency at the end of migration. By default the shadow collection is used only on" - + " cdc events, it is configurable to be used on backfill events via setting" - + " `useShadowTablesForBackfill` to true. The shadow collections by default uses prefix" - + " `shadow_`, if it can cause collection name collision with the source database, please" - + " configure that by setting `shadowCollectionPrefix`. The shadow collections are not" - + " deleted after migration and can be used for validation purposes at the end of the" - + " migration.\n", - "The pipeline by default processes backfill events first with batch write, which is" - + " optimized for performance, followed by cdc events. This is configurable via setting" - + " `processBackfillFirst` to false to process backfill and cdc events together.\n", + "By default, the template runs in high-throughput shadowless mode without shadow collections" + + " or distributed multi-document transactions. When legacy mode is explicitly selected by" + + " setting `useShadowTables` to true, the template creates an additional shadow collection" + + " for each collection to track event ordering.\n", "Any errors that occur during operation are recorded in error queues. The error" + " queue is a Cloud Storage folder which stores all the Datastream events that had" + " encountered errors." @@ -144,31 +141,7 @@ public class DataStreamMongoDBToFirestore { static final TupleTag> BYPASS_UDF_TAG = new TupleTag<>(); private static final String AVRO_SUFFIX = "avro"; private static final String JSON_SUFFIX = "json"; - public static final Set MAPPER_IGNORE_FIELDS = - new HashSet( - Arrays.asList( - "_metadata_stream", - "_metadata_schema", - "_metadata_table", - "_metadata_source", - "_metadata_ssn", - "_metadata_rs_id", - "_metadata_tx_id", - "_metadata_uuid", - "_metadata_dlq_reconsumed", - "_metadata_error", - "_metadata_retry_count", - "_metadata_timestamp", - "_metadata_read_timestamp", - "_metadata_read_method", - "_metadata_deleted", - "_metadata_primary_keys", - "_metadata_log_file", - "_metadata_log_position", - "_metadata_dataflow_timestamp", - "data", - "_metadata_timestamp_seconds", - "_metadata_timestamp_nanos")); + public static final Set MAPPER_IGNORE_FIELDS = DatastreamConstants.MAPPER_IGNORE_FIELDS; /** * Options supported by the pipeline. @@ -179,9 +152,77 @@ public interface Options extends StreamingOptions, DataflowPipelineWorkerPoolOptions, JavascriptTextTransformerOptions { - @TemplateParameter.Text( + + @TemplateParameter.Boolean( order = 10, optional = true, + description = "Use shadow tables for tracking event ordering", + helpText = + "When false (default), runs in high-throughput shadowless mode without shadow" + + " collections.") + @Default.Boolean(false) + Boolean getUseShadowTables(); + + void setUseShadowTables(Boolean value); + + @TemplateParameter.Integer( + order = 11, + optional = true, + description = "Batch size for bulk database writes", + helpText = + "Number of documents per bulkWrite RPC. For Firestore MongoDB compatibility, max 500." + + " Default: 500.") + @Default.Integer(500) + Integer getBatchSize(); + + void setBatchSize(Integer value); + + @TemplateParameter.Integer( + order = 13, + optional = true, + description = "Maximum concurrent asynchronous writes per worker", + helpText = + "Maximum concurrent in-flight bulk write operations per worker thread pool. Default: 10.") + @Default.Integer(10) + Integer getMaxConcurrentAsyncWrites(); + + void setMaxConcurrentAsyncWrites(Integer value); + + @TemplateParameter.Integer( + order = 14, + optional = true, + description = "Initial write rate per worker (docs/sec)", + helpText = + "Initial maximum write rate per worker during warm-up. Set <= 0 to disable. Default: 500.") + @Default.Integer(500) + Integer getInitialWriteRatePerWorker(); + + void setInitialWriteRatePerWorker(Integer value); + + @TemplateParameter.Integer( + order = 15, + optional = true, + description = "Write rate ramp-up duration in minutes", + helpText = + "Duration in minutes over which the write rate ramps up to target throughput. Default: 5.") + @Default.Integer(5) + Integer getWriteRateRampUpMinutes(); + + void setWriteRateRampUpMinutes(Integer value); + + @TemplateParameter.Integer( + order = 16, + optional = true, + description = "Max write rate per worker after ramp-up", + helpText = "Target maximum write rate per worker after completing ramp-up. Default: 2500.") + @Default.Integer(2500) + Integer getMaxWriteRatePerWorker(); + + void setMaxWriteRatePerWorker(Integer value); + + @TemplateParameter.Text( + order = 17, + optional = true, description = "Shadow collection prefix", helpText = "The prefix used to name shadow collections. Default: `shadow_`.") @Default.String(DatastreamConstants.DEFAULT_SHADOW_COLLECTION_PREFIX) @@ -391,16 +432,6 @@ public interface Options String getDatabaseCollection(); void setDatabaseCollection(String value); - - @TemplateParameter.Integer( - order = 11, - optional = true, - description = "Batch size", - helpText = "The batch size for writing to Database.") - @Default.Integer(500) - Integer getBatchSize(); - - void setBatchSize(Integer value); } /** @@ -425,77 +456,103 @@ public static void main(String[] args) { run(options); } - private static void validateOptions(Options options) { + public static void validateOptions(Options options) { + String connectionUri = options.getConnectionUri(); + if (connectionUri == null || connectionUri.trim().isEmpty()) { + throw new IllegalArgumentException( + "Connection URI (connectionUri) must be specified and non-empty. " + + "Expected 'mongodb://...' or 'mongodb+srv://...'"); + } + if (!connectionUri.startsWith("mongodb://") && !connectionUri.startsWith("mongodb+srv://")) { + throw new IllegalArgumentException( + "Invalid connectionUri: " + + connectionUri + + ". Must start with 'mongodb://' or 'mongodb+srv://'"); + } + + String databaseName = options.getDatabaseName(); + if (databaseName == null || databaseName.trim().isEmpty()) { + throw new IllegalArgumentException( + "Database name (databaseName) must be specified and non-empty."); + } + String inputFileFormat = options.getInputFileFormat(); - if (!(inputFileFormat.equals(AVRO_SUFFIX) || inputFileFormat.equals(JSON_SUFFIX))) { + if (inputFileFormat != null + && !inputFileFormat.isEmpty() + && !(inputFileFormat.equals(AVRO_SUFFIX) || inputFileFormat.equals(JSON_SUFFIX))) { throw new IllegalArgumentException( "Input file format must be one of: avro, json or left empty - found " + inputFileFormat); } + + if (options.getBatchSize() != null && options.getBatchSize() <= 0) { + throw new IllegalArgumentException( + "Batch size must be a positive integer - found " + options.getBatchSize()); + } + + if (options.getMaxConcurrentAsyncWrites() != null + && options.getMaxConcurrentAsyncWrites() <= 0) { + throw new IllegalArgumentException( + "Max concurrent async writes must be a positive integer - found " + + options.getMaxConcurrentAsyncWrites()); + } + + if (options.getInitialWriteRatePerWorker() != null + && options.getMaxWriteRatePerWorker() != null + && options.getInitialWriteRatePerWorker() > 0 + && options.getMaxWriteRatePerWorker() > 0 + && options.getInitialWriteRatePerWorker() > options.getMaxWriteRatePerWorker()) { + throw new IllegalArgumentException( + "Initial write rate per worker (" + + options.getInitialWriteRatePerWorker() + + ") cannot exceed max write rate per worker (" + + options.getMaxWriteRatePerWorker() + + ")"); + } + + if (options.getWriteRateRampUpMinutes() != null && options.getWriteRateRampUpMinutes() < 0) { + throw new IllegalArgumentException( + "Write rate ramp up minutes cannot be negative - found " + + options.getWriteRateRampUpMinutes()); + } } /** * Runs the pipeline with the supplied options. * - *

This pipeline processes all events (backfill/CDC) together, ordered by the timestamp field - * from the datastream records. Shadow collections are used to track event ordering and prevent - * duplicate processing. - * * @param options The execution parameters to the pipeline. */ public static void run(Options options) { try { + validateOptions(options); + LOG.info( "Starting pipeline execution with options: inputFilePattern={}, fileType={}," - + " databaseName={}", + + " databaseName={}, useShadowTables={}", options.getInputFilePattern(), options.getInputFileFormat(), - options.getDatabaseName()); + options.getDatabaseName(), + options.getUseShadowTables()); // Decode the connection string String connectionString = options.getConnectionUri(); - if (!connectionString.startsWith("mongodb://") - && !connectionString.startsWith("mongodb+srv://")) { - LOG.error( - "Invalid URL: {}, Must be in pattern of" - + " 'mongodb://:,:/database?options', or" - + " 'mongodb+srv:///database?options'", - connectionString); - throw new IllegalArgumentException("Invalid connectionUri: " + connectionString); - } - if (connectionString.contains("MONGODB-OIDC") + if (connectionString != null + && connectionString.contains("MONGODB-OIDC") && !connectionString.contains("TOKEN_RESOURCE")) { connectionString += ",TOKEN_RESOURCE:FIRESTORE"; } - LOG.info("Creating MongoDB client with connection string: {}", connectionString); - MongoClientSettings settings = - MongoClientSettings.builder() - .applyConnectionString(new com.mongodb.ConnectionString(connectionString)) - .applyToSocketSettings( - builder -> { - // How long the driver will wait to establish a connection - builder.connectTimeout(60, TimeUnit.SECONDS); - builder.readTimeout(60, TimeUnit.SECONDS); // Example: 60 seconds - }) - .applyToClusterSettings( - builder -> builder.serverSelectionTimeout(10, TimeUnit.MINUTES)) - .uuidRepresentation(UuidRepresentation.STANDARD) - .build(); - MongoClient mongoClient = MongoClients.create(settings); - LOG.info("MongoDB client created successfully"); // Choose processing mode based on options LOG.info("Starting pipeline execution"); - PipelineResult result; - if (options.getProcessBackfillFirst()) { - LOG.info("Using backfill-first processing mode"); - runWithBackfillFirst(options, connectionString); + if (!Boolean.TRUE.equals(options.getUseShadowTables())) { + LOG.info("Using high-throughput shadowless processing mode"); + runShadowless(options, connectionString); + } else if (Boolean.TRUE.equals(options.getProcessBackfillFirst())) { + LOG.info("Using legacy backfill-first processing mode with shadow tables"); + runLegacyWithBackfillFirst(options, connectionString); } else { - LOG.info("Using unified processing mode"); - runAllEventsTogether(options, connectionString); + LOG.info("Using legacy unified processing mode with shadow tables"); + runLegacyAllEventsTogether(options, connectionString); } - - mongoClient.close(); - LOG.info("MongoDB client closed"); } catch (Exception e) { LOG.error("Failed to run pipeline: {}", e.getMessage(), e); throw e; @@ -503,32 +560,259 @@ public static void run(Options options) { } /** - * Runs the pipeline with backfill events processed before CDC events. - * - *

This pipeline first processes all backfill events, then processes CDC events. This ensures - * that the initial state of the database is established before any changes are applied. Failures - * in backfill will be sent over to dlq and be processed with conflict resolving. + * Runs the pipeline in high-throughput shadowless mode with hierarchical stages. * * @param options The execution parameters to the pipeline. + * @param connectionString The MongoDB/Firestore connection URI. * @return The result of the pipeline execution. */ - private static PipelineResult runWithBackfillFirst(Options options, String connectionString) { + public static PipelineResult runShadowless(Options options, String connectionString) { + LOG.info("Creating shadowless pipeline DAG"); + Pipeline pipeline = Pipeline.create(options); + pipeline + .getCoderRegistry() + .registerCoderForClass(TimestampSortKey.class, TimestampSortKeyCoder.of()); + pipeline + .getCoderRegistry() + .registerCoderForClass( + MongoDbChangeEventContext.class, MongoDbChangeEventContextCoder.of()); + DeadLetterQueueManager dlqManager = buildDlqManager(options); + + /* + * Stage 1: Read/ + * - Read/DataStreamIO + * - Read/IngestAndNormalizeJson + * - Read/MergeWithReconsumedDlq + */ + LOG.info("Setting up Read/ stage"); + PCollection> jsonRecords = + ingestAndNormalizeJsonShadowless(options, dlqManager, pipeline) + .setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())); + + /* + * Stage 2: Process/ + * - Process/ApplyUdfToDataField (optional) + * - Process/CreateMongoDbChangeEventContext + * - Process/KeyByCollectionAndDocId + * - Process/GlobalWindows + * - Process/StatefulDeduplication + */ + LOG.info("Setting up Process/ stage"); + if (!Strings.isNullOrEmpty(options.getJavascriptTextTransformGcsPath())) { + LOG.info("Applying Javascript UDF in Process/ApplyUdfToDataField"); + jsonRecords = + jsonRecords.apply( + "Process/ApplyUdfToDataField", new ApplyUdfToDataField(options, dlqManager)); + } + + PCollectionTuple changeEventContexts = + jsonRecords.apply( + "Process/CreateMongoDbChangeEventContext", + ParDo.of(new CreateMongoDbChangeEventContextFn(options.getShadowCollectionPrefix())) + .withOutputTags( + CreateMongoDbChangeEventContextFn.SUCCESSFUL_CREATION_TAG, + TupleTagList.of(CreateMongoDbChangeEventContextFn.FAILED_CREATION_TAG))); + + changeEventContexts + .get(CreateMongoDbChangeEventContextFn.SUCCESSFUL_CREATION_TAG) + .setCoder(MongoDbChangeEventContextCoder.of()); + changeEventContexts + .get(CreateMongoDbChangeEventContextFn.FAILED_CREATION_TAG) + .setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())); + + writeFailedJsonToDlq( + options, + changeEventContexts, + dlqManager, + CreateMongoDbChangeEventContextFn.FAILED_CREATION_TAG, + "Process/WriteFailedContextCreationToDlq"); + + PCollection contexts = + changeEventContexts.get(CreateMongoDbChangeEventContextFn.SUCCESSFUL_CREATION_TAG); + + LOG.info("Configuring shadowless stateful deduplication by collection and doc ID"); + PCollection> keyedEvents = + contexts.apply( + "Process/KeyByCollectionAndDocId", + WithKeys.of( + (MongoDbChangeEventContext event) -> + event.getDataCollection() + + "#" + + Utils.documentIdToString(event.getDocumentId())) + .withKeyType(TypeDescriptors.strings())); + + PCollection dedupedEvents = + keyedEvents + .apply( + "Process/GlobalWindows", + Window.>into(new GlobalWindows())) + .apply("Process/StatefulDeduplication", ParDo.of(new StatefulDeduplicationFn())); + /* - * Stages: - * 1) Ingest and Normalize Data to FailsafeElement with JSON Strings - * 2) Convert json strings to MongoDbChangeEventContext - * 3) Split the MongoDbChangeEventContext into backfill and cdc events - * 4) Process backfill events with bulk writes, failed backfill will be sent to dlq and later processed with transactions and conflict resolving. - * 5) Process the cdc events with transactions + * Stage 3: Write/ + * - Write/AsyncBulkWriteToFirestore (MongoDbBulkTransforms.BulkWriteWithDlq) + * - Write/WriteToDlq_Retryable + * - Write/WriteToDlq_Severe */ + LOG.info("Setting up Write/ stage"); + PCollectionTuple writeResult = + dedupedEvents.apply( + "Write/AsyncBulkWriteToFirestore", + MongoDbBulkTransforms.bulkWriteWithDlq() + .withUri(connectionString) + .withDatabase(options.getDatabaseName()) + .withBatchSize(options.getBatchSize()) + .withMaxConcurrentAsyncWrites(options.getMaxConcurrentAsyncWrites()) + .withInitialWriteRatePerWorker(options.getInitialWriteRatePerWorker()) + .withWriteRateRampUpMinutes(options.getWriteRateRampUpMinutes()) + .withMaxWriteRatePerWorker(options.getMaxWriteRatePerWorker())); + + writeResult + .get(MongoDbBulkTransforms.SUCCESSFUL_WRITE_TAG) + .setCoder(MongoDbChangeEventContextCoder.of()); + writeResult + .get(MongoDbBulkTransforms.FAILED_WRITE_TAG) + .setCoder( + FailsafeElementCoder.of( + MongoDbChangeEventContextCoder.of(), MongoDbChangeEventContextCoder.of())); + writeResult + .get(MongoDbBulkTransforms.SEVERE_FAILED_WRITE_TAG) + .setCoder( + FailsafeElementCoder.of( + MongoDbChangeEventContextCoder.of(), MongoDbChangeEventContextCoder.of())); + + writeFailedEventsToDlq( + options, + writeResult, + dlqManager, + MongoDbBulkTransforms.FAILED_WRITE_TAG, + "Write/WriteToDlq_Retryable"); + + writeSevereEventsToDlq( + options, + writeResult, + dlqManager, + MongoDbBulkTransforms.SEVERE_FAILED_WRITE_TAG, + "Write/WriteToDlq_Severe"); + + LOG.info("Executing shadowless pipeline"); + return pipeline.run(); + } + + /** Read from input path and dlq to collect objects to process without reshuffle. */ + private static PCollection> ingestAndNormalizeJsonShadowless( + Options options, DeadLetterQueueManager dlqManager, Pipeline pipeline) { + LOG.info("Starting Read/ ingestion for shadowless mode"); + boolean isRegularMode = "regular".equals(options.getRunMode()); + PCollectionTuple reconsumedElements = + pipeline.apply("Read/PollAndReconsumeDLQ", new ReconsumeDlqTransform(options, dlqManager)); + + PCollection> dlqJsonRecords = + reconsumedElements + .get(DeadLetterQueueManager.RETRYABLE_ERRORS) + .setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())) + .apply( + "Read/Count DLQ Retries", + ParDo.of( + new DoFn, FailsafeElement>() { + private final Counter dlqRetries = + Metrics.counter(DataStreamMongoDBToFirestore.class, "dlqRetries"); + + @ProcessElement + public void processElement(ProcessContext c) { + dlqRetries.inc(); + c.output(c.element()); + } + })); + + reconsumedElements + .get(DeadLetterQueueManager.PERMANENT_ERRORS) + .setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())) + .apply( + "Read/Count Permanent Failures", + ParDo.of( + new DoFn, FailsafeElement>() { + private final Counter permanentFailures = + Metrics.counter(DataStreamMongoDBToFirestore.class, "permanentFailures"); + + @ProcessElement + public void processElement(ProcessContext c) { + permanentFailures.inc(); + c.output(c.element()); + } + })) + .apply( + "Read/Write Permanent Failures To DLQ - Sanitize", + MapElements.via(new StringDeadLetterQueueSanitizer())) + .setCoder(StringUtf8Coder.of()) + .apply( + "Read/Write Permanent Failures To DLQ", + DLQWriteTransform.WriteDLQ.newBuilder() + .withDlqDirectory(dlqManager.getSevereDlqDirectoryWithDateTime()) + .withTmpDirectory(dlqManager.getSevereDlqDirectory() + "tmp_severe/") + .setIncludePaneInfo(true) + .build()); + + if (isRegularMode) { + PCollection> datastreamJsonRecords = + pipeline.apply( + "Read/DataStreamIO", + new DataStreamIO( + options.getStreamName(), + options.getInputFilePattern(), + options.getInputFileFormat(), + options.getGcsPubSubSubscription(), + options.getRfcStartDateTime()) + .withFileReadConcurrency(options.getFileReadConcurrency()) + .withoutDatastreamRecordsReshuffle() + .withDirectoryWatchDuration( + Duration.standardMinutes(options.getDirectoryWatchDurationInMinutes()))); + + return PCollectionList.of(datastreamJsonRecords) + .and(dlqJsonRecords) + .apply("Read/MergeWithReconsumedDlq", Flatten.pCollections()) + .setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())); + } else { + return PCollectionList.of(dlqJsonRecords) + .apply("Read/MergeWithReconsumedDlq", Flatten.pCollections()) + .setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())); + } + } + + /** Legacy execution path: backfill events processed before CDC events with shadow tables. */ + public static PipelineResult runLegacyWithBackfillFirst( + Options options, String connectionString) { + return runWithBackfillFirst(options, connectionString); + } + + /** Legacy execution path: all events processed together with shadow tables. */ + public static PipelineResult runLegacyAllEventsTogether( + Options options, String connectionString) { + return runAllEventsTogether(options, connectionString); + } + + /** + * Runs the pipeline with backfill events processed before CDC events. + * + * @param options The execution parameters to the pipeline. + * @return The result of the pipeline execution. + */ + private static PipelineResult runWithBackfillFirst(Options options, String connectionString) { LOG.info("Creating pipeline with backfill-first processing"); Pipeline pipeline = Pipeline.create(options); + pipeline + .getCoderRegistry() + .registerCoderForClass(TimestampSortKey.class, TimestampSortKeyCoder.of()); + pipeline + .getCoderRegistry() + .registerCoderForClass( + MongoDbChangeEventContext.class, MongoDbChangeEventContextCoder.of()); LOG.info("Building Dead Letter Queue manager"); DeadLetterQueueManager dlqManager = buildDlqManager(options); // Stage 1: Ingest data from GCS - LOG.info("Ingesting data from GCS"); + LOG.info("Configuring data ingestion from GCS"); PCollection> jsonRecords = ingestAndNormalizeJson(options, dlqManager, pipeline) .setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())); @@ -542,21 +826,21 @@ private static PipelineResult runWithBackfillFirst(Options options, String conne } // Stage 2: Create MongoDbChangeEventContext objects - LOG.info("Creating MongoDbChangeEventContext objects"); + LOG.info("Configuring MongoDbChangeEventContext creation"); PCollectionTuple changeEventContexts = jsonRecords.apply( "Create MongoDbChangeEventContext objects", ParDo.of(new CreateMongoDbChangeEventContextFn(options.getShadowCollectionPrefix())) .withOutputTags( - CreateMongoDbChangeEventContextFn.successfulCreationTag, - TupleTagList.of(CreateMongoDbChangeEventContextFn.failedCreationTag))); + CreateMongoDbChangeEventContextFn.SUCCESSFUL_CREATION_TAG, + TupleTagList.of(CreateMongoDbChangeEventContextFn.FAILED_CREATION_TAG))); // Set coders changeEventContexts - .get(CreateMongoDbChangeEventContextFn.successfulCreationTag) - .setCoder(SerializableCoder.of(MongoDbChangeEventContext.class)); + .get(CreateMongoDbChangeEventContextFn.SUCCESSFUL_CREATION_TAG) + .setCoder(MongoDbChangeEventContextCoder.of()); changeEventContexts - .get(CreateMongoDbChangeEventContextFn.failedCreationTag) + .get(CreateMongoDbChangeEventContextFn.FAILED_CREATION_TAG) .setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())); // Handle failed creation with DLQ @@ -564,85 +848,83 @@ private static PipelineResult runWithBackfillFirst(Options options, String conne options, changeEventContexts, dlqManager, - CreateMongoDbChangeEventContextFn.failedCreationTag); + CreateMongoDbChangeEventContextFn.FAILED_CREATION_TAG); // Stage 3: Split events into backfill and CDC streams - LOG.info("Splitting events into backfill and CDC streams"); + LOG.info("Configuring event splitting into backfill and CDC streams"); PCollectionTuple splitEvents = changeEventContexts - .get(CreateMongoDbChangeEventContextFn.successfulCreationTag) + .get(CreateMongoDbChangeEventContextFn.SUCCESSFUL_CREATION_TAG) .apply( "Split Backfill and CDC", ParDo.of(new SplitBackfillAndCdcEventsFn()) .withOutputTags( - SplitBackfillAndCdcEventsFn.backfillTag, - TupleTagList.of(SplitBackfillAndCdcEventsFn.cdcTag))); + SplitBackfillAndCdcEventsFn.BACKFILL_TAG, + TupleTagList.of(SplitBackfillAndCdcEventsFn.CDC_TAG))); // Set coders for split events splitEvents - .get(SplitBackfillAndCdcEventsFn.backfillTag) - .setCoder(SerializableCoder.of(MongoDbChangeEventContext.class)); + .get(SplitBackfillAndCdcEventsFn.BACKFILL_TAG) + .setCoder(MongoDbChangeEventContextCoder.of()); splitEvents - .get(SplitBackfillAndCdcEventsFn.cdcTag) - .setCoder(SerializableCoder.of(MongoDbChangeEventContext.class)); + .get(SplitBackfillAndCdcEventsFn.CDC_TAG) + .setCoder(MongoDbChangeEventContextCoder.of()); // Stage 4: Process backfill events - LOG.info("Processing backfill events"); + LOG.info("Configuring backfill event processing"); PCollectionTuple backfillResult; if (options.getUseShadowTablesForBackfill()) { // Use shadow tables for backfill (same as CDC processing) backfillResult = splitEvents - .get(SplitBackfillAndCdcEventsFn.backfillTag) + .get(SplitBackfillAndCdcEventsFn.BACKFILL_TAG) .apply( "Process Backfill with Shadow Tables", ParDo.of(new ProcessChangeEventFn(connectionString, options.getDatabaseName())) .withOutputTags( - ProcessChangeEventFn.successfulWriteTag, - TupleTagList.of(ProcessChangeEventFn.failedWriteTag) - .and(ProcessChangeEventFn.severeFailedWriteTag))); + ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG, + TupleTagList.of(ProcessChangeEventFn.FAILED_WRITE_TAG) + .and(ProcessChangeEventFn.SEVERE_FAILED_WRITE_TAG))); } else { // Process backfill without shadow tables backfillResult = splitEvents - .get(SplitBackfillAndCdcEventsFn.backfillTag) + .get(SplitBackfillAndCdcEventsFn.BACKFILL_TAG) .apply( "Process Backfill without Shadow Tables", ParDo.of( new ProcessBackfillEventFn( connectionString, options.getDatabaseName(), options.getBatchSize())) .withOutputTags( - ProcessBackfillEventFn.successfulWriteTag, - TupleTagList.of(ProcessBackfillEventFn.failedWriteTag) - .and(ProcessBackfillEventFn.severeFailedWriteTag))); + ProcessBackfillEventFn.SUCCESSFUL_WRITE_TAG, + TupleTagList.of(ProcessBackfillEventFn.FAILED_WRITE_TAG) + .and(ProcessBackfillEventFn.SEVERE_FAILED_WRITE_TAG))); } // Set coders for backfill results backfillResult .get( options.getUseShadowTablesForBackfill() - ? ProcessChangeEventFn.successfulWriteTag - : ProcessBackfillEventFn.successfulWriteTag) - .setCoder(SerializableCoder.of(MongoDbChangeEventContext.class)); + ? ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG + : ProcessBackfillEventFn.SUCCESSFUL_WRITE_TAG) + .setCoder(MongoDbChangeEventContextCoder.of()); backfillResult .get( options.getUseShadowTablesForBackfill() - ? ProcessChangeEventFn.failedWriteTag - : ProcessBackfillEventFn.failedWriteTag) + ? ProcessChangeEventFn.FAILED_WRITE_TAG + : ProcessBackfillEventFn.FAILED_WRITE_TAG) .setCoder( FailsafeElementCoder.of( - SerializableCoder.of(MongoDbChangeEventContext.class), - SerializableCoder.of(MongoDbChangeEventContext.class))); + MongoDbChangeEventContextCoder.of(), MongoDbChangeEventContextCoder.of())); backfillResult .get( options.getUseShadowTablesForBackfill() - ? ProcessChangeEventFn.severeFailedWriteTag - : ProcessBackfillEventFn.severeFailedWriteTag) + ? ProcessChangeEventFn.SEVERE_FAILED_WRITE_TAG + : ProcessBackfillEventFn.SEVERE_FAILED_WRITE_TAG) .setCoder( FailsafeElementCoder.of( - SerializableCoder.of(MongoDbChangeEventContext.class), - SerializableCoder.of(MongoDbChangeEventContext.class))); + MongoDbChangeEventContextCoder.of(), MongoDbChangeEventContextCoder.of())); // Handle failed backfill writes with DLQ writeFailedEventsToDlq( @@ -650,8 +932,8 @@ private static PipelineResult runWithBackfillFirst(Options options, String conne backfillResult, dlqManager, options.getUseShadowTablesForBackfill() - ? ProcessChangeEventFn.failedWriteTag - : ProcessBackfillEventFn.failedWriteTag); + ? ProcessChangeEventFn.FAILED_WRITE_TAG + : ProcessBackfillEventFn.FAILED_WRITE_TAG); // Write severe backfill failures directly to severe DLQ writeSevereEventsToDlq( @@ -659,45 +941,43 @@ private static PipelineResult runWithBackfillFirst(Options options, String conne backfillResult, dlqManager, options.getUseShadowTablesForBackfill() - ? ProcessChangeEventFn.severeFailedWriteTag - : ProcessBackfillEventFn.severeFailedWriteTag); + ? ProcessChangeEventFn.SEVERE_FAILED_WRITE_TAG + : ProcessBackfillEventFn.SEVERE_FAILED_WRITE_TAG); // Stage 5: Process CDC events - LOG.info("Processing CDC events"); + LOG.info("Configuring CDC event processing"); PCollectionTuple cdcResult = splitEvents - .get(SplitBackfillAndCdcEventsFn.cdcTag) + .get(SplitBackfillAndCdcEventsFn.CDC_TAG) .apply( "Process CDC Events", ParDo.of(new ProcessChangeEventFn(connectionString, options.getDatabaseName())) .withOutputTags( - ProcessChangeEventFn.successfulWriteTag, - TupleTagList.of(ProcessChangeEventFn.failedWriteTag) - .and(ProcessChangeEventFn.severeFailedWriteTag))); + ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG, + TupleTagList.of(ProcessChangeEventFn.FAILED_WRITE_TAG) + .and(ProcessChangeEventFn.SEVERE_FAILED_WRITE_TAG))); // Set coders for CDC results cdcResult - .get(ProcessChangeEventFn.successfulWriteTag) - .setCoder(SerializableCoder.of(MongoDbChangeEventContext.class)); + .get(ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG) + .setCoder(MongoDbChangeEventContextCoder.of()); cdcResult - .get(ProcessChangeEventFn.failedWriteTag) + .get(ProcessChangeEventFn.FAILED_WRITE_TAG) .setCoder( FailsafeElementCoder.of( - SerializableCoder.of(MongoDbChangeEventContext.class), - SerializableCoder.of(MongoDbChangeEventContext.class))); + MongoDbChangeEventContextCoder.of(), MongoDbChangeEventContextCoder.of())); cdcResult - .get(ProcessChangeEventFn.severeFailedWriteTag) + .get(ProcessChangeEventFn.SEVERE_FAILED_WRITE_TAG) .setCoder( FailsafeElementCoder.of( - SerializableCoder.of(MongoDbChangeEventContext.class), - SerializableCoder.of(MongoDbChangeEventContext.class))); + MongoDbChangeEventContextCoder.of(), MongoDbChangeEventContextCoder.of())); // Handle failed CDC writes with DLQ - writeFailedEventsToDlq(options, cdcResult, dlqManager, ProcessChangeEventFn.failedWriteTag); + writeFailedEventsToDlq(options, cdcResult, dlqManager, ProcessChangeEventFn.FAILED_WRITE_TAG); // Write severe CDC failures directly to severe DLQ writeSevereEventsToDlq( - options, cdcResult, dlqManager, ProcessChangeEventFn.severeFailedWriteTag); + options, cdcResult, dlqManager, ProcessChangeEventFn.SEVERE_FAILED_WRITE_TAG); // Execute the pipeline LOG.info("Executing pipeline"); @@ -705,122 +985,94 @@ private static PipelineResult runWithBackfillFirst(Options options, String conne } /** - * Runs the pipeline with all events processed together. - * - *

This pipeline processes both backfill and CDC events in a unified flow, ordered primarily by - * their timestamps. Events with the same timestamp are ordered by type, with backfill events - * processed before CDC events. Shadow collections are used to track event ordering and prevent - * duplicate processing. + * Runs the pipeline with all events processed together using shadow tables. * * @param options The execution parameters to the pipeline. * @return The result of the pipeline execution. */ private static PipelineResult runAllEventsTogether(Options options, String connectionString) { - /* - * Stages: - * 1) Ingest and Normalize Data to FailsafeElement with JSON Strings - * 2) Convert json strings to MongoDbChangeEvents - * 3) Write the change events with transactions - */ - LOG.info("Creating pipeline"); Pipeline pipeline = Pipeline.create(options); + pipeline + .getCoderRegistry() + .registerCoderForClass(TimestampSortKey.class, TimestampSortKeyCoder.of()); + pipeline + .getCoderRegistry() + .registerCoderForClass( + MongoDbChangeEventContext.class, MongoDbChangeEventContextCoder.of()); LOG.info("Building Dead Letter Queue manager"); DeadLetterQueueManager dlqManager = buildDlqManager(options); - /* - * Stage 1: Ingest and Normalize Data to FailsafeElement with JSON Strings - * a) Read DataStream data from GCS into JSON String FailsafeElements (datastreamJsonRecords) - */ - LOG.info("Stage 1: Starting ingestion of data from GCS"); + LOG.info("Stage 1: Configuring data ingestion from GCS"); PCollection> jsonRecords = ingestAndNormalizeJson(options, dlqManager, pipeline) .setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())); - /* - * Optional Stage 1.5: Apply Javascript UDF to transform JSON strings - */ if (!Strings.isNullOrEmpty(options.getJavascriptTextTransformGcsPath())) { jsonRecords = jsonRecords.apply( "Apply UDF To Data Field", new ApplyUdfToDataField(options, dlqManager)); } - LOG.info("Stage 1: Completed ingestion of data from GCS"); - - /* - * Stage 2: Create MongoDbChangeEventContext objects with error handling - */ - LOG.info("Stage 2: Creating MongoDbChangeEventContext objects"); + LOG.info("Stage 2: Configuring MongoDbChangeEventContext creation"); PCollectionTuple changeEventContexts = jsonRecords.apply( "Create MongoDbChangeEventContext objects", ParDo.of(new CreateMongoDbChangeEventContextFn(options.getShadowCollectionPrefix())) .withOutputTags( - CreateMongoDbChangeEventContextFn.successfulCreationTag, - TupleTagList.of(CreateMongoDbChangeEventContextFn.failedCreationTag))); + CreateMongoDbChangeEventContextFn.SUCCESSFUL_CREATION_TAG, + TupleTagList.of(CreateMongoDbChangeEventContextFn.FAILED_CREATION_TAG))); - /* Set coder for successful creation */ changeEventContexts - .get(CreateMongoDbChangeEventContextFn.successfulCreationTag) - .setCoder(SerializableCoder.of(MongoDbChangeEventContext.class)); + .get(CreateMongoDbChangeEventContextFn.SUCCESSFUL_CREATION_TAG) + .setCoder(MongoDbChangeEventContextCoder.of()); - /* Set coder for failed creation */ changeEventContexts - .get(CreateMongoDbChangeEventContextFn.failedCreationTag) + .get(CreateMongoDbChangeEventContextFn.FAILED_CREATION_TAG) .setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())); - // Handle failed creation with DLQ LOG.info("Setting up DLQ handling for failed event creation"); writeFailedJsonToDlq( options, changeEventContexts, dlqManager, - CreateMongoDbChangeEventContextFn.failedCreationTag); + CreateMongoDbChangeEventContextFn.FAILED_CREATION_TAG); - /* Stage 3: Iterate through the success events and write with transactions */ - LOG.info("Stage 3: Processing change events and writing to the destination database"); + LOG.info("Stage 3: Configuring change event processing and destination database writing"); PCollectionTuple writeResult = changeEventContexts - .get(CreateMongoDbChangeEventContextFn.successfulCreationTag) - .setCoder(SerializableCoder.of(MongoDbChangeEventContext.class)) + .get(CreateMongoDbChangeEventContextFn.SUCCESSFUL_CREATION_TAG) + .setCoder(MongoDbChangeEventContextCoder.of()) .apply( "Transactional write events", ParDo.of(new ProcessChangeEventFn(connectionString, options.getDatabaseName())) .withOutputTags( - ProcessChangeEventFn.successfulWriteTag, - TupleTagList.of(ProcessChangeEventFn.failedWriteTag) - .and(ProcessChangeEventFn.severeFailedWriteTag))); + ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG, + TupleTagList.of(ProcessChangeEventFn.FAILED_WRITE_TAG) + .and(ProcessChangeEventFn.SEVERE_FAILED_WRITE_TAG))); - /* Set coder for successful writes */ writeResult - .get(ProcessChangeEventFn.successfulWriteTag) - .setCoder(SerializableCoder.of(MongoDbChangeEventContext.class)); + .get(ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG) + .setCoder(MongoDbChangeEventContextCoder.of()); - /* Set coder for failed writes */ writeResult - .get(ProcessChangeEventFn.failedWriteTag) + .get(ProcessChangeEventFn.FAILED_WRITE_TAG) .setCoder( FailsafeElementCoder.of( - SerializableCoder.of(MongoDbChangeEventContext.class), - SerializableCoder.of(MongoDbChangeEventContext.class))); + MongoDbChangeEventContextCoder.of(), MongoDbChangeEventContextCoder.of())); writeResult - .get(ProcessChangeEventFn.severeFailedWriteTag) + .get(ProcessChangeEventFn.SEVERE_FAILED_WRITE_TAG) .setCoder( FailsafeElementCoder.of( - SerializableCoder.of(MongoDbChangeEventContext.class), - SerializableCoder.of(MongoDbChangeEventContext.class))); + MongoDbChangeEventContextCoder.of(), MongoDbChangeEventContextCoder.of())); - /* Handle failed writes with DLQ */ LOG.info("Setting up DLQ handling for failed writes"); - writeFailedEventsToDlq(options, writeResult, dlqManager, ProcessChangeEventFn.failedWriteTag); - // Write severe failures directly to severe DLQ + writeFailedEventsToDlq(options, writeResult, dlqManager, ProcessChangeEventFn.FAILED_WRITE_TAG); writeSevereEventsToDlq( - options, writeResult, dlqManager, ProcessChangeEventFn.severeFailedWriteTag); + options, writeResult, dlqManager, ProcessChangeEventFn.SEVERE_FAILED_WRITE_TAG); - // Execute the pipeline and return the result. LOG.info("Executing pipeline"); return pipeline.run(); } @@ -834,18 +1086,16 @@ private static DeadLetterQueueManager buildDlqManager(Options options) { tempLocation = tempLocation.endsWith("/") ? tempLocation : tempLocation + "/"; LOG.info("Using temp location from pipeline options: {}", tempLocation); } else { - // If tempLocation is null, use a default location tempLocation = "/tmp/"; LOG.warn("TempLocation is null, using default location: {}", tempLocation); } } catch (Exception e) { - // If we can't get the temp location, use a default tempLocation = "/tmp/"; LOG.warn("Error getting tempLocation, using default location: {}", tempLocation, e); } String dlqDirectory = - options.getDeadLetterQueueDirectory().isEmpty() + Strings.isNullOrEmpty(options.getDeadLetterQueueDirectory()) ? tempLocation + "dlq/" : options.getDeadLetterQueueDirectory(); LOG.info("Dead-letter queue directory: {}", dlqDirectory); @@ -867,38 +1117,14 @@ private static DeadLetterQueueManager buildDlqManager(Options options) { /** Read from input path and dlq to collect objects to process. */ private static PCollection> ingestAndNormalizeJson( Options options, DeadLetterQueueManager dlqManager, Pipeline pipeline) { - LOG.info("Starting ingestion and normalization of JSON data"); + LOG.info("Configuring ingestion and normalization of JSON data"); PCollection> jsonRecords; - // Elements sent to the Dead Letter Queue are to be reconsumed. - // A DLQManager is to be created using PipelineOptions, and it is in charge - // of building pieces of the DLQ. PCollectionTuple reconsumedElements; boolean isRegularMode = "regular".equals(options.getRunMode()); LOG.info("Setting up DLQ reconsumption, mode: {}", isRegularMode ? "regular" : "retry"); - if (isRegularMode && (!Strings.isNullOrEmpty(options.getDlqGcsPubSubSubscription()))) { - LOG.info( - "Using PubSub notification for DLQ reconsumption with subscription: {}", - options.getDlqGcsPubSubSubscription()); - reconsumedElements = - dlqManager.getReconsumerDataTransformForFiles( - pipeline.apply( - "Read retry from PubSub", - new PubSubNotifiedDlqIO( - options.getDlqGcsPubSubSubscription(), - // file paths to ignore when re-consuming for retry - new ArrayList( - Arrays.asList("/severe/", "/tmp_retry", "/tmp_severe/", ".temp"))))); - } else { - LOG.info( - "Using periodic polling for DLQ reconsumption with retry minutes: {}", - options.getDlqRetryMinutes()); - reconsumedElements = - dlqManager.getReconsumerDataTransform( - pipeline.apply( - "Periodically polling from DLQ", - dlqManager.dlqReconsumer(options.getDlqRetryMinutes()))); - } + reconsumedElements = + pipeline.apply("PollAndReconsumeDLQ", new ReconsumeDlqTransform(options, dlqManager)); LOG.info("Processing retryable errors from DLQ"); PCollection> dlqJsonRecords = @@ -987,7 +1213,6 @@ public void processElement(ProcessContext c) { .setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())) .apply("Reshuffle records", Reshuffle.viaRandomKey()); } - LOG.info("Completed ingestion and normalization of JSON data"); return jsonRecords; } @@ -996,17 +1221,24 @@ private static void writeFailedJsonToDlq( PCollectionTuple results, DeadLetterQueueManager dlqManager, TupleTag> failedTag) { - LOG.info("Setting up DLQ for failed JSON processing"); - // Write failed writes to DLQ + writeFailedJsonToDlq( + options, results, dlqManager, failedTag, "Write Failed Json To DLQ - " + failedTag.getId()); + } + + private static void writeFailedJsonToDlq( + Options options, + PCollectionTuple results, + DeadLetterQueueManager dlqManager, + TupleTag> failedTag, + String stageName) { + LOG.info("Setting up DLQ for failed JSON processing: {}", stageName); results .get(failedTag) .setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())) - .apply( - "DLQ: Write Retryable Json Failures to GCS - " + failedTag.getId(), - MapElements.via(new StringDeadLetterQueueSanitizer())) + .apply(stageName + " - Sanitize", MapElements.via(new StringDeadLetterQueueSanitizer())) .setCoder(StringUtf8Coder.of()) .apply( - "Write Failed Json To DLQ - " + failedTag.getId(), + stageName, DLQWriteTransform.WriteDLQ.newBuilder() .withDlqDirectory(dlqManager.getSevereDlqDirectoryWithDateTime()) .withTmpDirectory(options.getDeadLetterQueueDirectory() + "/tmp_non_retry_json/") @@ -1020,20 +1252,26 @@ private static void writeFailedEventsToDlq( PCollectionTuple results, DeadLetterQueueManager dlqManager, TupleTag> failedTag) { - LOG.info("Setting up DLQ for failed MongoDB event processing"); - // Write failed writes to DLQ + writeFailedEventsToDlq(options, results, dlqManager, failedTag, "Write Events Failures To DLQ"); + } + + private static void writeFailedEventsToDlq( + Options options, + PCollectionTuple results, + DeadLetterQueueManager dlqManager, + TupleTag> failedTag, + String stageName) { + LOG.info("Setting up DLQ for failed MongoDB event processing: {}", stageName); results .get(failedTag) .setCoder( FailsafeElementCoder.of( - SerializableCoder.of(MongoDbChangeEventContext.class), - SerializableCoder.of(MongoDbChangeEventContext.class))) + MongoDbChangeEventContextCoder.of(), MongoDbChangeEventContextCoder.of())) .apply( - "DLQ: Write Retryable Events Failures to GCS", - MapElements.via(new MongoDbEventDeadLetterQueueSanitizer())) + stageName + " - Sanitize", MapElements.via(new MongoDbEventDeadLetterQueueSanitizer())) .setCoder(StringUtf8Coder.of()) .apply( - "Write Events Failures To DLQ", + stageName, DLQWriteTransform.WriteDLQ.newBuilder() .withDlqDirectory(dlqManager.getRetryDlqDirectoryWithDateTime()) .withTmpDirectory(options.getDeadLetterQueueDirectory() + "/tmp_retry_mongo_event/") @@ -1047,19 +1285,27 @@ private static void writeSevereEventsToDlq( PCollectionTuple results, DeadLetterQueueManager dlqManager, TupleTag> failedTag) { - LOG.info("Setting up Severe DLQ for failed MongoDB event processing"); + writeSevereEventsToDlq( + options, results, dlqManager, failedTag, "Write Severe Events Failures To DLQ"); + } + + private static void writeSevereEventsToDlq( + Options options, + PCollectionTuple results, + DeadLetterQueueManager dlqManager, + TupleTag> failedTag, + String stageName) { + LOG.info("Setting up Severe DLQ for failed MongoDB event processing: {}", stageName); results .get(failedTag) .setCoder( FailsafeElementCoder.of( - SerializableCoder.of(MongoDbChangeEventContext.class), - SerializableCoder.of(MongoDbChangeEventContext.class))) + MongoDbChangeEventContextCoder.of(), MongoDbChangeEventContextCoder.of())) .apply( - "DLQ: Write Severe Events Failures to GCS", - MapElements.via(new MongoDbEventDeadLetterQueueSanitizer())) + stageName + " - Sanitize", MapElements.via(new MongoDbEventDeadLetterQueueSanitizer())) .setCoder(StringUtf8Coder.of()) .apply( - "Write Severe Events Failures To DLQ", + stageName, DLQWriteTransform.WriteDLQ.newBuilder() .withDlqDirectory(dlqManager.getSevereDlqDirectoryWithDateTime()) .withTmpDirectory( @@ -1075,8 +1321,9 @@ public static class SplitBackfillAndCdcEventsFn private static final Logger LOG = LoggerFactory.getLogger(SplitBackfillAndCdcEventsFn.class); - public static TupleTag backfillTag = new TupleTag<>("backfill"); - public static TupleTag cdcTag = new TupleTag<>("cdc"); + public static final TupleTag BACKFILL_TAG = + new TupleTag<>("backfill"); + public static final TupleTag CDC_TAG = new TupleTag<>("cdc"); @ProcessElement public void processElement(ProcessContext c, MultiOutputReceiver out) { @@ -1084,10 +1331,10 @@ public void processElement(ProcessContext c, MultiOutputReceiver out) { if (isNonDlqBackfillEvent(event)) { LOG.debug("Classified event as backfill for document ID: {}", event.getDocumentId()); - out.get(backfillTag).output(event); + out.get(BACKFILL_TAG).output(event); } else { LOG.debug("Classified event as CDC for document ID: {}", event.getDocumentId()); - out.get(cdcTag).output(event); + out.get(CDC_TAG).output(event); } } @@ -1097,7 +1344,6 @@ private boolean isNonDlqBackfillEvent(MongoDbChangeEventContext event) { } JsonNode jsonNode = event.getChangeEvent(); - // Check for CDC-specific fields boolean hasCdcFields = jsonNode.has("_metadata_log_file") || jsonNode.has("_metadata_log_position") @@ -1105,13 +1351,11 @@ private boolean isNonDlqBackfillEvent(MongoDbChangeEventContext event) { || jsonNode.has("_metadata_ssn") || jsonNode.has("_metadata_rs_id"); - // Check for change type String changeType = null; if (jsonNode.has(DatastreamConstants.EVENT_CHANGE_TYPE_KEY)) { changeType = jsonNode.get(DatastreamConstants.EVENT_CHANGE_TYPE_KEY).asText(); } - // If it has CDC fields or a specific change type (not READ), it's a CDC event return !hasCdcFields && (changeType == null || "READ".equals(changeType)); } } @@ -1122,18 +1366,19 @@ public static class ProcessBackfillEventFn private static final Logger LOG = LoggerFactory.getLogger(ProcessBackfillEventFn.class); - public static TupleTag successfulWriteTag = + public static final TupleTag SUCCESSFUL_WRITE_TAG = new TupleTag<>("backfillSuccessfulWrite"); - public static TupleTag> - failedWriteTag = new TupleTag<>("backfillFailedWrite"); - public static TupleTag> - severeFailedWriteTag = new TupleTag<>("backfillSevereFailedWrite"); + public static final TupleTag< + FailsafeElement> + FAILED_WRITE_TAG = new TupleTag<>("backfillFailedWrite"); + public static final TupleTag< + FailsafeElement> + SEVERE_FAILED_WRITE_TAG = new TupleTag<>("backfillSevereFailedWrite"); private final String connectionString; private final String targetDatabaseName; private final int batchSize; - // Maps to store buffered operations by collection private transient Map> bufferedEvents; private transient Map> collectionMap; private transient MongoClient client; @@ -1169,9 +1414,8 @@ public void setup() { .applyConnectionString(new com.mongodb.ConnectionString(connectionString)) .applyToSocketSettings( builder -> { - // How long the driver will wait to establish a connection builder.connectTimeout(60, TimeUnit.SECONDS); - builder.readTimeout(60, TimeUnit.SECONDS); // Example: 60 seconds + builder.readTimeout(60, TimeUnit.SECONDS); }) .applyToClusterSettings( builder -> builder.serverSelectionTimeout(10, TimeUnit.MINUTES)) @@ -1194,12 +1438,10 @@ public void processElement(ProcessContext context, MultiOutputReceiver out) { MongoDbChangeEventContext element = context.element(); String collectionName = element.getDataCollection(); - // Buffer the event if (!bufferedEvents.containsKey(collectionName)) { LOG.debug("Creating new buffer for collection: {}", collectionName); bufferedEvents.put(collectionName, new ArrayList<>()); - // Initialize collection reference if needed if (!collectionMap.containsKey(collectionName)) { MongoDatabase database = client.getDatabase(targetDatabaseName); collectionMap.put(collectionName, database.getCollection(collectionName)); @@ -1208,7 +1450,6 @@ public void processElement(ProcessContext context, MultiOutputReceiver out) { bufferedEvents.get(collectionName).add(element); - // If we've reached batch size for this collection, process the batch if (bufferedEvents.get(collectionName).size() >= batchSize) { LOG.debug( "Batch size reached for collection {}, processing {} events", @@ -1220,7 +1461,6 @@ public void processElement(ProcessContext context, MultiOutputReceiver out) { @FinishBundle public void finishBundle(FinishBundleContext context) { - // Process any remaining batches for (String collectionName : bufferedEvents.keySet()) { if (!bufferedEvents.get(collectionName).isEmpty()) { LOG.debug( @@ -1241,19 +1481,15 @@ private void processBatch(String collectionName, MultiOutputReceiver out) { } try { - // Create bulk operation List> bulkOperations = new ArrayList<>(events.size()); - // Add operations to bulk for (MongoDbChangeEventContext event : events) { Object docId = event.getDocumentId(); Bson lookupById = eq("_id", docId); if (event.isDeleteEvent()) { - // Add delete operation bulkOperations.add(new DeleteOneModel<>(lookupById)); } else { - // Add upsert operation bulkOperations.add( new ReplaceOneModel<>( lookupById, @@ -1262,7 +1498,6 @@ private void processBatch(String collectionName, MultiOutputReceiver out) { } } - // Execute bulk write with ordered(false) to isolate failed documents BulkWriteResult result = collection.bulkWrite(bulkOperations, new BulkWriteOptions().ordered(false)); LOG.debug( @@ -1271,16 +1506,14 @@ private void processBatch(String collectionName, MultiOutputReceiver out) { result.getInsertedCount() + result.getModifiedCount() + result.getUpserts().size(), result.getDeletedCount()); - // Output successful events for (MongoDbChangeEventContext event : events) { - out.get(successfulWriteTag).output(event); + out.get(SUCCESSFUL_WRITE_TAG).output(event); successfulWrites.inc(); } } catch (MongoBulkWriteException e) { LOG.warn( "Bulk write partially failed for collection {}: {}", collectionName, e.getMessage()); - // Identify failed documents and route them to appropriate tag Set failedIndices = new HashSet<>(); for (BulkWriteError error : e.getWriteErrors()) { failedIndices.add(error.getIndex()); @@ -1290,21 +1523,18 @@ private void processBatch(String collectionName, MultiOutputReceiver out) { failedElement.setErrorMessage(error.getMessage()); failedElement.setStacktrace(Throwables.getStackTraceAsString(e)); - // Check if the error is permanent (e.g. code 2 for InvalidArgument when exceeding nesting - // limit) if (error.getCode() == ProcessChangeEventFn.INVALID_ARGUMENT) { - out.get(severeFailedWriteTag).output(failedElement); + out.get(SEVERE_FAILED_WRITE_TAG).output(failedElement); severeFailedWrites.inc(); } else { - out.get(failedWriteTag).output(failedElement); + out.get(FAILED_WRITE_TAG).output(failedElement); retriableFailedWrites.inc(); } } - // Output successful events that were not part of the failed indices for (int i = 0; i < events.size(); i++) { if (!failedIndices.contains(i)) { - out.get(successfulWriteTag).output(events.get(i)); + out.get(SUCCESSFUL_WRITE_TAG).output(events.get(i)); successfulWrites.inc(); } } @@ -1315,18 +1545,16 @@ private void processBatch(String collectionName, MultiOutputReceiver out) { e.getMessage(), e); - // On error, output all events as failed for (MongoDbChangeEventContext event : events) { FailsafeElement failedElement = FailsafeElement.of(event, event); failedElement.setErrorMessage(e.getMessage()); failedElement.setStacktrace(Throwables.getStackTraceAsString(e)); - out.get(failedWriteTag).output(failedElement); + out.get(FAILED_WRITE_TAG).output(failedElement); retriableFailedWrites.inc(); } } - // Clear the processed batch events.clear(); } @@ -1339,19 +1567,15 @@ private void processBatchFinish(String collectionName, FinishBundleContext conte } try { - // Create bulk operation List> bulkOperations = new ArrayList<>(events.size()); - // Add operations to bulk for (MongoDbChangeEventContext event : events) { Object docId = event.getDocumentId(); Bson lookupById = eq("_id", docId); if (event.isDeleteEvent()) { - // Add delete operation bulkOperations.add(new DeleteOneModel<>(lookupById)); } else { - // Add upsert operation bulkOperations.add( new ReplaceOneModel<>( lookupById, @@ -1360,7 +1584,6 @@ private void processBatchFinish(String collectionName, FinishBundleContext conte } } - // Execute bulk write with ordered(false) to isolate failed documents BulkWriteResult result = collection.bulkWrite(bulkOperations, new BulkWriteOptions().ordered(false)); LOG.debug( @@ -1369,16 +1592,14 @@ private void processBatchFinish(String collectionName, FinishBundleContext conte result.getInsertedCount() + result.getModifiedCount() + result.getUpserts().size(), result.getDeletedCount()); - // Output successful events for (MongoDbChangeEventContext event : events) { - context.output(successfulWriteTag, event, Instant.now(), GlobalWindow.INSTANCE); + context.output(SUCCESSFUL_WRITE_TAG, event, Instant.now(), GlobalWindow.INSTANCE); successfulWrites.inc(); } } catch (MongoBulkWriteException e) { LOG.warn( "Bulk write partially failed for collection {}: {}", collectionName, e.getMessage()); - // Identify failed documents and route them to appropriate tag Set failedIndices = new HashSet<>(); for (BulkWriteError error : e.getWriteErrors()) { failedIndices.add(error.getIndex()); @@ -1388,22 +1609,20 @@ private void processBatchFinish(String collectionName, FinishBundleContext conte failedElement.setErrorMessage(error.getMessage()); failedElement.setStacktrace(Throwables.getStackTraceAsString(e)); - // Check if the error is permanent (e.g. code 2 for InvalidArgument when exceeding nesting - // limit) if (error.getCode() == ProcessChangeEventFn.INVALID_ARGUMENT) { context.output( - severeFailedWriteTag, failedElement, Instant.now(), GlobalWindow.INSTANCE); + SEVERE_FAILED_WRITE_TAG, failedElement, Instant.now(), GlobalWindow.INSTANCE); severeFailedWrites.inc(); } else { - context.output(failedWriteTag, failedElement, Instant.now(), GlobalWindow.INSTANCE); + context.output(FAILED_WRITE_TAG, failedElement, Instant.now(), GlobalWindow.INSTANCE); retriableFailedWrites.inc(); } } - // Output successful events that were not part of the failed indices for (int i = 0; i < events.size(); i++) { if (!failedIndices.contains(i)) { - context.output(successfulWriteTag, events.get(i), Instant.now(), GlobalWindow.INSTANCE); + context.output( + SUCCESSFUL_WRITE_TAG, events.get(i), Instant.now(), GlobalWindow.INSTANCE); successfulWrites.inc(); } } @@ -1414,18 +1633,16 @@ private void processBatchFinish(String collectionName, FinishBundleContext conte e.getMessage(), e); - // On error, output all events as failed for (MongoDbChangeEventContext event : events) { FailsafeElement failedElement = FailsafeElement.of(event, event); failedElement.setErrorMessage(e.getMessage()); failedElement.setStacktrace(Throwables.getStackTraceAsString(e)); - context.output(failedWriteTag, failedElement, Instant.now(), GlobalWindow.INSTANCE); + context.output(FAILED_WRITE_TAG, failedElement, Instant.now(), GlobalWindow.INSTANCE); retriableFailedWrites.inc(); } } - // Clear the processed batch events.clear(); } @@ -1453,7 +1670,6 @@ public void processElement(ProcessContext c) { try { String fullEventJson = element.getPayload(); Document doc = Document.parse(fullEventJson); - // Handle events wrapped in a changeEvent field (common for reconsumed DLQ records) Document innerDoc = Utils.extractInnerEvent(doc); String changeType = innerDoc.getString(DatastreamConstants.EVENT_CHANGE_TYPE_KEY); @@ -1463,20 +1679,16 @@ public void processElement(ProcessContext c) { Object dataVal = innerDoc.get(MongoDbChangeEventContext.DATA_COL); - // Delete events don't have a 'data' field to transform, so we bypass the UDF. if ("DELETE".equalsIgnoreCase(changeType)) { c.output(BYPASS_UDF_TAG, element); return; } - // Update events with null data occurs when an updated document is later - // deleted. In this case, we skip the UDF transformation. if ("UPDATE".equalsIgnoreCase(changeType) && dataVal == null) { skippedUpdates.inc(); - return; // Skip by not outputting anything + return; } - // Extract and canonicalize the 'data' field for UDF input. String canonicalJson = Utils.getCanonicalJsonOfDataField(innerDoc); if (canonicalJson == null) { throw new IllegalArgumentException( @@ -1505,20 +1717,12 @@ public void processElement(ProcessContext c) { try { JsonNode fullEventNode = OBJECT_MAPPER.readTree(fullEventJson); - // Validate that the UDF output is a valid BSON document before proceeding. - // This ensures we don't write invalid data to the destination. Document.parse(transformedData); - // Merge the transformed data back into the full event JSON. - // The payload of the output FailsafeElement will contain this modified event JSON, - // which is used by downstream steps (like writing to Firestore) to process the update. - // The originalPayload remains the raw event for safety and DLQ purposes. JsonNode targetNode = Utils.extractInnerEvent(fullEventNode); ((ObjectNode) targetNode).put(MongoDbChangeEventContext.DATA_COL, transformedData); String modifiedEventJson = OBJECT_MAPPER.writeValueAsString(fullEventNode); - // Output the element with the preserved original payload and the modified event JSON - // containing UDF output. c.output(FailsafeElement.of(element.getOriginalPayload(), modifiedEventJson)); } catch (Exception e) { LOG.error("Error restoring UDF output, exception: {}", e.getMessage(), e); @@ -1554,7 +1758,6 @@ public PCollection> expand( .withOutputTags( UDF_SUCCESS_TAG, TupleTagList.of(PREPARE_FAILURE_TAG).and(BYPASS_UDF_TAG))); - // Handle failed preparation writeFailedJsonToDlq(options, preparedResult, dlqManager, PREPARE_FAILURE_TAG); PCollection> preparedInput = @@ -1597,11 +1800,42 @@ public PCollection> expand( .get(UDF_SUCCESS_TAG) .setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())); - // Merge the restored UDF output with the events that bypassed the UDF. - // Both streams now contain the full event JSON in the required format for downstream steps. return PCollectionList.of(restoredOutput) .and(bypassedElements) .apply("Merge Streams", Flatten.pCollections()); } } + + /** + * Composite PTransform that encapsulates DLQ polling and reconsumption logic under the Read/ + * stage. + */ + public static class ReconsumeDlqTransform extends PTransform { + private final Options options; + private final DeadLetterQueueManager dlqManager; + + public ReconsumeDlqTransform(Options options, DeadLetterQueueManager dlqManager) { + this.options = options; + this.dlqManager = dlqManager; + } + + @Override + public PCollectionTuple expand(PBegin input) { + boolean isRegularMode = "regular".equals(options.getRunMode()); + if (isRegularMode && (!Strings.isNullOrEmpty(options.getDlqGcsPubSubSubscription()))) { + return dlqManager.getReconsumerDataTransformForFiles( + input.apply( + "Read retry from PubSub", + new PubSubNotifiedDlqIO( + options.getDlqGcsPubSubSubscription(), + new ArrayList( + Arrays.asList("/severe/", "/tmp_retry", "/tmp_severe/", ".temp"))))); + } else { + return dlqManager.getReconsumerDataTransform( + input.apply( + "Periodically polling from DLQ", + dlqManager.dlqReconsumer(options.getDlqRetryMinutes()))); + } + } + } } diff --git a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/templates/datastream/DatastreamConstants.java b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/templates/datastream/DatastreamConstants.java index 95803936c0..dc29040194 100644 --- a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/templates/datastream/DatastreamConstants.java +++ b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/templates/datastream/DatastreamConstants.java @@ -15,8 +15,15 @@ */ package com.google.cloud.teleport.v2.templates.datastream; +import com.google.common.collect.ImmutableSet; +import java.util.Set; + /** Constants used in Datastream templates. */ -public class DatastreamConstants { +public final class DatastreamConstants { + + private DatastreamConstants() { + // Utility class; prevent instantiation + } // Common event metadata fields public static final String EVENT_SOURCE_METADATA = "_metadata_source"; @@ -37,11 +44,49 @@ public class DatastreamConstants { // Event types public static final String DELETE_EVENT = "DELETE"; public static final String UPDATE_EVENT = "UPDATE"; + public static final String READ_EVENT = "READ"; public static final String EMPTY_EVENT = ""; + // Read method metadata + public static final String EVENT_READ_METHOD_KEY = "_metadata_read_method"; + public static final String READ_METHOD_BACKFILL = "backfill"; + public static final String READ_METHOD_CDC = "cdc"; + // Default shadow collection prefix public static final String DEFAULT_SHADOW_COLLECTION_PREFIX = "shadow_"; /* Max DoFns per dataflow worker in a streaming pipeline. */ public static final int MAX_DOFN_PER_WORKER = 500; + + /** Internal Datastream metadata fields ignored/stripped from customer documents. */ + public static final Set DATASTREAM_METADATA_FIELDS = + ImmutableSet.of( + "_metadata_stream", + "_metadata_schema", + "_metadata_table", + "_metadata_source", + "_metadata_ssn", + "_metadata_rs_id", + "_metadata_tx_id", + "_metadata_uuid", + "_metadata_dlq_reconsumed", + "_metadata_error", + "_metadata_retry_count", + "_metadata_timestamp", + "_metadata_read_timestamp", + "_metadata_read_method", + "_metadata_deleted", + "_metadata_primary_keys", + "_metadata_log_file", + "_metadata_log_position", + "_metadata_dataflow_timestamp", + "_metadata_timestamp_seconds", + "_metadata_timestamp_nanos"); + + /** Fields ignored when generating shadow metadata documents. */ + public static final Set SHADOW_DOC_IGNORE_FIELDS = + ImmutableSet.builder().addAll(DATASTREAM_METADATA_FIELDS).add("data").build(); + + /** Retained for backwards compatibility. */ + public static final Set MAPPER_IGNORE_FIELDS = DATASTREAM_METADATA_FIELDS; } diff --git a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/templates/datastream/MongoDbChangeEventContext.java b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/templates/datastream/MongoDbChangeEventContext.java index 0e085bf9d3..9053ab0f0c 100644 --- a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/templates/datastream/MongoDbChangeEventContext.java +++ b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/templates/datastream/MongoDbChangeEventContext.java @@ -15,16 +15,17 @@ */ package com.google.cloud.teleport.v2.templates.datastream; -import static com.google.cloud.teleport.v2.templates.DataStreamMongoDBToFirestore.MAPPER_IGNORE_FIELDS; - import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.cloud.teleport.v2.transforms.MongoDbChangeEventContextCoder; import com.google.cloud.teleport.v2.transforms.Utils; import com.google.common.collect.ImmutableMap; +import java.io.IOException; import java.io.Serializable; import java.util.Objects; +import org.apache.beam.sdk.coders.DefaultCoder; import org.bson.Document; import org.bson.types.ObjectId; import org.slf4j.Logger; @@ -34,6 +35,7 @@ * MongoDB's implementation of ChangeEventContext that provides implementation for handling MongoDB * change events. */ +@DefaultCoder(MongoDbChangeEventContextCoder.class) public class MongoDbChangeEventContext implements Serializable { private static final Logger LOG = LoggerFactory.getLogger(MongoDbChangeEventContext.class); @@ -51,6 +53,7 @@ public class MongoDbChangeEventContext implements Serializable { private final JsonNode changeEvent; private final JsonNode originalChangeEvent; + private final String shadowCollectionPrefix; private final String dataCollection; private final String shadowCollection; private final Object documentId; @@ -59,8 +62,8 @@ public class MongoDbChangeEventContext implements Serializable { private final boolean isDeleteEvent; private final boolean isUpdateEvent; private final Document timestampDoc; - private boolean isDlqReconsumed; - private int retryCount; + private final boolean isDlqReconsumed; + private final int retryCount; /** Gets the change type from the event metadata. */ private String getChangeType(JsonNode changeEvent) { @@ -74,6 +77,46 @@ public String getChangeType() { return getChangeType(this.changeEvent); } + /** Determines if the event is a backfill snapshot event. */ + public boolean isBackfillEvent() { + if (changeEvent.has(DatastreamConstants.EVENT_READ_METHOD_KEY)) { + String readMethod = changeEvent.get(DatastreamConstants.EVENT_READ_METHOD_KEY).asText(); + if (DatastreamConstants.READ_METHOD_BACKFILL.equalsIgnoreCase(readMethod)) { + return true; + } + } + String changeType = getChangeType(); + return DatastreamConstants.READ_EVENT.equalsIgnoreCase(changeType) + || "BACKFILL".equalsIgnoreCase(changeType); + } + + /** Determines if the event is a live CDC event. */ + public boolean isCdcEvent() { + return !isBackfillEvent(); + } + + /** Gets epoch timestamp seconds. */ + public long getTimestampSeconds() { + if (timestampDoc != null && timestampDoc.containsKey(TIMESTAMP_SECONDS_COL)) { + Object val = timestampDoc.get(TIMESTAMP_SECONDS_COL); + if (val instanceof Number) { + return ((Number) val).longValue(); + } + } + return 0L; + } + + /** Gets sub-second timestamp (wall nanoseconds for backfill, oplog increment for CDC). */ + public long getTimestampSubSeconds() { + if (timestampDoc != null && timestampDoc.containsKey(TIMESTAMP_NANOS_COL)) { + Object val = timestampDoc.get(TIMESTAMP_NANOS_COL); + if (val instanceof Number) { + return ((Number) val).longValue(); + } + } + return 0L; + } + /** Determines if the event is a delete event based on metadata. */ private boolean isDeleteEvent(JsonNode changeEvent) { String changeType = getChangeType(changeEvent); @@ -105,17 +148,28 @@ public MongoDbChangeEventContext(JsonNode payload, String shadowCollectionPrefix public MongoDbChangeEventContext( JsonNode payload, JsonNode originalPayload, String shadowCollectionPrefix) throws JsonProcessingException { + this( + payload, + originalPayload, + shadowCollectionPrefix, + extractIsDlqReconsumed(payload), + extractRetryCount(payload)); + } + + private MongoDbChangeEventContext( + JsonNode payload, + JsonNode originalPayload, + String shadowCollectionPrefix, + boolean isDlqReconsumed, + int retryCount) + throws JsonProcessingException { // Extracts the actual change event. If wrapped in a DLQ structure like {"changeEvent": {...}}, // it extracts the inner object. this.changeEvent = Utils.extractInnerEvent(payload); this.originalChangeEvent = Utils.extractInnerEvent(originalPayload); - - this.retryCount = - changeEvent.has(DatastreamConstants.RETRY_COUNT) - ? changeEvent.get(DatastreamConstants.RETRY_COUNT).asInt() - : payload.has(DatastreamConstants.RETRY_COUNT) - ? payload.get(DatastreamConstants.RETRY_COUNT).asInt() - : 0; + this.shadowCollectionPrefix = shadowCollectionPrefix != null ? shadowCollectionPrefix : ""; + this.isDlqReconsumed = isDlqReconsumed; + this.retryCount = retryCount; // Extract collection name from the event if (changeEvent.has(DatastreamConstants.EVENT_SOURCE_METADATA)) { @@ -129,7 +183,7 @@ public MongoDbChangeEventContext( throw new IllegalStateException("Invalid event record without _metadata_source."); } - this.shadowCollection = shadowCollectionPrefix + this.dataCollection; + this.shadowCollection = this.shadowCollectionPrefix + this.dataCollection; // Extract document id if (changeEvent.has(DatastreamConstants.MONGODB_DOCUMENT_ID)) { @@ -143,11 +197,13 @@ public MongoDbChangeEventContext( this.documentId = docIdVal.asDouble(); } else if (docIdVal.isTextual()) { this.documentId = docIdVal.asText(); - } else if (docIdVal.isObject()) { - if (docIdVal.has(OID_FIELD_NAME) && docIdVal.get(OID_FIELD_NAME).isTextual()) { + } else if (docIdVal.isObject() || docIdVal.isArray()) { + if (docIdVal.isObject() + && docIdVal.has(OID_FIELD_NAME) + && docIdVal.get(OID_FIELD_NAME).isTextual()) { this.documentId = new ObjectId(docIdVal.get(OID_FIELD_NAME).asText()); } else { - // Support for generic Object-typed IDs or other complex BSON types (e.g., Binary) + // Support for generic Document (Map), Array (List), Binary, and composite BSON _id types Document wrapper = Document.parse("{ \"val\": " + docIdVal.toString() + " }"); this.documentId = wrapper.get("val"); } @@ -182,8 +238,61 @@ public MongoDbChangeEventContext( this.isUpdateEvent = isUpdateEvent(changeEvent); this.jsonStringData = dataAsJsonString(); - this.shadowDocument = generateShadowDocument(); - this.isDlqReconsumed = isDlqReconsumed(changeEvent); + this.shadowDocument = null; + } + + private static boolean extractIsDlqReconsumed(JsonNode payload) { + if (payload == null) { + return false; + } + JsonNode changeEvent = Utils.extractInnerEvent(payload); + if (changeEvent.has(DatastreamConstants.IS_DLQ_RECONSUMED)) { + return changeEvent + .get(DatastreamConstants.IS_DLQ_RECONSUMED) + .asText() + .equalsIgnoreCase("true"); + } + if (payload.has(DatastreamConstants.IS_DLQ_RECONSUMED)) { + return payload.get(DatastreamConstants.IS_DLQ_RECONSUMED).asText().equalsIgnoreCase("true"); + } + return false; + } + + private static int extractRetryCount(JsonNode payload) { + if (payload == null) { + return 0; + } + JsonNode changeEvent = Utils.extractInnerEvent(payload); + if (changeEvent.has(DatastreamConstants.RETRY_COUNT)) { + return changeEvent.get(DatastreamConstants.RETRY_COUNT).asInt(); + } + if (payload.has(DatastreamConstants.RETRY_COUNT)) { + return payload.get(DatastreamConstants.RETRY_COUNT).asInt(); + } + return 0; + } + + /** + * Reconstitutes a {@link MongoDbChangeEventContext} from its serialized components without Java + * reflection serialization. + */ + public static MongoDbChangeEventContext reconstitute( + String changeEventJson, + String originalChangeEventJson, + String shadowPrefix, + boolean isDlq, + int retryCount) + throws IOException { + if (changeEventJson == null) { + return null; + } + JsonNode changeEventNode = OBJECT_MAPPER.readTree(changeEventJson); + JsonNode originalChangeEventNode = + originalChangeEventJson != null + ? OBJECT_MAPPER.readTree(originalChangeEventJson) + : changeEventNode; + return new MongoDbChangeEventContext( + changeEventNode, originalChangeEventNode, shadowPrefix, isDlq, retryCount); } /** Creates a shadow document for tracking event ordering. */ @@ -203,7 +312,7 @@ public Document generateShadowDocument() throws JsonProcessingException { shadowDoc.put("processed_at", System.currentTimeMillis()); shadowDoc.put("is_from_dlq", isDlqReconsumed); - Utils.removeTableRowFields(shadowDoc, MAPPER_IGNORE_FIELDS); + Utils.removeTableRowFields(shadowDoc, DatastreamConstants.SHADOW_DOC_IGNORE_FIELDS); return shadowDoc; } @@ -224,6 +333,18 @@ public JsonNode getOriginalChangeEvent() { return originalChangeEvent; } + public String getChangeEventJsonString() { + return changeEvent != null ? changeEvent.toString() : null; + } + + public String getOriginalChangeEventJsonString() { + return originalChangeEvent != null ? originalChangeEvent.toString() : null; + } + + public String getShadowCollectionPrefix() { + return shadowCollectionPrefix; + } + public String getDataCollection() { return dataCollection; } @@ -245,6 +366,13 @@ public boolean isUpdateEvent() { } public Document getShadowDocument() { + if (this.shadowDocument == null && this.shadowCollection != null) { + try { + return generateShadowDocument(); + } catch (JsonProcessingException e) { + LOG.warn("Failed to generate shadow document: {}", e.getMessage()); + } + } return shadowDocument; } @@ -291,8 +419,8 @@ public String toString() { // Convert timestamp document to JSON if (this.timestampDoc != null) { ObjectNode timestampNode = OBJECT_MAPPER.createObjectNode(); - timestampNode.put(TIMESTAMP_SECONDS_COL, this.timestampDoc.getLong(TIMESTAMP_SECONDS_COL)); - timestampNode.put(TIMESTAMP_NANOS_COL, this.timestampDoc.getInteger(TIMESTAMP_NANOS_COL)); + timestampNode.put(TIMESTAMP_SECONDS_COL, getTimestampSeconds()); + timestampNode.put(TIMESTAMP_NANOS_COL, getTimestampSubSeconds()); jsonNode.set(TIMESTAMP_COL, timestampNode); } @@ -317,10 +445,39 @@ public String toString() { } } + @Override public boolean equals(Object other) { + if (this == other) { + return true; + } if (other instanceof MongoDbChangeEventContext) { - return Objects.equals(this.toString(), other.toString()); + MongoDbChangeEventContext o = (MongoDbChangeEventContext) other; + return Objects.equals(this.dataCollection, o.dataCollection) + && Objects.equals(this.shadowCollectionPrefix, o.shadowCollectionPrefix) + && Objects.equals(this.documentId, o.documentId) + && Objects.equals(this.timestampDoc, o.timestampDoc) + && this.isDeleteEvent == o.isDeleteEvent + && this.isUpdateEvent == o.isUpdateEvent + && this.isDlqReconsumed == o.isDlqReconsumed + && this.retryCount == o.retryCount + && Objects.equals(this.changeEvent, o.changeEvent) + && Objects.equals(this.originalChangeEvent, o.originalChangeEvent); } return false; } + + @Override + public int hashCode() { + return Objects.hash( + dataCollection, + shadowCollectionPrefix, + documentId, + timestampDoc, + isDeleteEvent, + isUpdateEvent, + isDlqReconsumed, + retryCount, + changeEvent, + originalChangeEvent); + } } diff --git a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/CreateMongoDbChangeEventContextFn.java b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/CreateMongoDbChangeEventContextFn.java index 1c095bf1b4..17f2ab9375 100644 --- a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/CreateMongoDbChangeEventContextFn.java +++ b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/CreateMongoDbChangeEventContextFn.java @@ -35,9 +35,9 @@ public class CreateMongoDbChangeEventContextFn LoggerFactory.getLogger(CreateMongoDbChangeEventContextFn.class); private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); - public static TupleTag successfulCreationTag = + public static final TupleTag SUCCESSFUL_CREATION_TAG = new TupleTag<>("successfulCreation"); - public static TupleTag> failedCreationTag = + public static final TupleTag> FAILED_CREATION_TAG = new TupleTag<>("failedCreation"); private final String shadowCollectionPrefix; @@ -56,12 +56,12 @@ public void processElement(ProcessContext context, MultiOutputReceiver out) { JsonNode originalNode = OBJECT_MAPPER.readTree(element.getOriginalPayload()); MongoDbChangeEventContext changeEventContext = new MongoDbChangeEventContext(jsonNode, originalNode, shadowCollectionPrefix); - out.get(successfulCreationTag).output(changeEventContext); + out.get(SUCCESSFUL_CREATION_TAG).output(changeEventContext); } catch (Exception e) { LOG.error("Error creating MongoDbChangeEventContext, exception: {}, element: {}", e, element); element.setErrorMessage(e.getMessage()); element.setStacktrace(Throwables.getStackTraceAsString(e)); - out.get(failedCreationTag).output(element); + out.get(FAILED_CREATION_TAG).output(element); contextCreationFailures.inc(); LOG.info("Failed element sent to DLQ"); } diff --git a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbBulkTransforms.java b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbBulkTransforms.java new file mode 100644 index 0000000000..fff01a51e3 --- /dev/null +++ b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbBulkTransforms.java @@ -0,0 +1,1208 @@ +/* + * 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 com.mongodb.client.model.Filters.eq; + +import com.google.cloud.teleport.v2.coders.FailsafeElementCoder; +import com.google.cloud.teleport.v2.templates.datastream.MongoDbChangeEventContext; +import com.google.cloud.teleport.v2.values.FailsafeElement; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Throwables; +import com.google.common.collect.ImmutableSet; +import com.google.common.util.concurrent.RateLimiter; +import com.mongodb.ConnectionString; +import com.mongodb.MongoBulkWriteException; +import com.mongodb.MongoClientSettings; +import com.mongodb.MongoException; +import com.mongodb.bulk.BulkWriteError; +import com.mongodb.bulk.BulkWriteResult; +import com.mongodb.bulk.WriteConcernError; +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoClients; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.model.BulkWriteOptions; +import com.mongodb.client.model.DeleteOneModel; +import com.mongodb.client.model.ReplaceOneModel; +import com.mongodb.client.model.ReplaceOptions; +import com.mongodb.client.model.WriteModel; +import java.io.Serializable; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; +import org.apache.beam.sdk.coders.SerializableCoder; +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.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.SerializableFunction; +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.PCollection; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.TupleTagList; +import org.bson.Document; +import org.bson.UuidRepresentation; +import org.bson.conversions.Bson; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** High-throughput asynchronous bulk transforms for writing CDC events to Firestore/MongoDB. */ +public class MongoDbBulkTransforms { + + private static final Logger LOG = LoggerFactory.getLogger(MongoDbBulkTransforms.class); + + // Permanent error codes that should not be retried in-memory + public static final int ERR_BAD_VALUE = 2; + public static final int ERR_UNAUTHORIZED = 13; + public static final int ERR_TYPE_MISMATCH = 14; + public static final int ERR_INVALID_LENGTH = 16; + public static final int ERR_NAMESPACE_NOT_FOUND = 26; + public static final int ERR_IMMUTABLE_FIELD = 66; + public static final int ERR_NETWORK_TIMEOUT = 89; + public static final int ERR_SHUTDOWN_IN_PROGRESS = 91; + public static final int ERR_WRITE_CONFLICT = 112; + public static final int ERR_DOCUMENT_VALIDATION_FAILURE = 121; + public static final int ERR_DUPLICATE_KEY = 11000; + public static final int ERR_KEY_TOO_LONG = 17280; + + public static final Set PERMANENT_ERROR_CODES = + ImmutableSet.of( + ERR_BAD_VALUE, + ERR_UNAUTHORIZED, + ERR_TYPE_MISMATCH, + ERR_INVALID_LENGTH, + ERR_NAMESPACE_NOT_FOUND, + ERR_IMMUTABLE_FIELD, + ERR_DOCUMENT_VALIDATION_FAILURE, + ERR_KEY_TOO_LONG); + + public static final TupleTag SUCCESSFUL_WRITE_TAG = + new TupleTag("successfulWrite") {}; + public static final TupleTag< + FailsafeElement> + FAILED_WRITE_TAG = + new TupleTag>( + "failedWrite") {}; + public static final TupleTag< + FailsafeElement> + SEVERE_FAILED_WRITE_TAG = + new TupleTag>( + "severeFailedWrite") {}; + + 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); + } + builder.applyToSocketSettings( + b -> { + b.connectTimeout(60, TimeUnit.SECONDS); + b.readTimeout(60, TimeUnit.SECONDS); + }); + builder.applyToConnectionPoolSettings( + b -> { + b.minSize(10); + b.maxSize(200); + b.maxWaitTime(15, TimeUnit.SECONDS); + }); + builder.applyToClusterSettings(b -> b.serverSelectionTimeout(10, TimeUnit.MINUTES)); + return MongoClients.create(builder.build()); + } + + public static boolean isPermanentErrorCode(int code) { + return PERMANENT_ERROR_CODES.contains(code); + } + + public static boolean isPermanentError(int code) { + return isPermanentErrorCode(code); + } + + public static BulkWriteWithDlq bulkWriteWithDlq() { + return new BulkWriteWithDlq(); + } + + /** PTransform encapsulating asynchronous bulk writing with retry and severe DLQ routing. */ + public static class BulkWriteWithDlq + extends PTransform, PCollectionTuple> { + + public static final TupleTag SUCCESS_TAG = SUCCESSFUL_WRITE_TAG; + public static final TupleTag< + FailsafeElement> + FAILED_TAG = FAILED_WRITE_TAG; + public static final TupleTag< + FailsafeElement> + SEVERE_FAILED_TAG = SEVERE_FAILED_WRITE_TAG; + + private String connectionString; + private String database; + private int batchSize = 500; + private int maxConcurrentAsyncWrites = 10; + private int initialWriteRatePerWorker = 500; + private int writeRateRampUpMinutes = 5; + private int writeRateRampUpSteps = 5; + private int maxWriteRatePerWorker = 2500; + private int maxWriteRetries = 3; + private int dlqMaxRetries = 3; + private SerializableFunction clientFactory = + MongoDbBulkTransforms::createMongoClient; + + public BulkWriteWithDlq withConnectionString(String connectionString) { + this.connectionString = connectionString; + return this; + } + + public BulkWriteWithDlq withUri(String uri) { + this.connectionString = uri; + return this; + } + + public BulkWriteWithDlq withDatabase(String database) { + this.database = database; + return this; + } + + public BulkWriteWithDlq withBatchSize(Integer batchSize) { + if (batchSize != null) { + this.batchSize = batchSize; + } + return this; + } + + public BulkWriteWithDlq withMaxConcurrentAsyncWrites(Integer maxConcurrentAsyncWrites) { + if (maxConcurrentAsyncWrites != null) { + this.maxConcurrentAsyncWrites = maxConcurrentAsyncWrites; + } + return this; + } + + public BulkWriteWithDlq withInitialWriteRatePerWorker(Integer initialWriteRatePerWorker) { + if (initialWriteRatePerWorker != null) { + this.initialWriteRatePerWorker = initialWriteRatePerWorker; + } + return this; + } + + public BulkWriteWithDlq withWriteRateRampUpMinutes(Integer writeRateRampUpMinutes) { + if (writeRateRampUpMinutes != null) { + this.writeRateRampUpMinutes = writeRateRampUpMinutes; + } + return this; + } + + public BulkWriteWithDlq withWriteRateRampUpSteps(Integer writeRateRampUpSteps) { + if (writeRateRampUpSteps != null) { + this.writeRateRampUpSteps = writeRateRampUpSteps; + } + return this; + } + + public BulkWriteWithDlq withMaxWriteRatePerWorker(Integer maxWriteRatePerWorker) { + if (maxWriteRatePerWorker != null) { + this.maxWriteRatePerWorker = maxWriteRatePerWorker; + } + return this; + } + + public BulkWriteWithDlq withMaxWriteRetries(Integer maxWriteRetries) { + if (maxWriteRetries != null) { + this.maxWriteRetries = maxWriteRetries; + } + return this; + } + + public BulkWriteWithDlq withDlqMaxRetries(Integer dlqMaxRetries) { + if (dlqMaxRetries != null) { + this.dlqMaxRetries = dlqMaxRetries; + } + return this; + } + + public BulkWriteWithDlq withClientFactory( + SerializableFunction clientFactory) { + if (clientFactory != null) { + this.clientFactory = clientFactory; + } + return this; + } + + @Override + public PCollectionTuple expand(PCollection input) { + PCollectionTuple result = + input.apply( + "AsyncBulkWriteFn", + ParDo.of( + new BulkWriteFn( + connectionString, + database, + batchSize, + maxConcurrentAsyncWrites, + initialWriteRatePerWorker, + writeRateRampUpMinutes, + writeRateRampUpSteps, + maxWriteRatePerWorker, + maxWriteRetries, + dlqMaxRetries, + clientFactory, + SUCCESS_TAG, + FAILED_TAG, + SEVERE_FAILED_TAG)) + .withOutputTags(SUCCESS_TAG, TupleTagList.of(FAILED_TAG).and(SEVERE_FAILED_TAG))); + + result.get(SUCCESS_TAG).setCoder(SerializableCoder.of(MongoDbChangeEventContext.class)); + result + .get(FAILED_TAG) + .setCoder( + FailsafeElementCoder.of( + SerializableCoder.of(MongoDbChangeEventContext.class), + SerializableCoder.of(MongoDbChangeEventContext.class))); + result + .get(SEVERE_FAILED_TAG) + .setCoder( + FailsafeElementCoder.of( + SerializableCoder.of(MongoDbChangeEventContext.class), + SerializableCoder.of(MongoDbChangeEventContext.class))); + + return result; + } + } + + /** DoFn implementing non-transactional async bulk write with rate limiter ramp-up and triage. */ + public static class BulkWriteFn + extends DoFn { + + private static final Logger LOG = LoggerFactory.getLogger(BulkWriteFn.class); + + private final String connectionString; + private final String database; + private final int batchSize; + private final int maxConcurrentAsyncWrites; + private final int initialWriteRatePerWorker; + private final int writeRateRampUpMinutes; + private final int writeRateRampUpSteps; + private final int maxWriteRatePerWorker; + private final int maxWriteRetries; + private final int dlqMaxRetries; + private final SerializableFunction clientFactory; + private final TupleTag successTag; + private final TupleTag> + failureTag; + private final TupleTag> + severeFailureTag; + + private transient MongoClient mongoClient; + private transient ExecutorService executorService; + private transient Semaphore semaphore; + private transient RateLimiter rateLimiter; + private transient long setupStartTimeMs; + private transient Map> currentBatches; + private transient List> inFlightFutures; + private transient ConcurrentLinkedQueue successQueue; + private transient ConcurrentLinkedQueue< + FailsafeElement> + failureQueue; + private transient ConcurrentLinkedQueue< + FailsafeElement> + severeFailureQueue; + private transient Map> collectionsMap; + private transient ThrottledLogger throttledLogger; + + private final Counter successfulWrites = Metrics.counter(BulkWriteFn.class, "successfulWrites"); + private final Counter retriableFailedWrites = + Metrics.counter(BulkWriteFn.class, "retriableFailedWrites"); + private final Counter severeFailedWrites = + Metrics.counter(BulkWriteFn.class, "severeFailedWrites"); + + public BulkWriteFn( + String connectionString, + String database, + int batchSize, + int maxConcurrentAsyncWrites, + int initialWriteRatePerWorker, + int writeRateRampUpMinutes, + int writeRateRampUpSteps, + int maxWriteRatePerWorker, + int maxWriteRetries, + int dlqMaxRetries, + SerializableFunction clientFactory, + TupleTag successTag, + TupleTag> failureTag, + TupleTag> + severeFailureTag) { + this.connectionString = connectionString; + this.database = database; + this.batchSize = batchSize > 0 ? batchSize : 500; + this.maxConcurrentAsyncWrites = maxConcurrentAsyncWrites > 0 ? maxConcurrentAsyncWrites : 10; + this.initialWriteRatePerWorker = initialWriteRatePerWorker; + this.writeRateRampUpMinutes = writeRateRampUpMinutes > 0 ? writeRateRampUpMinutes : 5; + this.writeRateRampUpSteps = writeRateRampUpSteps > 0 ? writeRateRampUpSteps : 5; + this.maxWriteRatePerWorker = maxWriteRatePerWorker > 0 ? maxWriteRatePerWorker : 2500; + this.maxWriteRetries = maxWriteRetries >= 0 ? maxWriteRetries : 3; + this.dlqMaxRetries = dlqMaxRetries >= 0 ? dlqMaxRetries : 3; + this.clientFactory = clientFactory; + this.successTag = successTag; + this.failureTag = failureTag; + this.severeFailureTag = severeFailureTag; + } + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + private String connectionString; + private String database; + private int batchSize = 500; + private int maxConcurrentAsyncWrites = 10; + private int initialWriteRatePerWorker = 500; + private int writeRateRampUpMinutes = 5; + private int writeRateRampUpSteps = 5; + private int maxWriteRatePerWorker = 2500; + private int maxWriteRetries = 3; + private int dlqMaxRetries = 3; + private SerializableFunction clientFactory = + MongoDbBulkTransforms::createMongoClient; + private TupleTag successTag = SUCCESSFUL_WRITE_TAG; + private TupleTag> + failureTag = FAILED_WRITE_TAG; + private TupleTag> + severeFailureTag = SEVERE_FAILED_WRITE_TAG; + + public Builder withConnectionString(String connectionString) { + this.connectionString = connectionString; + return this; + } + + public Builder withUri(String uri) { + this.connectionString = uri; + return this; + } + + public Builder withDatabase(String database) { + this.database = database; + return this; + } + + public Builder withBatchSize(Integer batchSize) { + if (batchSize != null) { + this.batchSize = batchSize; + } + return this; + } + + public Builder withMaxConcurrentAsyncWrites(Integer maxConcurrentAsyncWrites) { + if (maxConcurrentAsyncWrites != null) { + this.maxConcurrentAsyncWrites = maxConcurrentAsyncWrites; + } + 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 withMaxWriteRetries(Integer maxWriteRetries) { + if (maxWriteRetries != null) { + this.maxWriteRetries = maxWriteRetries; + } + return this; + } + + public Builder withDlqMaxRetries(Integer dlqMaxRetries) { + if (dlqMaxRetries != null) { + this.dlqMaxRetries = dlqMaxRetries; + } + return this; + } + + public Builder withClientFactory(SerializableFunction clientFactory) { + if (clientFactory != null) { + this.clientFactory = clientFactory; + } + return this; + } + + public Builder withSuccessTag(TupleTag successTag) { + this.successTag = successTag; + return this; + } + + public Builder withFailureTag( + TupleTag> + failureTag) { + this.failureTag = failureTag; + return this; + } + + public Builder withSevereFailureTag( + TupleTag> + severeFailureTag) { + this.severeFailureTag = severeFailureTag; + return this; + } + + public BulkWriteFn build() { + return new BulkWriteFn( + connectionString, + database, + batchSize, + maxConcurrentAsyncWrites, + initialWriteRatePerWorker, + writeRateRampUpMinutes, + writeRateRampUpSteps, + maxWriteRatePerWorker, + maxWriteRetries, + dlqMaxRetries, + clientFactory, + successTag, + failureTag, + severeFailureTag); + } + } + + @Setup + public void setup() { + if (mongoClient == null) { + if (clientFactory != null) { + mongoClient = clientFactory.apply(connectionString); + } else { + mongoClient = createMongoClient(connectionString); + } + } + + int threads = Math.max(1, maxConcurrentAsyncWrites); + this.executorService = Executors.newFixedThreadPool(threads); + this.semaphore = new Semaphore(threads); + + double initialRate = + (initialWriteRatePerWorker > 0) + ? initialWriteRatePerWorker + : Math.max(1.0, maxWriteRatePerWorker); + if (initialWriteRatePerWorker > 0 || maxWriteRatePerWorker > 0) { + this.rateLimiter = RateLimiter.create(Math.max(1.0, initialRate)); + } else { + this.rateLimiter = null; + } + + this.setupStartTimeMs = System.currentTimeMillis(); + this.currentBatches = new ConcurrentHashMap<>(); + this.inFlightFutures = new ArrayList<>(); + this.successQueue = new ConcurrentLinkedQueue<>(); + this.failureQueue = new ConcurrentLinkedQueue<>(); + this.severeFailureQueue = new ConcurrentLinkedQueue<>(); + this.collectionsMap = new ConcurrentHashMap<>(); + this.throttledLogger = new ThrottledLogger("BulkWriteFn", 30000L); + } + + @StartBundle + public void startBundle() { + if (currentBatches != null) { + currentBatches.clear(); + } + if (inFlightFutures != null) { + inFlightFutures.clear(); + } + if (successQueue != null) { + successQueue.clear(); + } + if (failureQueue != null) { + failureQueue.clear(); + } + if (severeFailureQueue != null) { + severeFailureQueue.clear(); + } + } + + @ProcessElement + public void processElement(ProcessContext context) { + MongoDbChangeEventContext event = context.element(); + if (event == null) { + return; + } + String collectionName = event.getDataCollection(); + List batch = + currentBatches.computeIfAbsent(collectionName, k -> new ArrayList<>()); + batch.add(event); + + if (batch.size() >= batchSize) { + flushBatch(collectionName, batch); + } + + drainQueues(context); + } + + @FinishBundle + public void finishBundle(FinishBundleContext context) { + if (currentBatches != null) { + for (Map.Entry> entry : currentBatches.entrySet()) { + if (!entry.getValue().isEmpty()) { + flushBatch(entry.getKey(), entry.getValue()); + } + } + } + + try { + if (inFlightFutures != null && !inFlightFutures.isEmpty()) { + CompletableFuture.allOf(inFlightFutures.toArray(new CompletableFuture[0])).join(); + } + } finally { + drainQueuesFinishBundle(context); + if (inFlightFutures != null) { + inFlightFutures.clear(); + } + } + } + + @Teardown + public void teardown() { + if (executorService != null) { + executorService.shutdown(); + try { + if (!executorService.awaitTermination(30, TimeUnit.SECONDS)) { + executorService.shutdownNow(); + } + } catch (InterruptedException e) { + executorService.shutdownNow(); + Thread.currentThread().interrupt(); + } + } + if (mongoClient != null) { + mongoClient.close(); + mongoClient = null; + } + } + + private void flushBatch(String collectionName, List batch) { + if (batch == null || batch.isEmpty()) { + return; + } + List batchToExecute = new ArrayList<>(batch); + batch.clear(); + + try { + semaphore.acquire(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted acquiring semaphore permit", e); + } + + CompletableFuture future; + try { + future = + CompletableFuture.runAsync( + () -> { + try { + executeBatch(collectionName, batchToExecute); + } finally { + semaphore.release(); + } + }, + executorService); + } catch (Throwable t) { + semaphore.release(); + throw t; + } + + inFlightFutures.add(future); + } + + private void applyRateLimiter(int permits) { + if (rateLimiter == null || permits <= 0) { + return; + } + if (initialWriteRatePerWorker > 0 && maxWriteRatePerWorker > 0) { + long elapsedMs = System.currentTimeMillis() - setupStartTimeMs; + long rampUpMs = writeRateRampUpMinutes * 60 * 1000L; + double currentRate; + if (rampUpMs <= 0 || elapsedMs >= rampUpMs) { + currentRate = maxWriteRatePerWorker; + } else { + double fraction = (double) elapsedMs / rampUpMs; + if (writeRateRampUpSteps > 0) { + fraction = Math.floor(fraction * writeRateRampUpSteps) / writeRateRampUpSteps; + } + currentRate = + initialWriteRatePerWorker + + (maxWriteRatePerWorker - initialWriteRatePerWorker) * fraction; + } + double targetRate = Math.max(1.0, currentRate); + if (Math.abs(rateLimiter.getRate() - targetRate) > 1e-3) { + rateLimiter.setRate(targetRate); + } + } + rateLimiter.acquire(permits); + } + + private void executeBatch(String collectionName, List batch) { + if (batch == null || batch.isEmpty()) { + return; + } + + List activeBatch = new ArrayList<>(); + Map> supersededPerDoc = new LinkedHashMap<>(); + + try { + applyRateLimiter(batch.size()); + + // Intra-batch coalescing: coalesce operations per document ID within the batch to ensure + // only + // the latest state is written + Map latestPerDoc = new LinkedHashMap<>(); + for (MongoDbChangeEventContext event : batch) { + if (event == null) { + continue; + } + Object docId = event.getDocumentId(); + MongoDbChangeEventContext existing = latestPerDoc.get(docId); + if (existing == null) { + latestPerDoc.put(docId, event); + } else { + TimestampSortKey eventKey = TimestampSortKey.of(event); + TimestampSortKey existingKey = TimestampSortKey.of(existing); + int cmp = + (existingKey == null) + ? 1 + : (eventKey == null ? -1 : eventKey.compareTo(existingKey)); + if (cmp > 0 || (cmp == 0 && event.getIsDlqReconsumed())) { + supersededPerDoc.computeIfAbsent(docId, k -> new ArrayList<>()).add(existing); + latestPerDoc.put(docId, event); + } else { + supersededPerDoc.computeIfAbsent(docId, k -> new ArrayList<>()).add(event); + } + } + } + + MongoCollection collection = getCollection(collectionName); + List> operations = new ArrayList<>(latestPerDoc.size()); + + for (MongoDbChangeEventContext event : latestPerDoc.values()) { + Object docId = event.getDocumentId(); + Bson lookupById = eq("_id", docId); + if (event.isDeleteEvent()) { + operations.add(new DeleteOneModel<>(lookupById)); + activeBatch.add(event); + } else { + Document doc = Utils.jsonToDocument(event.getDataAsJsonString(), docId); + if (doc == null) { + if (event.isUpdateEvent()) { + // Null data on update event occurs when doc was deleted right after update; skip + LOG.info( + "Skipping update event for document ID: {} because 'data' field is null", + docId); + successQueue.add(event); + successfulWrites.inc(); + List superseded = supersededPerDoc.remove(docId); + if (superseded != null && !superseded.isEmpty()) { + successQueue.addAll(superseded); + successfulWrites.inc(superseded.size()); + } + } else { + // Missing document data for non-delete/non-update event -> send to severe DLQ + FailsafeElement + severeElement = FailsafeElement.of(event, event); + severeElement.setErrorMessage("Missing or null document data for docId: " + docId); + severeFailureQueue.add(severeElement); + severeFailedWrites.inc(); + supersededPerDoc.remove(docId); + } + } else { + operations.add( + new ReplaceOneModel<>(lookupById, doc, new ReplaceOptions().upsert(true))); + activeBatch.add(event); + } + } + } + + if (operations.isEmpty()) { + return; + } + + BulkWriteResult result = + collection.bulkWrite(operations, new BulkWriteOptions().ordered(false)); + for (MongoDbChangeEventContext event : activeBatch) { + successQueue.add(event); + successfulWrites.inc(); + List superseded = + supersededPerDoc.remove(event.getDocumentId()); + if (superseded != null && !superseded.isEmpty()) { + successQueue.addAll(superseded); + successfulWrites.inc(superseded.size()); + } + } + } catch (MongoBulkWriteException mbwe) { + handleBulkWriteException(collectionName, activeBatch, supersededPerDoc, mbwe); + } catch (Exception e) { + List eventsToHandle = + (!activeBatch.isEmpty()) ? activeBatch : batch; + handleGeneralBatchException(collectionName, eventsToHandle, supersededPerDoc, e); + } + } + + private void handleBulkWriteException( + String collectionName, + List batch, + Map> supersededPerDoc, + MongoBulkWriteException mbwe) { + Set failedIndices = new HashSet<>(); + List transientEvents = new ArrayList<>(); + WriteConcernError writeConcernError = mbwe.getWriteConcernError(); + + for (BulkWriteError error : mbwe.getWriteErrors()) { + int idx = error.getIndex(); + failedIndices.add(idx); + if (idx >= 0 && idx < batch.size()) { + MongoDbChangeEventContext event = batch.get(idx); + int code = error.getCode(); + if (isPermanentErrorCode(code)) { + FailsafeElement severeElement = + FailsafeElement.of(event, event); + severeElement.setErrorMessage( + "Permanent write error (Code " + code + "): " + error.getMessage()); + severeElement.setStacktrace(Throwables.getStackTraceAsString(mbwe)); + severeFailureQueue.add(severeElement); + severeFailedWrites.inc(); + if (supersededPerDoc != null) { + supersededPerDoc.remove(event.getDocumentId()); + } + } else { + transientEvents.add(event); + } + } + } + + // If writeConcernError occurred, retry all non-permanently failed events in batch + if (writeConcernError != null) { + throttledLogger.logWarn( + LOG, + collectionName, + "Encountered WriteConcernError: {}, retrying unconfirmed batch", + writeConcernError.getMessage()); + for (int i = 0; i < batch.size(); i++) { + if (!failedIndices.contains(i)) { + transientEvents.add(batch.get(i)); + } + } + } else { + // Output successful documents only when no write concern error occurred + for (int i = 0; i < batch.size(); i++) { + if (!failedIndices.contains(i)) { + MongoDbChangeEventContext successfulEvent = batch.get(i); + successQueue.add(successfulEvent); + successfulWrites.inc(); + if (supersededPerDoc != null) { + List superseded = + supersededPerDoc.remove(successfulEvent.getDocumentId()); + if (superseded != null && !superseded.isEmpty()) { + successQueue.addAll(superseded); + successfulWrites.inc(superseded.size()); + } + } + } + } + } + + // Retry transient documents with backoff + if (!transientEvents.isEmpty()) { + retryTransientEvents(collectionName, transientEvents, supersededPerDoc); + } + } + + private void handleGeneralBatchException( + String collectionName, + List batch, + Map> supersededPerDoc, + Exception e) { + int code = 0; + if (e instanceof MongoException me) { + code = me.getCode(); + } + + if (isPermanentErrorCode(code)) { + throttledLogger.logError( + LOG, + collectionName, + "Permanent failure ({}) during bulkWrite: {}", + code, + e.getMessage()); + for (MongoDbChangeEventContext event : batch) { + FailsafeElement severeElement = + FailsafeElement.of(event, event); + severeElement.setErrorMessage("Permanent failure: " + e.getMessage()); + severeElement.setStacktrace(Throwables.getStackTraceAsString(e)); + severeFailureQueue.add(severeElement); + severeFailedWrites.inc(); + if (supersededPerDoc != null) { + supersededPerDoc.remove(event.getDocumentId()); + } + } + return; + } + + throttledLogger.logWarn( + LOG, + collectionName, + "Batch write encountered retryable error for collection {}: {}", + collectionName, + e.getMessage()); + retryTransientEvents(collectionName, batch, supersededPerDoc); + } + + private void retryTransientEvents( + String collectionName, + List events, + Map> supersededPerDoc) { + FluentBackoff backoff = + FluentBackoff.DEFAULT + .withInitialBackoff(Duration.standardSeconds(2)) + .withExponent(2.0) + .withMaxRetries(maxWriteRetries); + BackOff backoffInstance = backoff.backoff(); + Sleeper sleeper = Sleeper.DEFAULT; + + List currentRemaining = new ArrayList<>(events); + while (!currentRemaining.isEmpty()) { + List activeRetryBatch = new ArrayList<>(currentRemaining.size()); + try { + MongoCollection collection = getCollection(collectionName); + List> operations = new ArrayList<>(currentRemaining.size()); + for (MongoDbChangeEventContext event : currentRemaining) { + Object docId = event.getDocumentId(); + Bson lookupById = eq("_id", docId); + if (event.isDeleteEvent()) { + operations.add(new DeleteOneModel<>(lookupById)); + activeRetryBatch.add(event); + } else { + Document doc = Utils.jsonToDocument(event.getDataAsJsonString(), docId); + if (doc == null) { + if (event.isUpdateEvent()) { + successQueue.add(event); + successfulWrites.inc(); + if (supersededPerDoc != null) { + List superseded = supersededPerDoc.remove(docId); + if (superseded != null && !superseded.isEmpty()) { + successQueue.addAll(superseded); + successfulWrites.inc(superseded.size()); + } + } + } else { + FailsafeElement + severeElement = FailsafeElement.of(event, event); + severeElement.setErrorMessage( + "Missing or null document data on retry for docId: " + docId); + severeFailureQueue.add(severeElement); + severeFailedWrites.inc(); + if (supersededPerDoc != null) { + supersededPerDoc.remove(docId); + } + } + } else { + operations.add( + new ReplaceOneModel<>(lookupById, doc, new ReplaceOptions().upsert(true))); + activeRetryBatch.add(event); + } + } + } + + if (operations.isEmpty()) { + currentRemaining.clear(); + break; + } + + collection.bulkWrite(operations, new BulkWriteOptions().ordered(false)); + for (MongoDbChangeEventContext event : activeRetryBatch) { + successQueue.add(event); + successfulWrites.inc(); + if (supersededPerDoc != null) { + List superseded = + supersededPerDoc.remove(event.getDocumentId()); + if (superseded != null && !superseded.isEmpty()) { + successQueue.addAll(superseded); + successfulWrites.inc(superseded.size()); + } + } + } + currentRemaining.clear(); + break; + } catch (MongoBulkWriteException mbwe) { + Set failedIndices = new HashSet<>(); + List nextRetry = new ArrayList<>(); + for (BulkWriteError err : mbwe.getWriteErrors()) { + int idx = err.getIndex(); + failedIndices.add(idx); + if (idx >= 0 && idx < activeRetryBatch.size()) { + MongoDbChangeEventContext ev = activeRetryBatch.get(idx); + if (isPermanentErrorCode(err.getCode())) { + FailsafeElement + severeElement = FailsafeElement.of(ev, ev); + severeElement.setErrorMessage( + "Permanent write error on retry (Code " + + err.getCode() + + "): " + + err.getMessage()); + severeElement.setStacktrace(Throwables.getStackTraceAsString(mbwe)); + severeFailureQueue.add(severeElement); + severeFailedWrites.inc(); + if (supersededPerDoc != null) { + supersededPerDoc.remove(ev.getDocumentId()); + } + } else { + nextRetry.add(ev); + } + } + } + for (int i = 0; i < activeRetryBatch.size(); i++) { + if (!failedIndices.contains(i)) { + MongoDbChangeEventContext successfulEvent = activeRetryBatch.get(i); + successQueue.add(successfulEvent); + successfulWrites.inc(); + if (supersededPerDoc != null) { + List superseded = + supersededPerDoc.remove(successfulEvent.getDocumentId()); + if (superseded != null && !superseded.isEmpty()) { + successQueue.addAll(superseded); + successfulWrites.inc(superseded.size()); + } + } + } + } + currentRemaining = nextRetry; + } catch (Exception e) { + throttledLogger.logWarn( + LOG, + collectionName, + "Error during retry for collection {}: {}", + collectionName, + e.getMessage()); + } + + if (!currentRemaining.isEmpty()) { + try { + if (!BackOffUtils.next(sleeper, backoffInstance)) { + // Backoff exhausted -> route to DLQ + for (MongoDbChangeEventContext ev : currentRemaining) { + FailsafeElement + failedElement = FailsafeElement.of(ev, ev); + failedElement.setErrorMessage( + "Transient write error retries exhausted after " + + maxWriteRetries + + " attempts"); + failureQueue.add(failedElement); + retriableFailedWrites.inc(); + if (supersededPerDoc != null) { + supersededPerDoc.remove(ev.getDocumentId()); + } + } + break; + } + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + for (MongoDbChangeEventContext ev : currentRemaining) { + FailsafeElement failedElement = + FailsafeElement.of(ev, ev); + failedElement.setErrorMessage("Retry interrupted: " + ie.getMessage()); + failureQueue.add(failedElement); + retriableFailedWrites.inc(); + if (supersededPerDoc != null) { + supersededPerDoc.remove(ev.getDocumentId()); + } + } + break; + } + } + } + } + + private MongoCollection getCollection(String collectionName) { + return collectionsMap.computeIfAbsent( + collectionName, k -> mongoClient.getDatabase(database).getCollection(collectionName)); + } + + private void drainQueues(ProcessContext context) { + drainToOutput( + (tag, value) -> context.output((TupleTag) tag, value), + (tag, value) -> + context.output( + (TupleTag>) + tag, + value)); + } + + private void drainQueuesFinishBundle(FinishBundleContext context) { + drainToOutput( + (tag, value) -> { + long tsSeconds = value.getTimestampSeconds(); + Instant timestamp = + tsSeconds > 0 ? Instant.ofEpochMilli(tsSeconds * 1000L) : Instant.now(); + context.output( + (TupleTag) tag, value, timestamp, GlobalWindow.INSTANCE); + }, + (tag, value) -> { + MongoDbChangeEventContext orig = value.getOriginalPayload(); + long tsSeconds = (orig != null) ? orig.getTimestampSeconds() : 0; + Instant timestamp = + tsSeconds > 0 ? Instant.ofEpochMilli(tsSeconds * 1000L) : Instant.now(); + context.output( + (TupleTag>) + tag, + value, + timestamp, + GlobalWindow.INSTANCE); + }); + } + + private void drainToOutput( + java.util.function.BiConsumer, MongoDbChangeEventContext> successConsumer, + java.util.function.BiConsumer< + TupleTag, FailsafeElement> + failureConsumer) { + MongoDbChangeEventContext success; + while ((success = successQueue.poll()) != null) { + successConsumer.accept(successTag, success); + } + + FailsafeElement failure; + while ((failure = failureQueue.poll()) != null) { + failureConsumer.accept(failureTag, failure); + } + + FailsafeElement severe; + while ((severe = severeFailureQueue.poll()) != null) { + failureConsumer.accept(severeFailureTag, severe); + } + } + + @VisibleForTesting + public void setMongoClient(MongoClient client) { + this.mongoClient = client; + } + } + + // Backward compatibility alias + public static class FirestoreAsyncBulkWriterFn extends BulkWriteFn { + public FirestoreAsyncBulkWriterFn( + String connectionString, + String database, + int batchSize, + int maxConcurrentAsyncWrites, + int initialWriteRatePerWorker, + int writeRateRampUpMinutes, + int maxWriteRatePerWorker, + int maxWriteRetries, + SerializableFunction clientFactory) { + super( + connectionString, + database, + batchSize, + maxConcurrentAsyncWrites, + initialWriteRatePerWorker, + writeRateRampUpMinutes, + 5, + maxWriteRatePerWorker, + maxWriteRetries, + 3, + clientFactory, + SUCCESSFUL_WRITE_TAG, + FAILED_WRITE_TAG, + SEVERE_FAILED_WRITE_TAG); + } + + @VisibleForTesting + public FirestoreAsyncBulkWriterFn( + MongoClient mongoClient, + String databaseName, + int batchSize, + int maxConcurrentAsyncWrites, + int initialWriteRatePerWorker, + int writeRateRampUpMinutes, + int maxWriteRatePerWorker, + int maxWriteRetries) { + super( + "mongodb://localhost:27017", + databaseName, + batchSize, + maxConcurrentAsyncWrites, + initialWriteRatePerWorker, + writeRateRampUpMinutes, + 5, + maxWriteRatePerWorker, + maxWriteRetries, + 3, + new StubMongoClientFactory(mongoClient), + SUCCESSFUL_WRITE_TAG, + FAILED_WRITE_TAG, + SEVERE_FAILED_WRITE_TAG); + setMongoClient(mongoClient); + } + } + + private static class StubMongoClientFactory + implements SerializableFunction, Serializable { + private final transient MongoClient client; + + StubMongoClientFactory(MongoClient client) { + this.client = client; + } + + @Override + public MongoClient apply(String input) { + return client; + } + } +} diff --git a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbChangeEventContextCoder.java b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbChangeEventContextCoder.java new file mode 100644 index 0000000000..48da195762 --- /dev/null +++ b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbChangeEventContextCoder.java @@ -0,0 +1,85 @@ +/* + * 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.google.cloud.teleport.v2.templates.datastream.MongoDbChangeEventContext; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import org.apache.beam.sdk.coders.AtomicCoder; +import org.apache.beam.sdk.coders.BooleanCoder; +import org.apache.beam.sdk.coders.NullableCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; + +/** + * Deterministic binary coder for {@link MongoDbChangeEventContext}. + * + *

Encodes MongoDB change event context without Java reflection serialization overhead, + * preserving current and original payloads (for UDF transformations and DLQ auditing), collection + * names, DLQ reconsumption flags, and retry counters. + */ +public class MongoDbChangeEventContextCoder extends AtomicCoder { + + private static final MongoDbChangeEventContextCoder INSTANCE = + new MongoDbChangeEventContextCoder(); + private static final NullableCoder STRING_CODER = NullableCoder.of(StringUtf8Coder.of()); + private static final BooleanCoder BOOLEAN_CODER = BooleanCoder.of(); + private static final VarIntCoder VARINT_CODER = VarIntCoder.of(); + + private MongoDbChangeEventContextCoder() {} + + public static MongoDbChangeEventContextCoder of() { + return INSTANCE; + } + + @Override + public void encode(MongoDbChangeEventContext value, OutputStream outStream) throws IOException { + if (value == null) { + BOOLEAN_CODER.encode(false, outStream); + return; + } + BOOLEAN_CODER.encode(true, outStream); + STRING_CODER.encode(value.getChangeEventJsonString(), outStream); + STRING_CODER.encode(value.getOriginalChangeEventJsonString(), outStream); + STRING_CODER.encode(value.getShadowCollectionPrefix(), outStream); + BOOLEAN_CODER.encode(value.getIsDlqReconsumed(), outStream); + VARINT_CODER.encode(value.getRetryCount(), outStream); + } + + @Override + public MongoDbChangeEventContext decode(InputStream inStream) throws IOException { + boolean isPresent = BOOLEAN_CODER.decode(inStream); + if (!isPresent) { + return null; + } + String changeEventJson = STRING_CODER.decode(inStream); + String originalChangeEventJson = STRING_CODER.decode(inStream); + String shadowPrefix = STRING_CODER.decode(inStream); + boolean isDlq = BOOLEAN_CODER.decode(inStream); + int retryCount = VARINT_CODER.decode(inStream); + + return MongoDbChangeEventContext.reconstitute( + changeEventJson, originalChangeEventJson, shadowPrefix, isDlq, retryCount); + } + + @Override + public void verifyDeterministic() throws NonDeterministicException { + STRING_CODER.verifyDeterministic(); + BOOLEAN_CODER.verifyDeterministic(); + VARINT_CODER.verifyDeterministic(); + } +} diff --git a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/ProcessChangeEventFn.java b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/ProcessChangeEventFn.java index d16fecca0d..b12be33841 100644 --- a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/ProcessChangeEventFn.java +++ b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/ProcessChangeEventFn.java @@ -52,14 +52,16 @@ public class ProcessChangeEventFn private final int maxRetries = 3; // Maximum number of retry attempts private final long retryDelayMs = 1000; // Initial delay in milliseconds - public static TupleTag successfulWriteTag = + public static final TupleTag SUCCESSFUL_WRITE_TAG = new TupleTag<>("successfulWrite"); - public static TupleTag> - failedWriteTag = new TupleTag<>("failedWrite"); + public static final TupleTag< + FailsafeElement> + FAILED_WRITE_TAG = new TupleTag<>("failedWrite"); // Tag for severe failures that should not be retried (e.g. permanent errors or non-transient // transaction errors) - public static TupleTag> - severeFailedWriteTag = new TupleTag<>("severeFailedWrite"); + public static final TupleTag< + FailsafeElement> + SEVERE_FAILED_WRITE_TAG = new TupleTag<>("severeFailedWrite"); // Error code 2 corresponds to BadValue/InvalidArgument, which is treated as a permanent error. public static final int INVALID_ARGUMENT = 2; @@ -181,7 +183,7 @@ public void processElement(ProcessContext context, MultiOutputReceiver out) { outOfOrderSkips.inc(); } session.commitTransaction(); - out.get(successfulWriteTag).output(element); + out.get(SUCCESSFUL_WRITE_TAG).output(element); break; // Exit the retry loop on success } catch (Exception e) { lastException = e; @@ -217,7 +219,7 @@ public void processElement(ProcessContext context, MultiOutputReceiver out) { FailsafeElement.of(element, element); failedElement.setErrorMessage(e.getMessage()); failedElement.setStacktrace(Throwables.getStackTraceAsString(e)); - out.get(severeFailedWriteTag).output(failedElement); + out.get(SEVERE_FAILED_WRITE_TAG).output(failedElement); String errorIdentifier = "UnknownError"; if (e instanceof MongoWriteException writeException) { @@ -264,7 +266,7 @@ public void processElement(ProcessContext context, MultiOutputReceiver out) { FailsafeElement.of(element, element); failedElement.setErrorMessage(ie.getMessage()); failedElement.setStacktrace(Throwables.getStackTraceAsString(ie)); - out.get(failedWriteTag).output(failedElement); + out.get(FAILED_WRITE_TAG).output(failedElement); retriableFailedWrites.inc(); break; // Exit the retry loop if interrupted } @@ -282,7 +284,7 @@ public void processElement(ProcessContext context, MultiOutputReceiver out) { FailsafeElement.of(element, element); failedElement.setErrorMessage(e.getMessage()); failedElement.setStacktrace(Throwables.getStackTraceAsString(e)); - out.get(failedWriteTag).output(failedElement); + out.get(FAILED_WRITE_TAG).output(failedElement); retriableFailedWrites.inc(); LOG.info( "Failed element of id {} sent to retry DLQ after {} attempts", diff --git a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/StatefulDeduplicationFn.java b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/StatefulDeduplicationFn.java new file mode 100644 index 0000000000..f43c7b4cb6 --- /dev/null +++ b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/StatefulDeduplicationFn.java @@ -0,0 +1,113 @@ +/* + * 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.google.cloud.teleport.v2.templates.datastream.MongoDbChangeEventContext; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.metrics.Metrics; +import org.apache.beam.sdk.state.StateSpec; +import org.apache.beam.sdk.state.StateSpecs; +import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.values.KV; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Stateful DoFn that ensures monotonic timestamp ordering per document ID and drops out-of-order + * events in memory using Dataflow Streaming Engine state without database round-trips. + */ +public class StatefulDeduplicationFn + extends DoFn, MongoDbChangeEventContext> { + + private static final Logger LOG = LoggerFactory.getLogger(StatefulDeduplicationFn.class); + + @StateId("latestTimestamp") + private final StateSpec> latestTimestampSpec = + StateSpecs.value(TimestampSortKeyCoder.of()); + + private final Counter outOfOrderSkips = + Metrics.counter(StatefulDeduplicationFn.class, "outOfOrderSkips"); + private final Counter dedupOutputs = + Metrics.counter(StatefulDeduplicationFn.class, "dedupOutputs"); + private final Counter dlqEqualTimestampPassThrough = + Metrics.counter(StatefulDeduplicationFn.class, "dlqEqualTimestampPassThrough"); + + private final ThrottledLogger throttledLogger = new ThrottledLogger(30000L); + private transient java.util.Map bundleStateCache; + + @StartBundle + public void startBundle() { + bundleStateCache = new java.util.HashMap<>(); + } + + @FinishBundle + public void finishBundle() { + if (bundleStateCache != null) { + bundleStateCache.clear(); + } + } + + @ProcessElement + public void processElement( + ProcessContext context, + @StateId("latestTimestamp") ValueState latestTimestampState, + OutputReceiver out) { + KV element = context.element(); + if (element == null || element.getValue() == null) { + return; + } + + MongoDbChangeEventContext event = element.getValue(); + TimestampSortKey currentSortKey = TimestampSortKey.of(event); + + String key = element.getKey(); + TimestampSortKey latestSortKey = + (bundleStateCache != null && key != null) ? bundleStateCache.get(key) : null; + if (latestSortKey == null) { + latestSortKey = latestTimestampState.read(); + } + + int cmp = (latestSortKey == null) ? 1 : currentSortKey.compareTo(latestSortKey); + + if (cmp > 0) { + latestTimestampState.write(currentSortKey); + if (bundleStateCache != null && key != null) { + bundleStateCache.put(key, currentSortKey); + } + out.output(event); + dedupOutputs.inc(); + } else if (cmp == 0 && event.getIsDlqReconsumed()) { + // Reconsumed DLQ events with identical timestamp are allowed to pass through + out.output(event); + dedupOutputs.inc(); + dlqEqualTimestampPassThrough.inc(); + } else { + // Stale / out-of-order event - drop immediately without database RPCs + if (bundleStateCache != null && key != null && latestSortKey != null) { + bundleStateCache.put(key, latestSortKey); + } + outOfOrderSkips.inc(); + throttledLogger.logInfo( + LOG, + event.getDataCollection(), + "Dropped out-of-order event for docId: {}, currentSortKey: {}, latestSortKey: {}", + event.getDocumentId(), + currentSortKey, + latestSortKey); + } + } +} diff --git a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/ThrottledLogger.java b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/ThrottledLogger.java new file mode 100644 index 0000000000..ab8167bebd --- /dev/null +++ b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/ThrottledLogger.java @@ -0,0 +1,288 @@ +/* + * 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.io.Serializable; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicLong; +import org.apache.beam.sdk.metrics.Counter; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Utility for windowed rate-limited logging to prevent flooding Cloud Logging while accurately + * capturing error counts. + * + *

Ensures that metrics counters are called unconditionally while log messages are throttled to a + * configurable window (default 30 seconds) using a bounded ConcurrentHashMap. + */ +public class ThrottledLogger implements Serializable { + + private static final Logger LOG = LoggerFactory.getLogger(ThrottledLogger.class); + public static final long DEFAULT_THROTTLE_INTERVAL_MS = 30_000L; + private static final int MAX_ERROR_CATEGORIES = 200; + + private final String componentName; + private final long throttleIntervalMs; + private transient volatile AtomicLong totalErrors; + private transient volatile AtomicLong totalRetryable; + private transient volatile AtomicLong totalSevere; + private transient volatile ConcurrentHashMap errorCategories; + private transient volatile ConcurrentHashMap logStates; + private transient volatile AtomicLong lastLogTimestamp; + + public ThrottledLogger() { + this("ThrottledLogger", DEFAULT_THROTTLE_INTERVAL_MS); + } + + public ThrottledLogger(String componentName) { + this(componentName, DEFAULT_THROTTLE_INTERVAL_MS); + } + + public ThrottledLogger(long throttleIntervalMs) { + this("ThrottledLogger", throttleIntervalMs); + } + + public ThrottledLogger(String componentName, long throttleIntervalMs) { + this.componentName = componentName != null ? componentName : "ThrottledLogger"; + this.throttleIntervalMs = + throttleIntervalMs > 0 ? throttleIntervalMs : DEFAULT_THROTTLE_INTERVAL_MS; + init(); + } + + private void init() { + this.totalRetryable = new AtomicLong(0); + this.totalSevere = new AtomicLong(0); + this.errorCategories = new ConcurrentHashMap<>(); + this.logStates = new ConcurrentHashMap<>(); + this.lastLogTimestamp = new AtomicLong(System.currentTimeMillis()); + this.totalErrors = new AtomicLong(0); + } + + private void ensureInitialized() { + if (this.totalErrors == null) { + synchronized (this) { + if (this.totalErrors == null) { + init(); + } + } + } + } + + private void readObject(java.io.ObjectInputStream in) + throws java.io.IOException, ClassNotFoundException { + in.defaultReadObject(); + init(); + } + + public void recordRetryableError(String category, String message) { + ensureInitialized(); + totalErrors.incrementAndGet(); + totalRetryable.incrementAndGet(); + incrementCategory(category); + checkAndFlush(); + } + + public void recordSevereError(String category, String message) { + ensureInitialized(); + totalErrors.incrementAndGet(); + totalSevere.incrementAndGet(); + incrementCategory(category); + checkAndFlush(); + } + + public void recordError(String category, String message) { + recordRetryableError(category, message); + } + + private void incrementCategory(String category) { + String safeCategory = category != null ? category : "UNKNOWN"; + if (errorCategories.size() >= MAX_ERROR_CATEGORIES + && !errorCategories.containsKey(safeCategory)) { + safeCategory = "OTHER"; + } + errorCategories.computeIfAbsent(safeCategory, k -> new AtomicLong(0)).incrementAndGet(); + } + + private void checkAndFlush() { + long now = System.currentTimeMillis(); + long last = lastLogTimestamp.get(); + if (now - last >= throttleIntervalMs) { + if (lastLogTimestamp.compareAndSet(last, now)) { + flushSummary(); + } + } + } + + public void flushSummary() { + ensureInitialized(); + long errors = totalErrors.getAndSet(0); + if (errors == 0) { + return; + } + long retryable = totalRetryable.getAndSet(0); + long severe = totalSevere.getAndSet(0); + StringBuilder sb = new StringBuilder(); + errorCategories.forEach( + (cat, count) -> { + long val = count.getAndSet(0); + if (val > 0) { + if (sb.length() > 0) { + sb.append(", "); + } + sb.append(cat).append("=").append(val); + } + }); + + LOG.warn( + "[{}] Error Summary (throttled): Total={}, Retryable={}, Severe={}. Breakdown: [{}]", + componentName, + errors, + retryable, + severe, + sb); + } + + public long getTotalErrors() { + ensureInitialized(); + return totalErrors.get(); + } + + public long getTotalRetryable() { + ensureInitialized(); + return totalRetryable.get(); + } + + public long getTotalSevere() { + ensureInitialized(); + return totalSevere.get(); + } + + /** Evaluates if a log message should be emitted for the key in this window. */ + public boolean shouldLog(String key) { + ensureInitialized(); + String safeKey = key != null ? key : "DEFAULT"; + if (logStates.size() >= MAX_ERROR_CATEGORIES && !logStates.containsKey(safeKey)) { + safeKey = "OTHER"; + } + LogEntryState state = logStates.computeIfAbsent(safeKey, k -> new LogEntryState(0)); + long now = System.currentTimeMillis(); + long lastTime = state.lastLoggedTimeMs.get(); + if (now - lastTime >= throttleIntervalMs) { + if (state.lastLoggedTimeMs.compareAndSet(lastTime, now)) { + return true; + } + } + state.suppressedCount.incrementAndGet(); + return false; + } + + public long getAndResetSuppressedCount(String key) { + ensureInitialized(); + String safeKey = key != null ? key : "DEFAULT"; + if (logStates.size() >= MAX_ERROR_CATEGORIES && !logStates.containsKey(safeKey)) { + safeKey = "OTHER"; + } + LogEntryState state = logStates.get(safeKey); + return state != null ? state.suppressedCount.getAndSet(0) : 0; + } + + public void logInfo(Logger logger, String key, String message, Object... args) { + if (shouldLog(key)) { + long suppressed = getAndResetSuppressedCount(key); + if (suppressed > 0) { + logger.info( + message + + " [Suppressed " + + suppressed + + " similar logs in the last " + + (throttleIntervalMs / 1000) + + "s]", + args); + } else { + logger.info(message, args); + } + } + } + + public void logWarn(Logger logger, String key, String message, Object... args) { + if (shouldLog(key)) { + long suppressed = getAndResetSuppressedCount(key); + if (suppressed > 0) { + logger.warn( + message + + " [Suppressed " + + suppressed + + " similar logs in the last " + + (throttleIntervalMs / 1000) + + "s]", + args); + } else { + logger.warn(message, args); + } + } + } + + public void logError(Logger logger, String key, String message, Object... args) { + if (shouldLog(key)) { + long suppressed = getAndResetSuppressedCount(key); + if (suppressed > 0) { + logger.error( + message + + " [Suppressed " + + suppressed + + " similar logs in the last " + + (throttleIntervalMs / 1000) + + "s]", + args); + } else { + logger.error(message, args); + } + } + } + + public void logWarn(Logger logger, Counter counter, String key, String message, Object... args) { + if (counter != null) { + counter.inc(); + } + logWarn(logger, key, message, args); + } + + public void logError(Logger logger, Counter counter, String key, String message, Object... args) { + if (counter != null) { + counter.inc(); + } + logError(logger, key, message, args); + } + + public static class LogEntryState implements Serializable { + private final AtomicLong lastLoggedTimeMs; + private final AtomicLong suppressedCount; + + public LogEntryState(long initialTimeMs) { + this.lastLoggedTimeMs = new AtomicLong(initialTimeMs); + this.suppressedCount = new AtomicLong(0); + } + + public AtomicLong getLastLoggedTimeMs() { + return lastLoggedTimeMs; + } + + public AtomicLong getSuppressedCount() { + return suppressedCount; + } + } +} diff --git a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/TimestampSortKey.java b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/TimestampSortKey.java new file mode 100644 index 0000000000..99b24fe22e --- /dev/null +++ b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/TimestampSortKey.java @@ -0,0 +1,130 @@ +/* + * 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.google.cloud.teleport.v2.templates.datastream.MongoDbChangeEventContext; +import java.io.Serializable; +import java.util.Objects; +import org.apache.beam.sdk.coders.DefaultCoder; + +/** + * Composite monotonic sort key for Datastream MongoDB change events. + * + *

Disentangles cross-domain timestamp comparisons by ordering: + * + *

    + *
  1. Epoch seconds (MongoDB oplog timestamp seconds / backfill extraction seconds). + *
  2. Stream type precedence: Live CDC mutations (INSERT, UPDATE, DELETE) strictly supersede + * Backfill snapshot READ events within the same second. + *
  3. Sub-second ordering within the same stream type: + *
      + *
    • For CDC: MongoDB oplog increment counter. + *
    • For Backfill: Snapshot extraction wall-clock nanoseconds. + *
    + *
+ */ +@DefaultCoder(TimestampSortKeyCoder.class) +public class TimestampSortKey implements Serializable, Comparable { + + private final long seconds; + private final long subSeconds; + private final boolean isCdc; + + public TimestampSortKey(long seconds, long subSeconds, boolean isCdc) { + this.seconds = seconds; + this.subSeconds = subSeconds; + this.isCdc = isCdc; + } + + public static TimestampSortKey of(MongoDbChangeEventContext event) { + if (event == null) { + return null; + } + return new TimestampSortKey( + event.getTimestampSeconds(), event.getTimestampSubSeconds(), event.isCdcEvent()); + } + + public static TimestampSortKey of(long seconds, long subSeconds, boolean isCdc) { + return new TimestampSortKey(seconds, subSeconds, isCdc); + } + + public long getSeconds() { + return seconds; + } + + public long getTimestampSeconds() { + return seconds; + } + + public long getSubSeconds() { + return subSeconds; + } + + public long getTimestampSubSeconds() { + return subSeconds; + } + + public boolean isCdc() { + return isCdc; + } + + public boolean getIsCdc() { + return isCdc; + } + + @Override + public int compareTo(TimestampSortKey other) { + if (other == null) { + throw new NullPointerException("Cannot compare TimestampSortKey with null"); + } + // 1. Primary: Compare epoch seconds + if (this.seconds != other.seconds) { + return Long.compare(this.seconds, other.seconds); + } + // 2. Stream type precedence: Live CDC strictly supersedes Backfill snapshot within the same + // second + if (this.isCdc && !other.isCdc) { + return 1; + } + if (!this.isCdc && other.isCdc) { + return -1; + } + // 3. Sub-second ordering within the same stream type (nanos for backfill, oplog inc for CDC) + return Long.compare(this.subSeconds, other.subSeconds); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof TimestampSortKey)) { + return false; + } + TimestampSortKey that = (TimestampSortKey) o; + return seconds == that.seconds && subSeconds == that.subSeconds && isCdc == that.isCdc; + } + + @Override + public int hashCode() { + return Objects.hash(seconds, subSeconds, isCdc); + } + + @Override + public String toString() { + return seconds + ":" + subSeconds + ":" + (isCdc ? "cdc" : "backfill"); + } +} diff --git a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/TimestampSortKeyCoder.java b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/TimestampSortKeyCoder.java new file mode 100644 index 0000000000..186bbdfdfd --- /dev/null +++ b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/TimestampSortKeyCoder.java @@ -0,0 +1,78 @@ +/* + * 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.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import org.apache.beam.sdk.coders.AtomicCoder; +import org.apache.beam.sdk.coders.BigEndianLongCoder; +import org.apache.beam.sdk.coders.BooleanCoder; + +/** + * Deterministic binary coder for {@link TimestampSortKey}. + * + *

Encodes sort keys into a compact binary format: + * + *

    + *
  • Presence flag (1 byte boolean) + *
  • Epoch seconds (8 bytes via {@link BigEndianLongCoder}) + *
  • Sub-second ordering / nanoseconds (8 bytes via {@link BigEndianLongCoder}) + *
  • Stream type / isCdc flag (1 byte via {@link BooleanCoder}) + *
+ */ +public class TimestampSortKeyCoder extends AtomicCoder { + + private static final TimestampSortKeyCoder INSTANCE = new TimestampSortKeyCoder(); + private static final BigEndianLongCoder LONG_CODER = BigEndianLongCoder.of(); + private static final BooleanCoder BOOLEAN_CODER = BooleanCoder.of(); + + private TimestampSortKeyCoder() {} + + public static TimestampSortKeyCoder of() { + return INSTANCE; + } + + @Override + public void encode(TimestampSortKey value, OutputStream outStream) throws IOException { + if (value == null) { + BOOLEAN_CODER.encode(false, outStream); + return; + } + BOOLEAN_CODER.encode(true, outStream); + LONG_CODER.encode(value.getSeconds(), outStream); + LONG_CODER.encode(value.getSubSeconds(), outStream); + BOOLEAN_CODER.encode(value.isCdc(), outStream); + } + + @Override + public TimestampSortKey decode(InputStream inStream) throws IOException { + boolean isPresent = BOOLEAN_CODER.decode(inStream); + if (!isPresent) { + return null; + } + long seconds = LONG_CODER.decode(inStream); + long subSeconds = LONG_CODER.decode(inStream); + boolean isCdc = BOOLEAN_CODER.decode(inStream); + return TimestampSortKey.of(seconds, subSeconds, isCdc); + } + + @Override + public void verifyDeterministic() throws NonDeterministicException { + LONG_CODER.verifyDeterministic(); + BOOLEAN_CODER.verifyDeterministic(); + } +} diff --git a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/Utils.java b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/Utils.java index c3def2c86b..6fec8051f8 100644 --- a/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/Utils.java +++ b/v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/Utils.java @@ -21,21 +21,25 @@ import com.google.cloud.teleport.v2.templates.datastream.DatastreamConstants; import com.google.cloud.teleport.v2.templates.datastream.MongoDbChangeEventContext; import java.util.Base64; +import java.util.List; import java.util.Set; import org.bson.Document; import org.bson.json.JsonMode; import org.bson.json.JsonWriterSettings; import org.bson.types.Binary; +import org.bson.types.ObjectId; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -/** Utils used by the Datastream-mongodb-to-mongodb pipeline. */ +/** Utils used by the Datastream-mongodb-to-firestore pipeline. */ public final class Utils { private static final Logger LOG = LoggerFactory.getLogger(Utils.class); private static final JsonWriterSettings CANONICAL_JSON_SETTINGS = JsonWriterSettings.builder().outputMode(JsonMode.EXTENDED).build(); + private Utils() {} + public static void removeTableRowFields(Document doc, Set ignoreFields) { for (String ignoreField : ignoreFields) { doc.remove(ignoreField); @@ -44,41 +48,146 @@ public static void removeTableRowFields(Document doc, Set ignoreFields) /* Whether the first timestamp is later than the second timestamp. */ public static boolean isNewerTimestamp(Document ts1, Document ts2) { - long s1 = ts1.getLong(MongoDbChangeEventContext.TIMESTAMP_SECONDS_COL); - int n1 = ts1.getInteger(MongoDbChangeEventContext.TIMESTAMP_NANOS_COL); - long s2 = ts2.getLong(MongoDbChangeEventContext.TIMESTAMP_SECONDS_COL); - int n2 = ts2.getInteger(MongoDbChangeEventContext.TIMESTAMP_NANOS_COL); + if (ts1 == null) { + return false; + } + if (ts2 == null) { + return true; + } + long s1 = 0L; + int n1 = 0; + if (ts1.containsKey(MongoDbChangeEventContext.TIMESTAMP_SECONDS_COL)) { + Object s = ts1.get(MongoDbChangeEventContext.TIMESTAMP_SECONDS_COL); + if (s instanceof Number) { + s1 = ((Number) s).longValue(); + } + } + if (ts1.containsKey(MongoDbChangeEventContext.TIMESTAMP_NANOS_COL)) { + Object n = ts1.get(MongoDbChangeEventContext.TIMESTAMP_NANOS_COL); + if (n instanceof Number) { + n1 = ((Number) n).intValue(); + } + } + long s2 = 0L; + int n2 = 0; + if (ts2.containsKey(MongoDbChangeEventContext.TIMESTAMP_SECONDS_COL)) { + Object s = ts2.get(MongoDbChangeEventContext.TIMESTAMP_SECONDS_COL); + if (s instanceof Number) { + s2 = ((Number) s).longValue(); + } + } + if (ts2.containsKey(MongoDbChangeEventContext.TIMESTAMP_NANOS_COL)) { + Object n = ts2.get(MongoDbChangeEventContext.TIMESTAMP_NANOS_COL); + if (n instanceof Number) { + n2 = ((Number) n).intValue(); + } + } return s1 > s2 || (s1 == s2 && n1 > n2); } + public static long getTimestampNanos(Document timestampDoc) { + if (timestampDoc == null) { + return 0L; + } + long seconds = 0L; + if (timestampDoc.containsKey(MongoDbChangeEventContext.TIMESTAMP_SECONDS_COL)) { + Object s = timestampDoc.get(MongoDbChangeEventContext.TIMESTAMP_SECONDS_COL); + if (s instanceof Number) { + seconds = ((Number) s).longValue(); + } + } + long nanos = 0L; + if (timestampDoc.containsKey(MongoDbChangeEventContext.TIMESTAMP_NANOS_COL)) { + Object n = timestampDoc.get(MongoDbChangeEventContext.TIMESTAMP_NANOS_COL); + if (n instanceof Number) { + nanos = ((Number) n).longValue(); + } + } + return (seconds * 1_000_000_000L) + nanos; + } + public static Document jsonToDocument(String jsonString, Object documentId) { - Document rawDoc; + if (jsonString == null) { + return null; + } + Document rawDoc = null; try { - rawDoc = Document.parse(Document.parse(jsonString).get(DATA_COL).toString()); + Document parsed = Document.parse(jsonString); + if (parsed.containsKey(DATA_COL)) { + Object dataObj = parsed.get(DATA_COL); + if (dataObj instanceof Document) { + rawDoc = (Document) dataObj; + } else if (dataObj instanceof String) { + rawDoc = Document.parse((String) dataObj); + } else if (dataObj != null) { + rawDoc = Document.parse(dataObj.toString()); + } + } else { + // No 'data' wrapper field; the parsed document itself is the payload + rawDoc = parsed; + } } catch (Exception ex) { - LOG.info( - "Document parsing for {} failed due to {}, try casting.", jsonString, ex.getMessage()); - rawDoc = (Document) Document.parse(jsonString).get(DATA_COL); + LOG.debug("Document parsing for {} failed due to {}.", jsonString, ex.getMessage()); } if (rawDoc == null) { return null; } + removeTableRowFields( + rawDoc, + com.google.cloud.teleport.v2.templates.datastream.DatastreamConstants + .DATASTREAM_METADATA_FIELDS); rawDoc.put(MongoDbChangeEventContext.DOC_ID_COL, documentId); return rawDoc; } + /** + * Converts a MongoDB document ID into a type-tagged, collision-free string representation. + * + *

NOTE: This method does NOT generate a semantically equivalent string for database + * writes, and must NEVER be used as the destination document's {@code _id} value (which should + * retain native BSON types such as {@link org.bson.types.ObjectId}, {@link Document}, or {@link + * org.bson.types.Binary}). + * + *

Use Case: This method is strictly intended for generating internal Apache Beam + * grouping and shuffling keys (e.g. {@code collection + "#" + documentIdToString(docId)}) and + * diagnostic string logs. The type prefix (e.g. {@code str_}, {@code i64_}, {@code bin__}) + * ensures distinct BSON types with identical string forms (such as string {@code "123"} vs Long + * {@code 123L}) never collide in Beam's stateful deduplication and windowing operations. + * + * @param documentId the raw BSON document ID object + * @return a type-tagged string representation suitable for pipeline routing keys + */ public static String documentIdToString(Object documentId) { if (documentId == null) { return "null"; } if (documentId instanceof Binary) { Binary binary = (Binary) documentId; - return Base64.getEncoder().encodeToString(binary.getData()); + return "bin_" + binary.getType() + "_" + Base64.getEncoder().encodeToString(binary.getData()); } if (documentId instanceof Document) { - return ((Document) documentId).toJson(); + return "doc_" + ((Document) documentId).toJson(CANONICAL_JSON_SETTINGS); + } + if (documentId instanceof List) { + Document wrapper = new Document("arr", documentId); + return "list_" + wrapper.toJson(CANONICAL_JSON_SETTINGS); + } + if (documentId instanceof ObjectId) { + return "oid_" + ((ObjectId) documentId).toHexString(); + } + if (documentId instanceof Long) { + return "i64_" + documentId; + } + if (documentId instanceof Integer) { + return "i32_" + documentId; } - return documentId.toString(); + if (documentId instanceof Double) { + return "f64_" + documentId; + } + if (documentId instanceof Boolean) { + return "bool_" + documentId; + } + return "str_" + documentId.toString(); } public static String getCanonicalJsonOfDataField(Document fullEvent) { @@ -100,12 +209,18 @@ public static String getCanonicalJsonOfDataField(String jsonString) { } public static Document extractInnerEvent(Document doc) { + if (doc == null) { + return null; + } return doc.containsKey(DatastreamConstants.CHANGE_EVENT) ? (Document) doc.get(DatastreamConstants.CHANGE_EVENT) : doc; } public static JsonNode extractInnerEvent(JsonNode payload) { + if (payload == null) { + return null; + } return payload.has(DatastreamConstants.CHANGE_EVENT) ? payload.get(DatastreamConstants.CHANGE_EVENT) : payload; diff --git a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/templates/DataStreamMongoDBToFirestoreTest.java b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/templates/DataStreamMongoDBToFirestoreTest.java index 7c98736087..4637e0b4db 100644 --- a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/templates/DataStreamMongoDBToFirestoreTest.java +++ b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/templates/DataStreamMongoDBToFirestoreTest.java @@ -56,6 +56,237 @@ public final class DataStreamMongoDBToFirestoreTest { private static final TupleTag> SUCCESS_TAG = new TupleTag>() {}; + @Test + public void inputArgs_shadowlessDefaults() { + String[] args = new String[] {"--inputFilePattern=gs://test-bkt/"}; + DataStreamMongoDBToFirestore.Options options = + PipelineOptionsFactory.fromArgs(args) + .withValidation() + .as(DataStreamMongoDBToFirestore.Options.class); + + assertFalse(options.getUseShadowTables()); + assertEquals(Integer.valueOf(500), options.getBatchSize()); + assertEquals(Integer.valueOf(10), options.getMaxConcurrentAsyncWrites()); + assertEquals(Integer.valueOf(500), options.getInitialWriteRatePerWorker()); + assertEquals(Integer.valueOf(5), options.getWriteRateRampUpMinutes()); + assertEquals(Integer.valueOf(2500), options.getMaxWriteRatePerWorker()); + } + + @Test + public void inputArgs_customShadowlessOptions() { + String[] args = + new String[] { + "--useShadowTables=false", + "--batchSize=200", + "--maxConcurrentAsyncWrites=20", + "--initialWriteRatePerWorker=1000", + "--writeRateRampUpMinutes=10", + "--maxWriteRatePerWorker=5000" + }; + DataStreamMongoDBToFirestore.Options options = + PipelineOptionsFactory.fromArgs(args) + .withValidation() + .as(DataStreamMongoDBToFirestore.Options.class); + + assertFalse(options.getUseShadowTables()); + assertEquals(Integer.valueOf(200), options.getBatchSize()); + assertEquals(Integer.valueOf(20), options.getMaxConcurrentAsyncWrites()); + assertEquals(Integer.valueOf(1000), options.getInitialWriteRatePerWorker()); + assertEquals(Integer.valueOf(10), options.getWriteRateRampUpMinutes()); + assertEquals(Integer.valueOf(5000), options.getMaxWriteRatePerWorker()); + } + + @Test + public void validateOptions_validShadowlessOptions() { + String[] args = + new String[] { + "--inputFilePattern=gs://test-bkt/", + "--connectionUri=mongodb://localhost:27017", + "--batchSize=500", + "--maxConcurrentAsyncWrites=10", + "--initialWriteRatePerWorker=500", + "--writeRateRampUpMinutes=5", + "--maxWriteRatePerWorker=2500" + }; + DataStreamMongoDBToFirestore.Options options = + PipelineOptionsFactory.fromArgs(args) + .withValidation() + .as(DataStreamMongoDBToFirestore.Options.class); + + DataStreamMongoDBToFirestore.validateOptions(options); + } + + @Test + public void validateOptions_missingConnectionUri_throwsException() { + String[] args = new String[] {"--inputFilePattern=gs://test-bkt/"}; + DataStreamMongoDBToFirestore.Options options = + PipelineOptionsFactory.fromArgs(args) + .withValidation() + .as(DataStreamMongoDBToFirestore.Options.class); + + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, + () -> DataStreamMongoDBToFirestore.validateOptions(options)); + assertTrue(thrown.getMessage().contains("Connection URI (connectionUri) must be specified")); + } + + @Test + public void validateOptions_emptyConnectionUri_throwsException() { + String[] args = new String[] {"--inputFilePattern=gs://test-bkt/", "--connectionUri= "}; + DataStreamMongoDBToFirestore.Options options = + PipelineOptionsFactory.fromArgs(args) + .withValidation() + .as(DataStreamMongoDBToFirestore.Options.class); + + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, + () -> DataStreamMongoDBToFirestore.validateOptions(options)); + assertTrue(thrown.getMessage().contains("Connection URI (connectionUri) must be specified")); + } + + @Test + public void validateOptions_invalidScheme_throwsException() { + String[] args = + new String[] { + "--inputFilePattern=gs://test-bkt/", "--connectionUri=http://localhost:27017" + }; + DataStreamMongoDBToFirestore.Options options = + PipelineOptionsFactory.fromArgs(args) + .withValidation() + .as(DataStreamMongoDBToFirestore.Options.class); + + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, + () -> DataStreamMongoDBToFirestore.validateOptions(options)); + assertTrue(thrown.getMessage().contains("Must start with 'mongodb://' or 'mongodb+srv://'")); + } + + @Test + public void validateOptions_validConnectionUri_mongodb_success() { + String[] args = + new String[] { + "--inputFilePattern=gs://test-bkt/", + "--connectionUri=mongodb://user:pass@localhost:27017/db" + }; + DataStreamMongoDBToFirestore.Options options = + PipelineOptionsFactory.fromArgs(args) + .withValidation() + .as(DataStreamMongoDBToFirestore.Options.class); + + DataStreamMongoDBToFirestore.validateOptions(options); + } + + @Test + public void validateOptions_validConnectionUri_mongodbSrv_success() { + String[] args = + new String[] { + "--inputFilePattern=gs://test-bkt/", + "--connectionUri=mongodb+srv://cluster.example.com/db" + }; + DataStreamMongoDBToFirestore.Options options = + PipelineOptionsFactory.fromArgs(args) + .withValidation() + .as(DataStreamMongoDBToFirestore.Options.class); + + DataStreamMongoDBToFirestore.validateOptions(options); + } + + @Test + public void validateOptions_emptyDatabaseName_throwsException() { + String[] args = + new String[] { + "--inputFilePattern=gs://test-bkt/", + "--connectionUri=mongodb://localhost:27017", + "--databaseName= " + }; + DataStreamMongoDBToFirestore.Options options = + PipelineOptionsFactory.fromArgs(args) + .withValidation() + .as(DataStreamMongoDBToFirestore.Options.class); + + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, + () -> DataStreamMongoDBToFirestore.validateOptions(options)); + assertTrue(thrown.getMessage().contains("Database name (databaseName) must be specified")); + } + + @Test + public void validateOptions_invalidBatchSize_throwsException() { + String[] args = + new String[] { + "--inputFilePattern=gs://test-bkt/", + "--connectionUri=mongodb://localhost:27017", + "--batchSize=0" + }; + DataStreamMongoDBToFirestore.Options options = + PipelineOptionsFactory.fromArgs(args) + .withValidation() + .as(DataStreamMongoDBToFirestore.Options.class); + + assertThrows( + IllegalArgumentException.class, + () -> DataStreamMongoDBToFirestore.validateOptions(options)); + } + + @Test + public void validateOptions_invalidMaxConcurrentAsyncWrites_throwsException() { + String[] args = + new String[] { + "--inputFilePattern=gs://test-bkt/", + "--connectionUri=mongodb://localhost:27017", + "--maxConcurrentAsyncWrites=-1" + }; + DataStreamMongoDBToFirestore.Options options = + PipelineOptionsFactory.fromArgs(args) + .withValidation() + .as(DataStreamMongoDBToFirestore.Options.class); + + assertThrows( + IllegalArgumentException.class, + () -> DataStreamMongoDBToFirestore.validateOptions(options)); + } + + @Test + public void validateOptions_invalidRateRampUpBounds_throwsException() { + String[] args = + new String[] { + "--inputFilePattern=gs://test-bkt/", + "--connectionUri=mongodb://localhost:27017", + "--initialWriteRatePerWorker=5000", + "--maxWriteRatePerWorker=2000" + }; + DataStreamMongoDBToFirestore.Options options = + PipelineOptionsFactory.fromArgs(args) + .withValidation() + .as(DataStreamMongoDBToFirestore.Options.class); + + assertThrows( + IllegalArgumentException.class, + () -> DataStreamMongoDBToFirestore.validateOptions(options)); + } + + @Test + public void validateOptions_invalidRampUpMinutes_throwsException() { + String[] args = + new String[] { + "--inputFilePattern=gs://test-bkt/", + "--connectionUri=mongodb://localhost:27017", + "--writeRateRampUpMinutes=-1" + }; + DataStreamMongoDBToFirestore.Options options = + PipelineOptionsFactory.fromArgs(args) + .withValidation() + .as(DataStreamMongoDBToFirestore.Options.class); + + assertThrows( + IllegalArgumentException.class, + () -> DataStreamMongoDBToFirestore.validateOptions(options)); + } + @Test public void inputArgs_inputFilePattern() { String[] args = new String[] {"--inputFilePattern=gs://test-bkt/"}; diff --git a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/templates/ProcessBackfillEventFnTest.java b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/templates/ProcessBackfillEventFnTest.java index f6ac101b00..e95af1e8db 100644 --- a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/templates/ProcessBackfillEventFnTest.java +++ b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/templates/ProcessBackfillEventFnTest.java @@ -69,11 +69,12 @@ public void setUp() { when(mockClient.getDatabase(DATABASE_NAME)).thenReturn(mockDatabase); when(mockDatabase.getCollection(COLLECTION_NAME)).thenReturn(mockCollection); - when(mockReceiver.get(DataStreamMongoDBToFirestore.ProcessBackfillEventFn.successfulWriteTag)) + when(mockReceiver.get(DataStreamMongoDBToFirestore.ProcessBackfillEventFn.SUCCESSFUL_WRITE_TAG)) .thenReturn(mockSuccessReceiver); - when(mockReceiver.get(DataStreamMongoDBToFirestore.ProcessBackfillEventFn.failedWriteTag)) + when(mockReceiver.get(DataStreamMongoDBToFirestore.ProcessBackfillEventFn.FAILED_WRITE_TAG)) .thenReturn(mockFailureReceiver); - when(mockReceiver.get(DataStreamMongoDBToFirestore.ProcessBackfillEventFn.severeFailedWriteTag)) + when(mockReceiver.get( + DataStreamMongoDBToFirestore.ProcessBackfillEventFn.SEVERE_FAILED_WRITE_TAG)) .thenReturn(mockSevereFailureReceiver); fn = @@ -239,7 +240,7 @@ public void testFinishBundle_processesRemainingEvents() { // Verify output in finishBundle verify(mockFinishBundleContext, times(1)) .output( - eq(DataStreamMongoDBToFirestore.ProcessBackfillEventFn.successfulWriteTag), + eq(DataStreamMongoDBToFirestore.ProcessBackfillEventFn.SUCCESSFUL_WRITE_TAG), eq(event1), any(org.joda.time.Instant.class), any(org.apache.beam.sdk.transforms.windowing.GlobalWindow.class)); diff --git a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/templates/datastream/MongoDbChangeEventContextTest.java b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/templates/datastream/MongoDbChangeEventContextTest.java index 76029ef81d..c6c74cb5d9 100644 --- a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/templates/datastream/MongoDbChangeEventContextTest.java +++ b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/templates/datastream/MongoDbChangeEventContextTest.java @@ -32,6 +32,7 @@ import com.fasterxml.jackson.databind.node.ObjectNode; import com.google.cloud.teleport.v2.transforms.Utils; import com.google.common.collect.ImmutableMap; +import java.util.List; import org.bson.Document; import org.bson.types.Binary; import org.bson.types.ObjectId; @@ -130,7 +131,7 @@ public void setUp() throws JsonProcessingException { "_metadata_source": { "collection": "test_collection" }, - "_id": "[1, 2, 3]", + "_id": true, "_metadata_timestamp_seconds": 1683782270, "_metadata_timestamp_nanos": 123456789 }\ @@ -412,8 +413,8 @@ public void testConstructorBinaryId() throws JsonProcessingException { assertEquals(0, binaryId.getType()); } - @Test(expected = IllegalArgumentException.class) - public void testConstructorArrayIdThrows() throws JsonProcessingException { + @Test + public void testConstructorArrayId() throws JsonProcessingException { JsonNode eventWithArrayId = OBJECT_MAPPER.readTree( """ @@ -429,7 +430,15 @@ public void testConstructorArrayIdThrows() throws JsonProcessingException { "_metadata_timestamp_nanos": 123456789 }\ """); - new MongoDbChangeEventContext(eventWithArrayId, SHADOW_PREFIX); + MongoDbChangeEventContext context = + new MongoDbChangeEventContext(eventWithArrayId, SHADOW_PREFIX); + assertNotNull(context.getDocumentId()); + assertTrue(context.getDocumentId() instanceof List); + List list = (List) context.getDocumentId(); + assertEquals(3, list.size()); + assertEquals(1, list.get(0)); + assertEquals(2, list.get(1)); + assertEquals(3, list.get(2)); } @Test(expected = IllegalArgumentException.class) diff --git a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/CreateMongoDbChangeEventContextFnTest.java b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/CreateMongoDbChangeEventContextFnTest.java index d78f179537..2f1fdfa0fe 100644 --- a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/CreateMongoDbChangeEventContextFnTest.java +++ b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/CreateMongoDbChangeEventContextFnTest.java @@ -60,9 +60,9 @@ public void setUp() throws Exception { mockFailureReceiver = mock(OutputReceiver.class); // Stub the get() method of MultiOutputReceiver - when(mockReceiver.get(CreateMongoDbChangeEventContextFn.successfulCreationTag)) + when(mockReceiver.get(CreateMongoDbChangeEventContextFn.SUCCESSFUL_CREATION_TAG)) .thenReturn(mockSuccessReceiver); - when(mockReceiver.get(CreateMongoDbChangeEventContextFn.failedCreationTag)) + when(mockReceiver.get(CreateMongoDbChangeEventContextFn.FAILED_CREATION_TAG)) .thenReturn(mockFailureReceiver); String validPayload = @@ -94,7 +94,7 @@ public void testProcessElementSuccess() { ArgumentCaptor successCaptor = ArgumentCaptor.forClass(MongoDbChangeEventContext.class); - verify(mockReceiver).get(CreateMongoDbChangeEventContextFn.successfulCreationTag); + verify(mockReceiver).get(CreateMongoDbChangeEventContextFn.SUCCESSFUL_CREATION_TAG); verify(mockSuccessReceiver, times(1)).output(successCaptor.capture()); MongoDbChangeEventContext actualContext = successCaptor.getValue(); @@ -111,7 +111,7 @@ public void testProcessElementFailureInvalidJson() throws Exception { ArgumentCaptor> failureCaptor = ArgumentCaptor.forClass(FailsafeElement.class); - verify(mockReceiver).get(CreateMongoDbChangeEventContextFn.failedCreationTag); + verify(mockReceiver).get(CreateMongoDbChangeEventContextFn.FAILED_CREATION_TAG); verify(mockFailureReceiver, times(1)).output(failureCaptor.capture()); assertEquals(failureElement, failureCaptor.getValue()); diff --git a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbBulkTransformsTest.java b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbBulkTransformsTest.java new file mode 100644 index 0000000000..cb990c3a1f --- /dev/null +++ b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbBulkTransformsTest.java @@ -0,0 +1,556 @@ +/* + * 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.assertNotNull; +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.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.cloud.teleport.v2.templates.datastream.MongoDbChangeEventContext; +import com.google.cloud.teleport.v2.values.FailsafeElement; +import com.mongodb.MongoBulkWriteException; +import com.mongodb.ServerAddress; +import com.mongodb.bulk.BulkWriteError; +import com.mongodb.bulk.BulkWriteResult; +import com.mongodb.client.MongoClient; +import com.mongodb.client.MongoCollection; +import com.mongodb.client.MongoDatabase; +import com.mongodb.client.model.BulkWriteOptions; +import com.mongodb.client.model.WriteModel; +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.SerializableFunction; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.bson.BsonDocument; +import org.bson.Document; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link MongoDbBulkTransforms}. */ +@RunWith(JUnit4.class) +public class MongoDbBulkTransformsTest { + + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + private static MongoClient mockClient; + private static MongoDatabase mockDatabase; + private static MongoCollection mockCollection; + + @Before + public void setUp() { + mockClient = mock(MongoClient.class); + mockDatabase = mock(MongoDatabase.class); + mockCollection = mock(MongoCollection.class); + + when(mockClient.getDatabase(anyString())).thenReturn(mockDatabase); + when(mockDatabase.getCollection(anyString())).thenReturn(mockCollection); + } + + private static class MockClientFactory + implements SerializableFunction, Serializable { + @Override + public MongoClient apply(String uri) { + return mockClient; + } + } + + private MongoDbChangeEventContext createEventContext( + String docId, String changeType, boolean isDelete) throws Exception { + String payload = + String.format( + "{" + + "\"_metadata_source\": {\"collection\": \"users\"}," + + "\"_id\": \"\\\"%s\\\"\"," + + "\"_metadata_timestamp_seconds\": 1000," + + "\"_metadata_timestamp_nanos\": 0," + + "\"_metadata_change_type\": \"%s\"," + + "\"data\": \"{\\\"name\\\": \\\"user_%s\\\"}\"" + + "}", + docId, changeType, docId); + return new MongoDbChangeEventContext(OBJECT_MAPPER.readTree(payload), "shadow_"); + } + + @Test + public void testSuccessfulBulkWrites() throws Exception { + MongoDbChangeEventContext insertEvent = createEventContext("doc1", "INSERT", false); + MongoDbChangeEventContext deleteEvent = createEventContext("doc2", "DELETE", true); + + when(mockCollection.bulkWrite(anyList(), any(BulkWriteOptions.class))) + .thenReturn(mock(BulkWriteResult.class)); + + PCollectionTuple result = + pipeline + .apply( + Create.of(insertEvent, deleteEvent).withCoder(MongoDbChangeEventContextCoder.of())) + .apply( + MongoDbBulkTransforms.bulkWriteWithDlq() + .withConnectionString("mongodb://localhost:27017") + .withDatabase("test_db") + .withBatchSize(2) + .withClientFactory(new MockClientFactory())); + + PAssert.that(result.get(MongoDbBulkTransforms.SUCCESSFUL_WRITE_TAG)) + .containsInAnyOrder(insertEvent, deleteEvent); + PAssert.that(result.get(MongoDbBulkTransforms.FAILED_WRITE_TAG)).empty(); + PAssert.that(result.get(MongoDbBulkTransforms.SEVERE_FAILED_WRITE_TAG)).empty(); + + pipeline.run(); + } + + @Test + public void testErrorCodeClassification_code2_routesToSevereDlq() throws Exception { + MongoDbChangeEventContext badDoc = createEventContext("bad_doc", "INSERT", false); + + BulkWriteError error = + new BulkWriteError( + MongoDbBulkTransforms.ERR_BAD_VALUE, + "BadValue: value exceeds limit", + new BsonDocument(), + 0); + MongoBulkWriteException exception = + new MongoBulkWriteException( + mock(BulkWriteResult.class), + Collections.singletonList(error), + null, + new ServerAddress("localhost", 27017), + Collections.emptySet()); + + doThrow(exception).when(mockCollection).bulkWrite(anyList(), any(BulkWriteOptions.class)); + + PCollectionTuple result = + pipeline + .apply(Create.of(badDoc).withCoder(MongoDbChangeEventContextCoder.of())) + .apply( + MongoDbBulkTransforms.bulkWriteWithDlq() + .withConnectionString("mongodb://localhost:27017") + .withDatabase("test_db") + .withBatchSize(1) + .withClientFactory(new MockClientFactory())); + + PAssert.that(result.get(MongoDbBulkTransforms.SUCCESSFUL_WRITE_TAG)).empty(); + PAssert.that(result.get(MongoDbBulkTransforms.FAILED_WRITE_TAG)).empty(); + + PCollection> severeOut = + result.get(MongoDbBulkTransforms.SEVERE_FAILED_WRITE_TAG); + PAssert.that(severeOut) + .satisfies( + elements -> { + int count = 0; + for (FailsafeElement elem : + elements) { + count++; + assertNotNull(elem); + assertTrue(elem.getErrorMessage().contains("Code 2")); + } + assertEquals(1, count); + return null; + }); + + pipeline.run(); + } + + @Test + public void testErrorCodeClassification_code121_routesToSevereDlq() throws Exception { + MongoDbChangeEventContext invalidDoc = createEventContext("invalid_doc", "UPDATE", false); + + BulkWriteError error = + new BulkWriteError( + MongoDbBulkTransforms.ERR_DOCUMENT_VALIDATION_FAILURE, + "Document failed validation", + new BsonDocument(), + 0); + MongoBulkWriteException exception = + new MongoBulkWriteException( + mock(BulkWriteResult.class), + Collections.singletonList(error), + null, + new ServerAddress("localhost", 27017), + Collections.emptySet()); + + doThrow(exception).when(mockCollection).bulkWrite(anyList(), any(BulkWriteOptions.class)); + + PCollectionTuple result = + pipeline + .apply(Create.of(invalidDoc).withCoder(MongoDbChangeEventContextCoder.of())) + .apply( + MongoDbBulkTransforms.bulkWriteWithDlq() + .withConnectionString("mongodb://localhost:27017") + .withDatabase("test_db") + .withBatchSize(1) + .withClientFactory(new MockClientFactory())); + + PAssert.that(result.get(MongoDbBulkTransforms.SUCCESSFUL_WRITE_TAG)).empty(); + PAssert.that(result.get(MongoDbBulkTransforms.FAILED_WRITE_TAG)).empty(); + + PCollection> severeOut = + result.get(MongoDbBulkTransforms.SEVERE_FAILED_WRITE_TAG); + PAssert.that(severeOut) + .satisfies( + elements -> { + int count = 0; + for (FailsafeElement elem : + elements) { + count++; + assertTrue(elem.getErrorMessage().contains("Code 121")); + } + assertEquals(1, count); + return null; + }); + + pipeline.run(); + } + + @Test + public void testErrorCodeClassification_transientError_retriesAndRoutesToRetryableDlq() + throws Exception { + MongoDbChangeEventContext transientDoc = createEventContext("transient_doc", "UPDATE", false); + + BulkWriteError error = + new BulkWriteError( + MongoDbBulkTransforms.ERR_WRITE_CONFLICT, + "WriteConflict: retryable conflict", + new BsonDocument(), + 0); + MongoBulkWriteException exception = + new MongoBulkWriteException( + mock(BulkWriteResult.class), + Collections.singletonList(error), + null, + new ServerAddress("localhost", 27017), + Collections.emptySet()); + + doThrow(exception).when(mockCollection).bulkWrite(anyList(), any(BulkWriteOptions.class)); + + PCollectionTuple result = + pipeline + .apply(Create.of(transientDoc).withCoder(MongoDbChangeEventContextCoder.of())) + .apply( + MongoDbBulkTransforms.bulkWriteWithDlq() + .withConnectionString("mongodb://localhost:27017") + .withDatabase("test_db") + .withBatchSize(1) + .withMaxWriteRetries(1) + .withClientFactory(new MockClientFactory())); + + PAssert.that(result.get(MongoDbBulkTransforms.SUCCESSFUL_WRITE_TAG)).empty(); + PAssert.that(result.get(MongoDbBulkTransforms.SEVERE_FAILED_WRITE_TAG)).empty(); + + PCollection> + retryableOut = result.get(MongoDbBulkTransforms.FAILED_WRITE_TAG); + PAssert.that(retryableOut) + .satisfies( + elements -> { + int count = 0; + for (FailsafeElement elem : + elements) { + count++; + assertTrue( + elem.getErrorMessage().contains("Transient write error retries exhausted")); + } + assertEquals(1, count); + return null; + }); + + pipeline.run(); + } + + private MongoDbChangeEventContext createEventContextWithTimestamp( + String docId, long seconds, int nanos, String changeType) throws Exception { + String payload = + String.format( + "{" + + "\"_metadata_source\": {\"collection\": \"users\"}," + + "\"_id\": \"\\\"%s\\\"\"," + + "\"_metadata_timestamp_seconds\": %d," + + "\"_metadata_timestamp_nanos\": %d," + + "\"_metadata_change_type\": \"%s\"," + + "\"_metadata_read_method\": \"cdc\"," + + "\"data\": \"{\\\"name\\\": \\\"user_%s\\\"}\"" + + "}", + docId, seconds, nanos, changeType, docId); + return new MongoDbChangeEventContext(OBJECT_MAPPER.readTree(payload), "shadow_"); + } + + @Test + public void testCoalescedBatch_successfulWrite_emitsActiveAndSupersededEvents() throws Exception { + MongoDbChangeEventContext v1 = createEventContextWithTimestamp("doc1", 1000L, 0, "INSERT"); + MongoDbChangeEventContext v2 = createEventContextWithTimestamp("doc1", 1000L, 100, "UPDATE"); + MongoDbChangeEventContext v3 = createEventContextWithTimestamp("doc1", 1001L, 0, "UPDATE"); + + when(mockCollection.bulkWrite(anyList(), any(BulkWriteOptions.class))) + .thenReturn(mock(BulkWriteResult.class)); + + PCollectionTuple result = + pipeline + .apply(Create.of(v1, v2, v3).withCoder(MongoDbChangeEventContextCoder.of())) + .apply( + MongoDbBulkTransforms.bulkWriteWithDlq() + .withConnectionString("mongodb://localhost:27017") + .withDatabase("test_db") + .withBatchSize(10) + .withClientFactory(new MockClientFactory())); + + PAssert.that(result.get(MongoDbBulkTransforms.SUCCESSFUL_WRITE_TAG)) + .containsInAnyOrder(v1, v2, v3); + PAssert.that(result.get(MongoDbBulkTransforms.FAILED_WRITE_TAG)).empty(); + PAssert.that(result.get(MongoDbBulkTransforms.SEVERE_FAILED_WRITE_TAG)).empty(); + + pipeline.run(); + } + + @Test + public void testCoalescedBatch_permanentFailure_doesNotEmitSupersededEvents() throws Exception { + MongoDbChangeEventContext v1 = createEventContextWithTimestamp("bad_doc", 1000L, 0, "INSERT"); + MongoDbChangeEventContext v2 = createEventContextWithTimestamp("bad_doc", 1000L, 100, "UPDATE"); + + when(mockCollection.bulkWrite(anyList(), any(BulkWriteOptions.class))) + .thenAnswer( + invocation -> { + List> ops = invocation.getArgument(0); + List errors = new ArrayList<>(); + for (int i = 0; i < ops.size(); i++) { + errors.add( + new BulkWriteError( + MongoDbBulkTransforms.ERR_BAD_VALUE, + "BadValue: value exceeds limit", + new BsonDocument(), + i)); + } + throw new MongoBulkWriteException( + mock(BulkWriteResult.class), + errors, + null, + new ServerAddress("localhost", 27017), + Collections.emptySet()); + }); + + PCollectionTuple result = + pipeline + .apply(Create.of(v1, v2).withCoder(MongoDbChangeEventContextCoder.of())) + .apply( + MongoDbBulkTransforms.bulkWriteWithDlq() + .withConnectionString("mongodb://localhost:27017") + .withDatabase("test_db") + .withBatchSize(10) + .withClientFactory(new MockClientFactory())); + + PAssert.that(result.get(MongoDbBulkTransforms.SUCCESSFUL_WRITE_TAG)).empty(); + PAssert.that(result.get(MongoDbBulkTransforms.FAILED_WRITE_TAG)).empty(); + + PCollection> severeOut = + result.get(MongoDbBulkTransforms.SEVERE_FAILED_WRITE_TAG); + PAssert.that(severeOut) + .satisfies( + elements -> { + int count = 0; + for (FailsafeElement elem : + elements) { + count++; + assertNotNull(elem); + assertTrue(elem.getErrorMessage().contains("Code 2")); + } + assertTrue(count >= 1); + return null; + }); + + pipeline.run(); + } + + @Test + public void testPartialBatchFailure_emitsSupersededEventsOnlyForSuccessfulDocs() + throws Exception { + MongoDbChangeEventContext docAv1 = createEventContextWithTimestamp("docA", 1000L, 0, "INSERT"); + MongoDbChangeEventContext docAv2 = + createEventContextWithTimestamp("docA", 1000L, 100, "UPDATE"); + MongoDbChangeEventContext docBv1 = createEventContextWithTimestamp("docB", 1000L, 0, "INSERT"); + MongoDbChangeEventContext docBv2 = + createEventContextWithTimestamp("docB", 1000L, 100, "UPDATE"); + + when(mockCollection.bulkWrite(anyList(), any(BulkWriteOptions.class))) + .thenAnswer( + invocation -> { + List> ops = invocation.getArgument(0); + List errors = new ArrayList<>(); + for (int i = 0; i < ops.size(); i++) { + WriteModel op = ops.get(i); + if (op.toString().contains("docB")) { + errors.add( + new BulkWriteError( + MongoDbBulkTransforms.ERR_BAD_VALUE, + "BadValue for docB", + new BsonDocument(), + i)); + } + } + if (!errors.isEmpty()) { + throw new MongoBulkWriteException( + mock(BulkWriteResult.class), + errors, + null, + new ServerAddress("localhost", 27017), + Collections.emptySet()); + } + return mock(BulkWriteResult.class); + }); + + PCollectionTuple result = + pipeline + .apply( + Create.of(docAv1, docAv2, docBv1, docBv2) + .withCoder(MongoDbChangeEventContextCoder.of())) + .apply( + MongoDbBulkTransforms.bulkWriteWithDlq() + .withConnectionString("mongodb://localhost:27017") + .withDatabase("test_db") + .withBatchSize(10) + .withClientFactory(new MockClientFactory())); + + // docA was successful -> docA events should be emitted as successful + PAssert.that(result.get(MongoDbBulkTransforms.SUCCESSFUL_WRITE_TAG)) + .satisfies( + elements -> { + int count = 0; + for (MongoDbChangeEventContext elem : elements) { + count++; + assertEquals("docA", elem.getDocumentId()); + } + assertTrue(count >= 1); + return null; + }); + PAssert.that(result.get(MongoDbBulkTransforms.FAILED_WRITE_TAG)).empty(); + + // docB failed -> only docB events should be emitted to severe DLQ, no docA in severe DLQ + PCollection> severeOut = + result.get(MongoDbBulkTransforms.SEVERE_FAILED_WRITE_TAG); + PAssert.that(severeOut) + .satisfies( + elements -> { + int count = 0; + for (FailsafeElement elem : + elements) { + count++; + assertNotNull(elem); + assertEquals("docB", elem.getOriginalPayload().getDocumentId()); + assertTrue(elem.getErrorMessage().contains("Code 2")); + } + assertTrue(count >= 1); + return null; + }); + + pipeline.run(); + } + + @Test + public void testGeneralException_inCollectionOrSetup_routesToDlqWithoutPipelineCrash() + throws Exception { + MongoDbChangeEventContext event = createEventContext("doc1", "INSERT", false); + + when(mockDatabase.getCollection(anyString())) + .thenThrow(new com.mongodb.MongoException(13, "Unauthorized collection access")); + + PCollectionTuple result = + pipeline + .apply(Create.of(event).withCoder(MongoDbChangeEventContextCoder.of())) + .apply( + MongoDbBulkTransforms.bulkWriteWithDlq() + .withConnectionString("mongodb://localhost:27017") + .withDatabase("test_db") + .withBatchSize(1) + .withClientFactory(new MockClientFactory())); + + PAssert.that(result.get(MongoDbBulkTransforms.SUCCESSFUL_WRITE_TAG)).empty(); + PAssert.that(result.get(MongoDbBulkTransforms.FAILED_WRITE_TAG)).empty(); + + PCollection> severeOut = + result.get(MongoDbBulkTransforms.SEVERE_FAILED_WRITE_TAG); + PAssert.that(severeOut) + .satisfies( + elements -> { + int count = 0; + for (FailsafeElement elem : + elements) { + count++; + assertNotNull(elem); + assertTrue(elem.getErrorMessage().contains("Permanent failure")); + } + assertEquals(1, count); + return null; + }); + + pipeline.run(); + } + + @Test + public void testPartialBatchFailure_transientErrorRetriedAndSucceeds() throws Exception { + MongoDbChangeEventContext docA = createEventContextWithTimestamp("docA", 1000L, 0, "INSERT"); + MongoDbChangeEventContext docB = createEventContextWithTimestamp("docB", 1000L, 0, "INSERT"); + + // First call: docB fails with transient WriteConflict (Code 112), docA succeeds + // Second call (retry): docB succeeds + BulkWriteError transientError = + new BulkWriteError( + MongoDbBulkTransforms.ERR_WRITE_CONFLICT, + "WriteConflict", + new org.bson.BsonDocument(), + 1); + MongoBulkWriteException partialException = + new MongoBulkWriteException( + mock(BulkWriteResult.class), + Collections.singletonList(transientError), + null, + new com.mongodb.ServerAddress("localhost", 27017), + Collections.emptySet()); + + when(mockCollection.bulkWrite(anyList(), any(BulkWriteOptions.class))) + .thenThrow(partialException) + .thenReturn(mock(BulkWriteResult.class)); + + PCollectionTuple result = + pipeline + .apply(Create.of(docA, docB).withCoder(MongoDbChangeEventContextCoder.of())) + .apply( + MongoDbBulkTransforms.bulkWriteWithDlq() + .withConnectionString("mongodb://localhost:27017") + .withDatabase("test_db") + .withBatchSize(2) + .withClientFactory(new MockClientFactory())); + + // Both should ultimately succeed without routing to DLQ + PAssert.that(result.get(MongoDbBulkTransforms.SUCCESSFUL_WRITE_TAG)) + .containsInAnyOrder(docA, docB); + PAssert.that(result.get(MongoDbBulkTransforms.FAILED_WRITE_TAG)).empty(); + PAssert.that(result.get(MongoDbBulkTransforms.SEVERE_FAILED_WRITE_TAG)).empty(); + + pipeline.run(); + } +} diff --git a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbChangeEventContextCoderTest.java b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbChangeEventContextCoderTest.java new file mode 100644 index 0000000000..0e28f5b389 --- /dev/null +++ b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/MongoDbChangeEventContextCoderTest.java @@ -0,0 +1,343 @@ +/* + * 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.assertNull; +import static org.junit.Assert.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.cloud.teleport.v2.templates.datastream.MongoDbChangeEventContext; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import org.apache.beam.sdk.testing.CoderProperties; +import org.apache.beam.sdk.util.CoderUtils; +import org.bson.Document; +import org.bson.types.ObjectId; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link MongoDbChangeEventContextCoder}. */ +@RunWith(JUnit4.class) +public class MongoDbChangeEventContextCoderTest { + + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + private static final String SHADOW_PREFIX = "shadow_"; + + private final MongoDbChangeEventContextCoder coder = MongoDbChangeEventContextCoder.of(); + + private JsonNode insertEventJson; + private JsonNode updateEventJson; + private JsonNode deleteEventJson; + private JsonNode readBackfillEventJson; + + @Before + public void setUp() throws Exception { + insertEventJson = + OBJECT_MAPPER.readTree( + """ + { + "_metadata_source": { + "collection": "users" + }, + "_id": "{\\\"$oid\\\": \\\"645c9a7e7b8b1a0e9c0f8b3a\\\"}", + "data": { + "name": "Alice", + "age": 30 + }, + "_metadata_timestamp_seconds": 1683782270, + "_metadata_timestamp_nanos": 123456789, + "op": "i" + }\ + """); + + updateEventJson = + OBJECT_MAPPER.readTree( + """ + { + "_metadata_source": { + "collection": "users" + }, + "_id": "{\\\"$oid\\\": \\\"645c9a7e7b8b1a0e9c0f8b3a\\\"}", + "data": { + "name": "Alice Smith", + "age": 31 + }, + "_metadata_timestamp_seconds": 1683782275, + "_metadata_timestamp_nanos": 200, + "op": "u", + "_metadata_change_type": "UPDATE" + }\ + """); + + deleteEventJson = + OBJECT_MAPPER.readTree( + """ + { + "_metadata_source": { + "collection": "users" + }, + "_id": "{\\\"$oid\\\": \\\"645c9a7e7b8b1a0e9c0f8b3a\\\"}", + "_metadata_timestamp_seconds": 1683782280, + "_metadata_timestamp_nanos": 300, + "op": "d", + "_metadata_change_type": "DELETE" + }\ + """); + + readBackfillEventJson = + OBJECT_MAPPER.readTree( + """ + { + "_metadata_source": { + "collection": "users" + }, + "_id": "{\\\"$oid\\\": \\\"645c9a7e7b8b1a0e9c0f8b3a\\\"}", + "data": { + "name": "Alice", + "age": 30 + }, + "_metadata_timestamp_seconds": 1683782260, + "_metadata_timestamp_nanos": 999999000, + "op": "r", + "_metadata_change_type": "READ", + "_metadata_read_method": "backfill" + }\ + """); + } + + @Test + public void testEncodeDecodeRoundTrip_insertEvent() throws Exception { + MongoDbChangeEventContext context = + new MongoDbChangeEventContext(insertEventJson, SHADOW_PREFIX); + MongoDbChangeEventContext decoded = CoderUtils.clone(coder, context); + + assertEquals(context, decoded); + assertEquals("users", decoded.getDataCollection()); + assertEquals("shadow_users", decoded.getShadowCollection()); + assertEquals(SHADOW_PREFIX, decoded.getShadowCollectionPrefix()); + assertTrue(decoded.getDocumentId() instanceof ObjectId); + assertEquals("645c9a7e7b8b1a0e9c0f8b3a", decoded.getDocumentId().toString()); + assertEquals(1683782270L, decoded.getTimestampSeconds()); + assertEquals(123456789L, decoded.getTimestampSubSeconds()); + assertFalse(decoded.isDeleteEvent()); + assertNotNull(decoded.getDataAsJsonString()); + assertFalse(decoded.getIsDlqReconsumed()); + assertEquals(0, decoded.getRetryCount()); + assertEquals(context.getChangeEvent(), decoded.getChangeEvent()); + assertEquals(context.getOriginalChangeEvent(), decoded.getOriginalChangeEvent()); + CoderProperties.coderDecodeEncodeEqual(coder, context); + } + + @Test + public void testEncodeDecodeRoundTrip_updateEvent() throws Exception { + MongoDbChangeEventContext context = + new MongoDbChangeEventContext(updateEventJson, SHADOW_PREFIX); + MongoDbChangeEventContext decoded = CoderUtils.clone(coder, context); + + assertEquals(context, decoded); + assertTrue(decoded.isUpdateEvent()); + assertFalse(decoded.isDeleteEvent()); + assertEquals(1683782275L, decoded.getTimestampSeconds()); + assertEquals(200L, decoded.getTimestampSubSeconds()); + CoderProperties.coderDecodeEncodeEqual(coder, context); + } + + @Test + public void testEncodeDecodeRoundTrip_deleteEvent() throws Exception { + MongoDbChangeEventContext context = + new MongoDbChangeEventContext(deleteEventJson, SHADOW_PREFIX); + MongoDbChangeEventContext decoded = CoderUtils.clone(coder, context); + + assertEquals(context, decoded); + assertTrue(decoded.isDeleteEvent()); + assertNull(decoded.getDataAsJsonString()); + assertEquals(1683782280L, decoded.getTimestampSeconds()); + CoderProperties.coderDecodeEncodeEqual(coder, context); + } + + @Test + public void testEncodeDecodeRoundTrip_readBackfillEvent() throws Exception { + MongoDbChangeEventContext context = + new MongoDbChangeEventContext(readBackfillEventJson, SHADOW_PREFIX); + MongoDbChangeEventContext decoded = CoderUtils.clone(coder, context); + + assertEquals(context, decoded); + assertTrue(decoded.isBackfillEvent()); + assertFalse(decoded.isCdcEvent()); + assertEquals(999999000L, decoded.getTimestampSubSeconds()); + CoderProperties.coderDecodeEncodeEqual(coder, context); + } + + @Test + public void testEncodeDecodeRoundTrip_udfPayloadPreservation() throws Exception { + JsonNode originalEvent = insertEventJson.deepCopy(); + JsonNode modifiedEvent = insertEventJson.deepCopy(); + ((ObjectNode) modifiedEvent.get("data")).put("transformedByUdf", true); + + MongoDbChangeEventContext context = + new MongoDbChangeEventContext(modifiedEvent, originalEvent, SHADOW_PREFIX); + MongoDbChangeEventContext decoded = CoderUtils.clone(coder, context); + + assertEquals(context, decoded); + // Verify modified event has the UDF-added field + assertTrue(decoded.getChangeEvent().get("data").has("transformedByUdf")); + assertTrue(decoded.getChangeEvent().get("data").get("transformedByUdf").asBoolean()); + + // Verify original event DOES NOT have the UDF-added field (original is preserved) + assertFalse(decoded.getOriginalChangeEvent().get("data").has("transformedByUdf")); + CoderProperties.coderDecodeEncodeEqual(coder, context); + } + + @Test + public void testEncodeDecodeRoundTrip_dlqReconsumedWithRetries() throws Exception { + String dlqPayload = + """ + { + "_metadata_source": { + "collection": "orders" + }, + "_id": "\\\"order_123\\\"", + "data": { + "amount": 99.99 + }, + "_metadata_timestamp_seconds": 1683782270, + "_metadata_timestamp_nanos": 100, + "isDlqReconsumed": "true", + "_metadata_retry_count": 3 + }\ + """; + MongoDbChangeEventContext context = + new MongoDbChangeEventContext(OBJECT_MAPPER.readTree(dlqPayload), SHADOW_PREFIX); + MongoDbChangeEventContext decoded = CoderUtils.clone(coder, context); + + assertEquals(context, decoded); + assertTrue(decoded.getIsDlqReconsumed()); + assertEquals(3, decoded.getRetryCount()); + CoderProperties.coderDecodeEncodeEqual(coder, context); + } + + @Test + public void testEncodeDecodeRoundTrip_variousDocumentIdTypes() throws Exception { + // String doc ID + String stringIdPayload = + """ + { + "_metadata_source": {"collection": "c1"}, + "_id": "\\\"str_id_val\\\"", + "_metadata_timestamp_seconds": 1000, + "_metadata_timestamp_nanos": 1, + "data": {} + }\ + """; + MongoDbChangeEventContext strContext = + new MongoDbChangeEventContext(OBJECT_MAPPER.readTree(stringIdPayload), SHADOW_PREFIX); + assertEquals("str_id_val", CoderUtils.clone(coder, strContext).getDocumentId()); + + // Long doc ID + String longIdPayload = + """ + { + "_metadata_source": {"collection": "c2"}, + "_id": 9223372036854775806, + "_metadata_timestamp_seconds": 1000, + "_metadata_timestamp_nanos": 1, + "data": {} + }\ + """; + MongoDbChangeEventContext longContext = + new MongoDbChangeEventContext(OBJECT_MAPPER.readTree(longIdPayload), SHADOW_PREFIX); + assertEquals(9223372036854775806L, CoderUtils.clone(coder, longContext).getDocumentId()); + + // Integer doc ID + String intIdPayload = + """ + { + "_metadata_source": {"collection": "c3"}, + "_id": 42, + "_metadata_timestamp_seconds": 1000, + "_metadata_timestamp_nanos": 1, + "data": {} + }\ + """; + MongoDbChangeEventContext intContext = + new MongoDbChangeEventContext(OBJECT_MAPPER.readTree(intIdPayload), SHADOW_PREFIX); + assertEquals(42, CoderUtils.clone(coder, intContext).getDocumentId()); + + // Double doc ID + String doubleIdPayload = + """ + { + "_metadata_source": {"collection": "c4"}, + "_id": "123.456", + "_metadata_timestamp_seconds": 1000, + "_metadata_timestamp_nanos": 1, + "data": {} + }\ + """; + MongoDbChangeEventContext doubleContext = + new MongoDbChangeEventContext(OBJECT_MAPPER.readTree(doubleIdPayload), SHADOW_PREFIX); + assertEquals(123.456, (Double) CoderUtils.clone(coder, doubleContext).getDocumentId(), 0.001); + + // Composite Document (Map) doc ID + String compositeIdPayload = + """ + { + "_metadata_source": {"collection": "c5"}, + "_id": "{\\\"tenant\\\": \\\"acme\\\", \\\"uid\\\": 100}", + "_metadata_timestamp_seconds": 1000, + "_metadata_timestamp_nanos": 1, + "data": {} + }\ + """; + MongoDbChangeEventContext compContext = + new MongoDbChangeEventContext(OBJECT_MAPPER.readTree(compositeIdPayload), SHADOW_PREFIX); + Object docId = CoderUtils.clone(coder, compContext).getDocumentId(); + assertTrue(docId instanceof Document); + assertEquals("acme", ((Document) docId).get("tenant")); + assertEquals(100, ((Document) docId).get("uid")); + } + + @Test + public void testNullValueEncoding() throws Exception { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + coder.encode(null, out); + byte[] bytes = out.toByteArray(); + + assertEquals(1, bytes.length); + + ByteArrayInputStream in = new ByteArrayInputStream(bytes); + MongoDbChangeEventContext decoded = coder.decode(in); + assertNull(decoded); + } + + @Test + public void testVerifyDeterministic() throws Exception { + coder.verifyDeterministic(); + + MongoDbChangeEventContext context1 = + new MongoDbChangeEventContext(insertEventJson, SHADOW_PREFIX); + MongoDbChangeEventContext context2 = + new MongoDbChangeEventContext(insertEventJson, SHADOW_PREFIX); + CoderProperties.coderDeterministic(coder, context1, context2); + } +} diff --git a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/ProcessChangeEventFnTest.java b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/ProcessChangeEventFnTest.java index bdb911644a..50134cfd9f 100644 --- a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/ProcessChangeEventFnTest.java +++ b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/ProcessChangeEventFnTest.java @@ -140,9 +140,10 @@ public void setUp() throws Exception { when(mockShadowCollection.find(mockSession, LOOKUP_BY_DOC_ID)).thenReturn(mockFindIterable); // Mock the MultiOutputReceiver's get() method - when(mockReceiver.get(ProcessChangeEventFn.successfulWriteTag)).thenReturn(mockSuccessReceiver); - when(mockReceiver.get(ProcessChangeEventFn.failedWriteTag)).thenReturn(mockFailureReceiver); - when(mockReceiver.get(ProcessChangeEventFn.severeFailedWriteTag)) + when(mockReceiver.get(ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG)) + .thenReturn(mockSuccessReceiver); + when(mockReceiver.get(ProcessChangeEventFn.FAILED_WRITE_TAG)).thenReturn(mockFailureReceiver); + when(mockReceiver.get(ProcessChangeEventFn.SEVERE_FAILED_WRITE_TAG)) .thenReturn(mockSevereFailureReceiver); } @@ -164,7 +165,7 @@ mockSession, LOOKUP_BY_DOC_ID, mockShadowDocElement, new ReplaceOptions().upsert ArgumentCaptor successCaptor = ArgumentCaptor.forClass(MongoDbChangeEventContext.class); - verify(mockReceiver).get(ProcessChangeEventFn.successfulWriteTag); + verify(mockReceiver).get(ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG); verify(mockSuccessReceiver, times(1)).output(successCaptor.capture()); verify(mockSession, never()).abortTransaction(); @@ -201,7 +202,7 @@ mockSession, LOOKUP_BY_DOC_ID, mockShadowDocElement, new ReplaceOptions().upsert verify(mockSession).commitTransaction(); ArgumentCaptor successCaptor = ArgumentCaptor.forClass(MongoDbChangeEventContext.class); - verify(mockReceiver).get(ProcessChangeEventFn.successfulWriteTag); + verify(mockReceiver).get(ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG); verify(mockSuccessReceiver, times(1)).output(successCaptor.capture()); verify(mockSession, never()).abortTransaction(); } @@ -218,7 +219,7 @@ public void testProcessElementInsertOrUpdateOlderExisting() { verify(mockShadowCollection, never()).replaceOne(any(), any(), any(), any()); ArgumentCaptor successCaptor = ArgumentCaptor.forClass(MongoDbChangeEventContext.class); - verify(mockReceiver).get(ProcessChangeEventFn.successfulWriteTag); + verify(mockReceiver).get(ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG); verify(mockSuccessReceiver, times(1)).output(successCaptor.capture()); verify(mockSession, never()).abortTransaction(); @@ -260,10 +261,10 @@ mockSession, LOOKUP_BY_DOC_ID, mockShadowDocElement, new ReplaceOptions().upsert verify(mockShadowCollection).find(mockSession, LOOKUP_BY_DOC_ID); verify(mockDataCollection).deleteOne(mockSession, LOOKUP_BY_DOC_ID); verify(mockSession).commitTransaction(); - verify(mockReceiver).get(ProcessChangeEventFn.successfulWriteTag); + verify(mockReceiver).get(ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG); ArgumentCaptor successCaptor = ArgumentCaptor.forClass(MongoDbChangeEventContext.class); - verify(mockReceiver).get(ProcessChangeEventFn.successfulWriteTag); + verify(mockReceiver).get(ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG); verify(mockSuccessReceiver, times(1)).output(successCaptor.capture()); verify(mockSession, never()).abortTransaction(); } @@ -286,7 +287,7 @@ mockSession, LOOKUP_BY_DOC_ID, mockShadowDocElement, new ReplaceOptions().upsert verify(mockSession).commitTransaction(); ArgumentCaptor successCaptor = ArgumentCaptor.forClass(MongoDbChangeEventContext.class); - verify(mockReceiver).get(ProcessChangeEventFn.successfulWriteTag); + verify(mockReceiver).get(ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG); verify(mockSuccessReceiver, times(1)).output(successCaptor.capture()); verify(mockSession, never()).abortTransaction(); } @@ -304,7 +305,7 @@ public void testProcessElementDeleteOlderExisting() { verify(mockShadowCollection, never()).replaceOne(any(), any(), any(), any()); ArgumentCaptor successCaptor = ArgumentCaptor.forClass(MongoDbChangeEventContext.class); - verify(mockReceiver).get(ProcessChangeEventFn.successfulWriteTag); + verify(mockReceiver).get(ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG); verify(mockSuccessReceiver, times(1)).output(successCaptor.capture()); verify(mockSession, never()).abortTransaction(); } @@ -322,7 +323,7 @@ public void testProcessElementTransientError_mixedErrors() { verify(mockShadowCollection, times(2)).find(mockSession, LOOKUP_BY_DOC_ID); ArgumentCaptor failureCaptor = ArgumentCaptor.forClass(MongoDbChangeEventContext.class); - verify(mockReceiver).get(ProcessChangeEventFn.severeFailedWriteTag); + verify(mockReceiver).get(ProcessChangeEventFn.SEVERE_FAILED_WRITE_TAG); verify(mockSevereFailureReceiver, times(1)).output(failureCaptor.capture()); verify(mockSession, never()).commitTransaction(); } @@ -338,7 +339,7 @@ public void testProcessElementTransientError_retryTillMaximum() { verify(mockShadowCollection, times(4)).find(mockSession, LOOKUP_BY_DOC_ID); ArgumentCaptor failureCaptor = ArgumentCaptor.forClass(MongoDbChangeEventContext.class); - verify(mockReceiver).get(ProcessChangeEventFn.failedWriteTag); + verify(mockReceiver).get(ProcessChangeEventFn.FAILED_WRITE_TAG); verify(mockFailureReceiver, times(1)).output(failureCaptor.capture()); verify(mockSession, never()).commitTransaction(); @@ -376,7 +377,7 @@ public void testProcessElementPermanentError_Code2() { processFn.processElement(mockContext, mockReceiver); verify(mockShadowCollection, times(1)).find(mockSession, LOOKUP_BY_DOC_ID); - verify(mockReceiver).get(ProcessChangeEventFn.severeFailedWriteTag); + verify(mockReceiver).get(ProcessChangeEventFn.SEVERE_FAILED_WRITE_TAG); verify(mockSevereFailureReceiver, times(1)).output(any()); verify(mockSession, never()).commitTransaction(); @@ -412,7 +413,7 @@ mockSession, LOOKUP_BY_DOC_ID, mockShadowDocElement, new ReplaceOptions().upsert verify(mockShadowCollection).find(mockSession, LOOKUP_BY_DOC_ID); verify(mockSession).commitTransaction(); - verify(mockReceiver).get(ProcessChangeEventFn.successfulWriteTag); + verify(mockReceiver).get(ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG); verify(mockSuccessReceiver, times(1)).output(any()); MetricsContainerImpl container = @@ -450,7 +451,7 @@ public void testProcessElementTransientWriteError_retry() { processFn.processElement(mockContext, mockReceiver); verify(mockShadowCollection, times(2)).find(mockSession, LOOKUP_BY_DOC_ID); - verify(mockReceiver).get(ProcessChangeEventFn.successfulWriteTag); + verify(mockReceiver).get(ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG); } @Test @@ -472,7 +473,7 @@ public void testProcessElementTransientCommandError_retry() { processFn.processElement(mockContext, mockReceiver); verify(mockShadowCollection, times(2)).find(mockSession, LOOKUP_BY_DOC_ID); - verify(mockReceiver).get(ProcessChangeEventFn.successfulWriteTag); + verify(mockReceiver).get(ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG); } @Test @@ -486,7 +487,7 @@ public void testProcessElementSevereCommandError() { processFn.processElement(mockContext, mockReceiver); verify(mockShadowCollection, times(1)).find(mockSession, LOOKUP_BY_DOC_ID); - verify(mockReceiver).get(ProcessChangeEventFn.severeFailedWriteTag); + verify(mockReceiver).get(ProcessChangeEventFn.SEVERE_FAILED_WRITE_TAG); verify(mockSevereFailureReceiver, times(1)).output(any()); verify(mockSession, never()).commitTransaction(); } @@ -506,7 +507,7 @@ public void testProcessElement_updateEventWithNullDataSkips() { .replaceOne(any(), any(), any(), any(ReplaceOptions.class)); verify(mockSession).commitTransaction(); verify(mockSession, never()).abortTransaction(); - verify(mockReceiver).get(ProcessChangeEventFn.successfulWriteTag); + verify(mockReceiver).get(ProcessChangeEventFn.SUCCESSFUL_WRITE_TAG); verify(mockSuccessReceiver, times(1)).output(mockElement); MetricsContainerImpl container = @@ -539,7 +540,7 @@ public void testProcessElement_insertEventWithNullDataRoutesToSevereDlq() { verify(mockDataCollection, never()).replaceOne(any(), any(), any(), any(ReplaceOptions.class)); verify(mockShadowCollection, never()) .replaceOne(any(), any(), any(), any(ReplaceOptions.class)); - verify(mockReceiver).get(ProcessChangeEventFn.severeFailedWriteTag); + verify(mockReceiver).get(ProcessChangeEventFn.SEVERE_FAILED_WRITE_TAG); verify(mockSevereFailureReceiver, times(1)).output(any()); } diff --git a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/StatefulDeduplicationFnTest.java b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/StatefulDeduplicationFnTest.java new file mode 100644 index 0000000000..40cce5a732 --- /dev/null +++ b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/StatefulDeduplicationFnTest.java @@ -0,0 +1,376 @@ +/* + * 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 com.fasterxml.jackson.databind.ObjectMapper; +import com.google.cloud.teleport.v2.templates.datastream.MongoDbChangeEventContext; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.metrics.MetricNameFilter; +import org.apache.beam.sdk.metrics.MetricQueryResults; +import org.apache.beam.sdk.metrics.MetricsFilter; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.testing.TestStream; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.windowing.GlobalWindows; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.TimestampedValue; +import org.joda.time.Instant; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link StatefulDeduplicationFn}. */ +@RunWith(JUnit4.class) +public class StatefulDeduplicationFnTest { + + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + + private MongoDbChangeEventContext createEventContext( + String docId, long seconds, int nanos, String changeType, boolean isDlqReconsumed) + throws Exception { + return createEventContext(docId, seconds, nanos, changeType, isDlqReconsumed, "cdc"); + } + + private MongoDbChangeEventContext createEventContext( + String docId, + long seconds, + int nanos, + String changeType, + boolean isDlqReconsumed, + String readMethod) + throws Exception { + String payload = + String.format( + "{" + + "\"_metadata_source\": {\"collection\": \"users\"}," + + "\"_id\": \"\\\"%s\\\"\"," + + "\"_metadata_timestamp_seconds\": %d," + + "\"_metadata_timestamp_nanos\": %d," + + "\"_metadata_change_type\": \"%s\"," + + "\"_metadata_read_method\": \"%s\"" + + (isDlqReconsumed ? ",\"isDlqReconsumed\": \"true\"" : "") + + ",\"data\": \"{\\\"name\\\": \\\"user_%s\\\"}\"" + + "}", + docId, + seconds, + nanos, + changeType, + readMethod, + docId); + return new MongoDbChangeEventContext(OBJECT_MAPPER.readTree(payload), "shadow_"); + } + + @Test + public void testInOrderEvents_emitsAll() throws Exception { + MongoDbChangeEventContext event1 = createEventContext("doc1", 1000L, 100, "INSERT", false); + MongoDbChangeEventContext event2 = createEventContext("doc1", 1000L, 200, "UPDATE", false); + + TestStream> stream = + TestStream.create(KvCoder.of(StringUtf8Coder.of(), MongoDbChangeEventContextCoder.of())) + .addElements( + TimestampedValue.of(KV.of("users#doc1", event1), new Instant(100)), + TimestampedValue.of(KV.of("users#doc1", event2), new Instant(200))) + .advanceWatermarkToInfinity(); + + PCollection result = + pipeline + .apply(stream) + .apply(Window.into(new GlobalWindows())) + .apply(ParDo.of(new StatefulDeduplicationFn())); + + PAssert.that(result).containsInAnyOrder(event1, event2); + pipeline.run(); + } + + @Test + public void testOutOfOrderEvents_dropsStaleEvent() throws Exception { + MongoDbChangeEventContext eventNewer = createEventContext("doc1", 1000L, 200, "UPDATE", false); + MongoDbChangeEventContext eventStale = createEventContext("doc1", 1000L, 100, "INSERT", false); + + TestStream> stream = + TestStream.create(KvCoder.of(StringUtf8Coder.of(), MongoDbChangeEventContextCoder.of())) + .addElements(TimestampedValue.of(KV.of("users#doc1", eventNewer), new Instant(100))) + .addElements(TimestampedValue.of(KV.of("users#doc1", eventStale), new Instant(200))) + .advanceWatermarkToInfinity(); + + PCollection result = + pipeline + .apply(stream) + .apply(Window.into(new GlobalWindows())) + .apply(ParDo.of(new StatefulDeduplicationFn())); + + PAssert.that(result).containsInAnyOrder(eventNewer); + PipelineResult pipelineResult = pipeline.run(); + + MetricQueryResults metrics = + pipelineResult + .metrics() + .queryMetrics( + MetricsFilter.builder() + .addNameFilter( + MetricNameFilter.named(StatefulDeduplicationFn.class, "outOfOrderSkips")) + .build()); + + long count = 0; + if (metrics.getCounters().iterator().hasNext()) { + count = metrics.getCounters().iterator().next().getAttempted(); + } + assertEquals(1L, count); + } + + @Test + public void testDeleteAndStaleUpdate_preventsZombieResurrection() throws Exception { + MongoDbChangeEventContext insertEvent = createEventContext("doc1", 1000L, 100, "INSERT", false); + MongoDbChangeEventContext deleteEvent = createEventContext("doc1", 1000L, 300, "DELETE", false); + MongoDbChangeEventContext staleUpdate = createEventContext("doc1", 1000L, 200, "UPDATE", false); + + TestStream> stream = + TestStream.create(KvCoder.of(StringUtf8Coder.of(), MongoDbChangeEventContextCoder.of())) + .addElements(TimestampedValue.of(KV.of("users#doc1", insertEvent), new Instant(100))) + .addElements(TimestampedValue.of(KV.of("users#doc1", deleteEvent), new Instant(200))) + .addElements(TimestampedValue.of(KV.of("users#doc1", staleUpdate), new Instant(300))) + .advanceWatermarkToInfinity(); + + PCollection result = + pipeline + .apply(stream) + .apply(Window.into(new GlobalWindows())) + .apply(ParDo.of(new StatefulDeduplicationFn())); + + PAssert.that(result).containsInAnyOrder(insertEvent, deleteEvent); + PipelineResult pipelineResult = pipeline.run(); + + MetricQueryResults metrics = + pipelineResult + .metrics() + .queryMetrics( + MetricsFilter.builder() + .addNameFilter( + MetricNameFilter.named(StatefulDeduplicationFn.class, "outOfOrderSkips")) + .build()); + + long skips = 0; + if (metrics.getCounters().iterator().hasNext()) { + skips = metrics.getCounters().iterator().next().getAttempted(); + } + assertEquals(1L, skips); + } + + @Test + public void testDlqEqualTimestamp_passThrough() throws Exception { + MongoDbChangeEventContext regularEvent = + createEventContext("doc1", 1000L, 100, "INSERT", false); + MongoDbChangeEventContext dlqReconsumedEvent = + createEventContext("doc1", 1000L, 100, "INSERT", true); + + TestStream> stream = + TestStream.create(KvCoder.of(StringUtf8Coder.of(), MongoDbChangeEventContextCoder.of())) + .addElements(TimestampedValue.of(KV.of("users#doc1", regularEvent), new Instant(100))) + .addElements( + TimestampedValue.of(KV.of("users#doc1", dlqReconsumedEvent), new Instant(200))) + .advanceWatermarkToInfinity(); + + PCollection result = + pipeline + .apply(stream) + .apply(Window.into(new GlobalWindows())) + .apply(ParDo.of(new StatefulDeduplicationFn())); + + PAssert.that(result).containsInAnyOrder(regularEvent, dlqReconsumedEvent); + PipelineResult pipelineResult = pipeline.run(); + + MetricQueryResults metrics = + pipelineResult + .metrics() + .queryMetrics( + MetricsFilter.builder() + .addNameFilter( + MetricNameFilter.named( + StatefulDeduplicationFn.class, "dlqEqualTimestampPassThrough")) + .build()); + + long passThroughCount = 0; + if (metrics.getCounters().iterator().hasNext()) { + passThroughCount = metrics.getCounters().iterator().next().getAttempted(); + } + assertEquals(1L, passThroughCount); + } + + @Test + public void testDuplicateEqualTimestamp_nonDlq_dropped() throws Exception { + MongoDbChangeEventContext event1 = createEventContext("doc1", 1000L, 100, "INSERT", false); + MongoDbChangeEventContext duplicateEvent = + createEventContext("doc1", 1000L, 100, "INSERT", false); + + TestStream> stream = + TestStream.create(KvCoder.of(StringUtf8Coder.of(), MongoDbChangeEventContextCoder.of())) + .addElements(TimestampedValue.of(KV.of("users#doc1", event1), new Instant(100))) + .addElements(TimestampedValue.of(KV.of("users#doc1", duplicateEvent), new Instant(200))) + .advanceWatermarkToInfinity(); + + PCollection result = + pipeline + .apply(stream) + .apply(Window.into(new GlobalWindows())) + .apply(ParDo.of(new StatefulDeduplicationFn())); + + PAssert.that(result).containsInAnyOrder(event1); + pipeline.run(); + } + + @Test + public void testCompositeMapAndArrayDocumentId() throws Exception { + String mapPayload = + "{" + + "\"_metadata_source\": {\"collection\": \"accounts\"}," + + "\"_id\": \"{\\\"tenant\\\": \\\"t1\\\", \\\"account\\\": 12345}\"," + + "\"_metadata_timestamp_seconds\": 1000," + + "\"_metadata_timestamp_nanos\": 100," + + "\"_metadata_change_type\": \"INSERT\"," + + "\"data\": \"{\\\"balance\\\": 500}\"" + + "}"; + MongoDbChangeEventContext mapEvent = + new MongoDbChangeEventContext(OBJECT_MAPPER.readTree(mapPayload), "shadow_"); + + String arrayPayload = + "{" + + "\"_metadata_source\": {\"collection\": \"items\"}," + + "\"_id\": \"[1, \\\"partA\\\", 2]\"," + + "\"_metadata_timestamp_seconds\": 1000," + + "\"_metadata_timestamp_nanos\": 200," + + "\"_metadata_change_type\": \"UPDATE\"," + + "\"data\": \"{\\\"qty\\\": 10}\"" + + "}"; + MongoDbChangeEventContext arrayEvent = + new MongoDbChangeEventContext(OBJECT_MAPPER.readTree(arrayPayload), "shadow_"); + + String mapKey = "accounts#" + Utils.documentIdToString(mapEvent.getDocumentId()); + String arrayKey = "items#" + Utils.documentIdToString(arrayEvent.getDocumentId()); + + TestStream> stream = + TestStream.create(KvCoder.of(StringUtf8Coder.of(), MongoDbChangeEventContextCoder.of())) + .addElements(TimestampedValue.of(KV.of(mapKey, mapEvent), new Instant(100))) + .addElements(TimestampedValue.of(KV.of(arrayKey, arrayEvent), new Instant(200))) + .advanceWatermarkToInfinity(); + + PCollection result = + pipeline + .apply(stream) + .apply(Window.into(new GlobalWindows())) + .apply(ParDo.of(new StatefulDeduplicationFn())); + + PAssert.that(result).containsInAnyOrder(mapEvent, arrayEvent); + pipeline.run(); + } + + @Test + public void testBackfillSnapshotDoesNotResurrectCdcDelete() throws Exception { + // Exact 30-orphan scenario: CDC DELETE at ts=1786382543:4752 followed by Backfill READ at + // ts=1786382543:967669000 + MongoDbChangeEventContext cdcDelete = + createEventContext("orphanDoc", 1786382543L, 4752, "DELETE", false, "cdc"); + MongoDbChangeEventContext backfillRead = + createEventContext("orphanDoc", 1786382543L, 967669000, "READ", false, "backfill"); + + TestStream> stream = + TestStream.create(KvCoder.of(StringUtf8Coder.of(), MongoDbChangeEventContextCoder.of())) + .addElements(TimestampedValue.of(KV.of("users#orphanDoc", cdcDelete), new Instant(100))) + .addElements( + TimestampedValue.of(KV.of("users#orphanDoc", backfillRead), new Instant(200))) + .advanceWatermarkToInfinity(); + + PCollection result = + pipeline + .apply(stream) + .apply(Window.into(new GlobalWindows())) + .apply(ParDo.of(new StatefulDeduplicationFn())); + + // Only the CDC delete should be emitted; the backfill read snapshot must be dropped as stale + PAssert.that(result).containsInAnyOrder(cdcDelete); + PipelineResult pipelineResult = pipeline.run(); + + MetricQueryResults metrics = + pipelineResult + .metrics() + .queryMetrics( + MetricsFilter.builder() + .addNameFilter( + MetricNameFilter.named(StatefulDeduplicationFn.class, "outOfOrderSkips")) + .build()); + + long skips = 0; + if (metrics.getCounters().iterator().hasNext()) { + skips = metrics.getCounters().iterator().next().getAttempted(); + } + assertEquals(1L, skips); + } + + @Test + public void testCdcUpdateSupersedesConcurrentBackfillRead() throws Exception { + MongoDbChangeEventContext backfillRead = + createEventContext("doc1", 1786382543L, 967669000, "READ", false, "backfill"); + MongoDbChangeEventContext cdcUpdate = + createEventContext("doc1", 1786382543L, 100, "UPDATE", false, "cdc"); + + TestStream> stream = + TestStream.create(KvCoder.of(StringUtf8Coder.of(), MongoDbChangeEventContextCoder.of())) + .addElements(TimestampedValue.of(KV.of("users#doc1", backfillRead), new Instant(100))) + .addElements(TimestampedValue.of(KV.of("users#doc1", cdcUpdate), new Instant(200))) + .advanceWatermarkToInfinity(); + + PCollection result = + pipeline + .apply(stream) + .apply(Window.into(new GlobalWindows())) + .apply(ParDo.of(new StatefulDeduplicationFn())); + + PAssert.that(result).containsInAnyOrder(backfillRead, cdcUpdate); + pipeline.run(); + } + + @Test + public void testDeleteAndRecreate_allowsNewerInsert() throws Exception { + MongoDbChangeEventContext insertV1 = createEventContext("doc1", 1000L, 100, "INSERT", false); + MongoDbChangeEventContext deleteV2 = createEventContext("doc1", 1000L, 200, "DELETE", false); + MongoDbChangeEventContext insertV3 = createEventContext("doc1", 1001L, 50, "INSERT", false); + + TestStream> stream = + TestStream.create(KvCoder.of(StringUtf8Coder.of(), MongoDbChangeEventContextCoder.of())) + .addElements(TimestampedValue.of(KV.of("users#doc1", insertV1), new Instant(100))) + .addElements(TimestampedValue.of(KV.of("users#doc1", deleteV2), new Instant(200))) + .addElements(TimestampedValue.of(KV.of("users#doc1", insertV3), new Instant(300))) + .advanceWatermarkToInfinity(); + + PCollection result = + pipeline + .apply(stream) + .apply(Window.into(new GlobalWindows())) + .apply(ParDo.of(new StatefulDeduplicationFn())); + + PAssert.that(result).containsInAnyOrder(insertV1, deleteV2, insertV3); + pipeline.run(); + } +} diff --git a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/ThrottledLoggerTest.java b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/ThrottledLoggerTest.java new file mode 100644 index 0000000000..375ca72b06 --- /dev/null +++ b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/ThrottledLoggerTest.java @@ -0,0 +1,177 @@ +/* + * 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.assertTrue; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.ObjectInputStream; +import java.io.ObjectOutputStream; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link ThrottledLogger}. */ +@RunWith(JUnit4.class) +public class ThrottledLoggerTest { + + @Test + public void testBasicRecordingAndGetters() { + ThrottledLogger logger = new ThrottledLogger("TestComponent", 30_000L); + logger.recordRetryableError("NETWORK", "Connection reset"); + logger.recordRetryableError("NETWORK", "Timeout"); + logger.recordSevereError("VALIDATION", "Invalid schema"); + + assertEquals(3L, logger.getTotalErrors()); + assertEquals(2L, logger.getTotalRetryable()); + assertEquals(1L, logger.getTotalSevere()); + } + + @Test + public void testFlushSummaryResetsCounts() { + ThrottledLogger logger = new ThrottledLogger("TestComponent", 30_000L); + logger.recordRetryableError("NETWORK", "Connection reset"); + logger.recordSevereError("SCHEMA", "Bad field"); + + assertEquals(2L, logger.getTotalErrors()); + logger.flushSummary(); + assertEquals(0L, logger.getTotalErrors()); + assertEquals(0L, logger.getTotalRetryable()); + assertEquals(0L, logger.getTotalSevere()); + } + + @Test + public void testShouldLogThrottling() { + ThrottledLogger logger = new ThrottledLogger("TestComponent", 5000L); + assertTrue(logger.shouldLog("key1")); + assertFalse(logger.shouldLog("key1")); + assertFalse(logger.shouldLog("key1")); + + long suppressed = logger.getAndResetSuppressedCount("key1"); + assertEquals(2L, suppressed); + assertEquals(0L, logger.getAndResetSuppressedCount("key1")); + } + + @Test + public void testSerializationAndConcurrentAccess() throws Exception { + ThrottledLogger original = new ThrottledLogger("ConcurrentLogger", 30_000L); + + // Serialize and deserialize to simulate Beam worker distribution where transient fields start + // null + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try (ObjectOutputStream oos = new ObjectOutputStream(baos)) { + oos.writeObject(original); + } + + ByteArrayInputStream bais = new ByteArrayInputStream(baos.toByteArray()); + ThrottledLogger deserialized; + try (ObjectInputStream ois = new ObjectInputStream(bais)) { + deserialized = (ThrottledLogger) ois.readObject(); + } + + int threadCount = 20; + int incrementsPerThread = 1000; + ExecutorService executor = Executors.newFixedThreadPool(threadCount); + CountDownLatch latch = new CountDownLatch(threadCount); + + for (int i = 0; i < threadCount; i++) { + final int threadId = i; + executor.submit( + () -> { + try { + for (int j = 0; j < incrementsPerThread; j++) { + if (threadId % 2 == 0) { + deserialized.recordRetryableError("CAT_RETRY", "retry error"); + } else { + deserialized.recordSevereError("CAT_SEVERE", "severe error"); + } + } + } finally { + latch.countDown(); + } + }); + } + + assertTrue(latch.await(10, TimeUnit.SECONDS)); + executor.shutdown(); + + long expectedTotal = (long) threadCount * incrementsPerThread; + assertEquals(expectedTotal, deserialized.getTotalErrors()); + assertEquals(expectedTotal / 2, deserialized.getTotalRetryable()); + assertEquals(expectedTotal / 2, deserialized.getTotalSevere()); + } + + @Test + public void testNullCategoryAndKeySafety() { + ThrottledLogger logger = new ThrottledLogger("NullSafetyLogger", 5000L); + logger.recordRetryableError(null, "Null category message"); + logger.recordSevereError(null, "Null category severe message"); + + assertEquals(2L, logger.getTotalErrors()); + assertTrue(logger.shouldLog(null)); + assertFalse(logger.shouldLog(null)); + assertEquals(1L, logger.getAndResetSuppressedCount(null)); + } + + @Test + public void testLogInfoWarnErrorMethods() { + ThrottledLogger logger = new ThrottledLogger("LogMethodsLogger", 5000L); + org.slf4j.Logger mockLogger = org.mockito.Mockito.mock(org.slf4j.Logger.class); + + logger.logInfo(mockLogger, "infoKey", "info message"); + logger.logInfo(mockLogger, "infoKey", "info message 2"); + org.mockito.Mockito.verify(mockLogger, org.mockito.Mockito.times(1)) + .info("info message", new Object[] {}); + + logger.logWarn(mockLogger, "warnKey", "warn message"); + logger.logWarn(mockLogger, "warnKey", "warn message 2"); + org.mockito.Mockito.verify(mockLogger, org.mockito.Mockito.times(1)) + .warn("warn message", new Object[] {}); + + logger.logError(mockLogger, "errorKey", "error message"); + logger.logError(mockLogger, "errorKey", "error message 2"); + org.mockito.Mockito.verify(mockLogger, org.mockito.Mockito.times(1)) + .error("error message", new Object[] {}); + } + + @Test + public void testCategoryAndKeyOverflowRoutesToOther() { + ThrottledLogger logger = new ThrottledLogger("OverflowLogger", 5000L); + + // Register 250 distinct categories (limit is 200) + for (int i = 0; i < 250; i++) { + logger.recordRetryableError("CAT_" + i, "error msg " + i); + } + assertEquals(250L, logger.getTotalErrors()); + + // Register 250 distinct log keys (limit is 200) + int loggedCount = 0; + for (int i = 0; i < 250; i++) { + if (logger.shouldLog("KEY_" + i)) { + loggedCount++; + } + } + // The first 200 unique keys + 1 for "OTHER" log once; remaining 49 are suppressed under "OTHER" + assertEquals(201, loggedCount); + } +} diff --git a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/TimestampSortKeyCoderTest.java b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/TimestampSortKeyCoderTest.java new file mode 100644 index 0000000000..7114eab5f9 --- /dev/null +++ b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/TimestampSortKeyCoderTest.java @@ -0,0 +1,104 @@ +/* + * 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 java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import org.apache.beam.sdk.testing.CoderProperties; +import org.apache.beam.sdk.util.CoderUtils; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link TimestampSortKeyCoder}. */ +@RunWith(JUnit4.class) +public class TimestampSortKeyCoderTest { + + private final TimestampSortKeyCoder coder = TimestampSortKeyCoder.of(); + + @Test + public void testEncodeDecodeRoundTrip_cdcEvent() throws Exception { + TimestampSortKey key = TimestampSortKey.of(1786382543L, 4752, true); + TimestampSortKey decoded = CoderUtils.clone(coder, key); + + assertEquals(key, decoded); + assertEquals(1786382543L, decoded.getSeconds()); + assertEquals(4752, decoded.getSubSeconds()); + assertEquals(true, decoded.isCdc()); + CoderProperties.coderDecodeEncodeEqual(coder, key); + } + + @Test + public void testEncodeDecodeRoundTrip_backfillEvent() throws Exception { + TimestampSortKey key = TimestampSortKey.of(1786382543L, 967669000, false); + TimestampSortKey decoded = CoderUtils.clone(coder, key); + + assertEquals(key, decoded); + assertEquals(1786382543L, decoded.getSeconds()); + assertEquals(967669000, decoded.getSubSeconds()); + assertEquals(false, decoded.isCdc()); + CoderProperties.coderDecodeEncodeEqual(coder, key); + } + + @Test + public void testEncodeDecodeRoundTrip_boundaryValues() throws Exception { + TimestampSortKey minKey = TimestampSortKey.of(0L, 0L, false); + CoderProperties.coderDecodeEncodeEqual(coder, minKey); + + TimestampSortKey maxKey = TimestampSortKey.of(Long.MAX_VALUE, Long.MAX_VALUE, true); + CoderProperties.coderDecodeEncodeEqual(coder, maxKey); + + TimestampSortKey largeSubSecKey = TimestampSortKey.of(1786382543L, 5_000_000_000L, true); + TimestampSortKey decoded = CoderUtils.clone(coder, largeSubSecKey); + assertEquals(largeSubSecKey, decoded); + assertEquals(5_000_000_000L, decoded.getSubSeconds()); + } + + @Test + public void testNullValueEncoding() throws Exception { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + coder.encode(null, out); + byte[] bytes = out.toByteArray(); + + assertEquals(1, bytes.length); + + ByteArrayInputStream in = new ByteArrayInputStream(bytes); + TimestampSortKey decoded = coder.decode(in); + assertNull(decoded); + } + + @Test + public void testVerifyDeterministic() throws Exception { + coder.verifyDeterministic(); + + TimestampSortKey key1 = TimestampSortKey.of(1000L, 50L, true); + TimestampSortKey key2 = TimestampSortKey.of(1000L, 50L, true); + CoderProperties.coderDeterministic(coder, key1, key2); + } + + @Test + public void testBinaryEncodingExactSize() throws Exception { + TimestampSortKey key = TimestampSortKey.of(1786382543L, 4752L, true); + ByteArrayOutputStream out = new ByteArrayOutputStream(); + coder.encode(key, out); + + // 1 byte presence + 8 bytes long + 8 bytes long + 1 byte boolean = 18 bytes + assertEquals(18, out.toByteArray().length); + } +} diff --git a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/TimestampSortKeyTest.java b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/TimestampSortKeyTest.java new file mode 100644 index 0000000000..2f14f570df --- /dev/null +++ b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/TimestampSortKeyTest.java @@ -0,0 +1,107 @@ +/* + * 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.assertNotEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link TimestampSortKey}. */ +@RunWith(JUnit4.class) +public class TimestampSortKeyTest { + + @Test + public void testNullEventReturnsNullKey() { + assertNull(TimestampSortKey.of(null)); + } + + @Test(expected = NullPointerException.class) + public void testCompareTo_nullThrowsNullPointerException() { + TimestampSortKey t1 = new TimestampSortKey(1000L, 500L, true); + t1.compareTo(null); + } + + @Test + public void testCompareTo_higherSecondsWins() { + TimestampSortKey t1 = new TimestampSortKey(1000L, 500L, true); + TimestampSortKey t2 = new TimestampSortKey(1001L, 100L, true); + + assertTrue(t2.compareTo(t1) > 0); + assertTrue(t1.compareTo(t2) < 0); + } + + @Test + public void testCompareTo_sameSecondCdcBeatsBackfill() { + TimestampSortKey cdcKey = new TimestampSortKey(1000L, 10L, true); + TimestampSortKey backfillKey = new TimestampSortKey(1000L, 999999999L, false); + + assertTrue(cdcKey.compareTo(backfillKey) > 0); + assertTrue(backfillKey.compareTo(cdcKey) < 0); + } + + @Test + public void testCompareTo_sameSecondSameStreamTypeComparesSubSeconds() { + TimestampSortKey cdc1 = new TimestampSortKey(1000L, 10L, true); + TimestampSortKey cdc2 = new TimestampSortKey(1000L, 20L, true); + + assertTrue(cdc2.compareTo(cdc1) > 0); + assertTrue(cdc1.compareTo(cdc2) < 0); + + TimestampSortKey bf1 = new TimestampSortKey(1000L, 100L, false); + TimestampSortKey bf2 = new TimestampSortKey(1000L, 200L, false); + + assertTrue(bf2.compareTo(bf1) > 0); + assertTrue(bf1.compareTo(bf2) < 0); + } + + @Test + public void testCompareTo_equalKeysReturnZero() { + TimestampSortKey k1 = new TimestampSortKey(1000L, 100L, true); + TimestampSortKey k2 = new TimestampSortKey(1000L, 100L, true); + + assertEquals(0, k1.compareTo(k2)); + assertEquals(0, k2.compareTo(k1)); + } + + @Test + public void testEqualsAndHashCode() { + TimestampSortKey k1 = new TimestampSortKey(1000L, 100L, true); + TimestampSortKey k2 = new TimestampSortKey(1000L, 100L, true); + TimestampSortKey k3 = new TimestampSortKey(1000L, 100L, false); + TimestampSortKey k4 = new TimestampSortKey(1001L, 100L, true); + + assertEquals(k1, k2); + assertEquals(k1.hashCode(), k2.hashCode()); + assertNotEquals(k1, k3); + assertNotEquals(k1, k4); + assertNotEquals(k1, null); + assertNotEquals(k1, "otherType"); + } + + @Test + public void testToString() { + TimestampSortKey cdcKey = new TimestampSortKey(1000L, 100L, true); + TimestampSortKey backfillKey = new TimestampSortKey(1000L, 500L, false); + + assertEquals("1000:100:cdc", cdcKey.toString()); + assertEquals("1000:500:backfill", backfillKey.toString()); + } +} diff --git a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/UtilsTest.java b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/UtilsTest.java index 0bee3cc76d..6cde98b437 100644 --- a/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/UtilsTest.java +++ b/v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/transforms/UtilsTest.java @@ -17,6 +17,7 @@ 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 com.fasterxml.jackson.databind.JsonNode; @@ -25,6 +26,7 @@ import com.google.cloud.teleport.v2.templates.datastream.DatastreamConstants; import com.google.cloud.teleport.v2.templates.datastream.MongoDbChangeEventContext; import java.util.HashSet; +import java.util.List; import java.util.Set; import org.bson.Document; import org.bson.types.Binary; @@ -135,6 +137,34 @@ public void testIsNewerTimestamp_sameTimestamp() { assertFalse(Utils.isNewerTimestamp(ts1, ts2)); } + @Test + public void testIsNewerTimestamp_nullSafety() { + Document ts1 = + new Document(MongoDbChangeEventContext.TIMESTAMP_SECONDS_COL, 1L) + .append(MongoDbChangeEventContext.TIMESTAMP_NANOS_COL, 100); + + assertTrue(Utils.isNewerTimestamp(ts1, null)); + assertFalse(Utils.isNewerTimestamp(null, ts1)); + assertFalse(Utils.isNewerTimestamp(null, null)); + } + + @Test + public void testGetTimestampNanos() { + assertEquals(0L, Utils.getTimestampNanos(null)); + assertEquals(0L, Utils.getTimestampNanos(new Document())); + + Document validTs = + new Document(MongoDbChangeEventContext.TIMESTAMP_SECONDS_COL, 1700000000L) + .append(MongoDbChangeEventContext.TIMESTAMP_NANOS_COL, 500); + assertEquals(1700000000000000500L, Utils.getTimestampNanos(validTs)); + + // Number types as Integer / Double + Document intTs = + new Document(MongoDbChangeEventContext.TIMESTAMP_SECONDS_COL, 100) + .append(MongoDbChangeEventContext.TIMESTAMP_NANOS_COL, 20); + assertEquals(100000000020L, Utils.getTimestampNanos(intTs)); + } + @Test public void testJsonToDocument() { String jsonString = @@ -172,20 +202,34 @@ public void testJsonToDocument() { @Test public void testDocumentIdToString() { - assertEquals("test_id", Utils.documentIdToString("test_id")); - assertEquals("123", Utils.documentIdToString(123L)); - assertEquals("123.456", Utils.documentIdToString(123.456)); - assertEquals("true", Utils.documentIdToString(true)); + assertEquals("str_test_id", Utils.documentIdToString("test_id")); + assertEquals("str_123", Utils.documentIdToString("123")); + assertEquals("i64_123", Utils.documentIdToString(123L)); + assertEquals("i32_123", Utils.documentIdToString(123)); + assertEquals("f64_123.456", Utils.documentIdToString(123.456)); + assertEquals("bool_true", Utils.documentIdToString(true)); + assertEquals("bool_false", Utils.documentIdToString(false)); assertEquals("null", Utils.documentIdToString(null)); + // Verify string "123" vs Long 123L vs Integer 123 are distinct and do not collide + assertFalse(Utils.documentIdToString("123").equals(Utils.documentIdToString(123L))); + assertFalse(Utils.documentIdToString("123").equals(Utils.documentIdToString(123))); + assertFalse(Utils.documentIdToString(123L).equals(Utils.documentIdToString(123))); + ObjectId objectId = new ObjectId("645c9a7e7b8b1a0e9c0f8b3a"); - assertEquals("645c9a7e7b8b1a0e9c0f8b3a", Utils.documentIdToString(objectId)); + assertEquals("oid_645c9a7e7b8b1a0e9c0f8b3a", Utils.documentIdToString(objectId)); Binary binary = new Binary(new byte[] {1, 2, 3}); - assertEquals("AQID", Utils.documentIdToString(binary)); + assertEquals("bin_0_AQID", Utils.documentIdToString(binary)); Document doc = new Document("a", 1).append("b", "test"); - assertEquals("{\"a\": 1, \"b\": \"test\"}", Utils.documentIdToString(doc)); + assertEquals( + "doc_{\"a\": {\"$numberInt\": \"1\"}, \"b\": \"test\"}", Utils.documentIdToString(doc)); + + List list = java.util.Arrays.asList(1, "partA", 2); + assertEquals( + "list_{\"arr\": [{\"$numberInt\": \"1\"}, \"partA\", {\"$numberInt\": \"2\"}]}", + Utils.documentIdToString(list)); } @Test @@ -249,4 +293,35 @@ public void testExtractInnerEvent_JsonNode() throws Exception { // Test case where the JSON node is NOT wrapped assertEquals(inner, Utils.extractInnerEvent(inner)); } + + @Test + public void testJsonToDocument_preservesCustomerDataField() { + String jsonString = + "{\"_id\":\"{\\\"$oid\\\": \\\"6a7c79fbfd2b2eb4aca6d225\\\"}\"," + + "\"data\":\"{\\\"_id\\\":{\\\"$oid\\\": \\\"6a7c79fbfd2b2eb4aca6d225\\\"}," + + "\\\"value\\\":0.7018,\\\"version\\\":14,\\\"data\\\":\\\"xxxxxxxxxxxxxxxxxxxx\\\"}\"," + + "\"_metadata_stream\":\"test-stream\",\"_metadata_timestamp\":1745957556}"; + + Document result = Utils.jsonToDocument(jsonString, "6a7c79fbfd2b2eb4aca6d225"); + assertNotNull(result); + assertEquals("6a7c79fbfd2b2eb4aca6d225", result.get("_id")); + assertEquals(0.7018, result.getDouble("value"), 1e-4); + assertEquals(14, result.getInteger("version").intValue()); + assertEquals("xxxxxxxxxxxxxxxxxxxx", result.getString("data")); + assertFalse(result.containsKey("_metadata_stream")); + assertFalse(result.containsKey("_metadata_timestamp")); + } + + @Test + public void testJsonToDocument_preservesDataFieldInObjectFormat() { + String jsonString = + "{\"data\":{\"name\":\"John\",\"data\":\"custom_payload_data\",\"_metadata_stream\":\"str\"}}"; + Document result = Utils.jsonToDocument(jsonString, "id1"); + + assertNotNull(result); + assertEquals("John", result.get("name")); + assertEquals("custom_payload_data", result.get("data")); + assertEquals("id1", result.get("_id")); + assertFalse(result.containsKey("_metadata_stream")); + } }