Skip to content

Commit 1379803

Browse files
namita-lnamita-l
authored
Refactor DLQ retry logic and preserve original event in MongoDB to Firestore template (#3808)
1 parent a292e17 commit 1379803

7 files changed

Lines changed: 200 additions & 89 deletions

File tree

‎v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/templates/DataStreamMongoDBToFirestore.java‎

Lines changed: 10 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1456,13 +1456,6 @@ public void processElement(ProcessContext c) {
14561456
// Handle events wrapped in a changeEvent field (common for reconsumed DLQ records)
14571457
Document innerDoc = Utils.extractInnerEvent(doc);
14581458

1459-
// Skip UDF if it has already been processed (e.g. on DLQ retry)
1460-
Boolean udfProcessed = innerDoc.getBoolean("_metadata_udf_processed");
1461-
if (udfProcessed != null && udfProcessed) {
1462-
c.output(BYPASS_UDF_TAG, element);
1463-
return;
1464-
}
1465-
14661459
String changeType = innerDoc.getString(DatastreamConstants.EVENT_CHANGE_TYPE_KEY);
14671460
if (changeType == null) {
14681461
changeType = "";
@@ -1516,10 +1509,17 @@ public void processElement(ProcessContext c) {
15161509
// This ensures we don't write invalid data to the destination.
15171510
Document.parse(transformedData);
15181511

1519-
((ObjectNode) fullEventNode).put(MongoDbChangeEventContext.DATA_COL, transformedData);
1512+
// Merge the transformed data back into the full event JSON.
1513+
// The payload of the output FailsafeElement will contain this modified event JSON,
1514+
// which is used by downstream steps (like writing to Firestore) to process the update.
1515+
// The originalPayload remains the raw event for safety and DLQ purposes.
1516+
JsonNode targetNode = Utils.extractInnerEvent(fullEventNode);
1517+
((ObjectNode) targetNode).put(MongoDbChangeEventContext.DATA_COL, transformedData);
15201518

15211519
String modifiedEventJson = OBJECT_MAPPER.writeValueAsString(fullEventNode);
1522-
c.output(FailsafeElement.of(modifiedEventJson, modifiedEventJson));
1520+
// Output the element with the preserved original payload and the modified event JSON
1521+
// containing UDF output.
1522+
c.output(FailsafeElement.of(element.getOriginalPayload(), modifiedEventJson));
15231523
} catch (Exception e) {
15241524
LOG.error("Error restoring UDF output, exception: {}", e.getMessage(), e);
15251525
FailsafeElement<String, String> failedElement =
@@ -1531,22 +1531,6 @@ public void processElement(ProcessContext c) {
15311531
}
15321532
}
15331533

1534-
public static class MarkUdfProcessedFn
1535-
extends DoFn<FailsafeElement<String, String>, FailsafeElement<String, String>> {
1536-
@ProcessElement
1537-
public void processElement(ProcessContext c) {
1538-
FailsafeElement<String, String> element = c.element();
1539-
try {
1540-
JsonNode node = OBJECT_MAPPER.readTree(element.getPayload());
1541-
((ObjectNode) node).put("_metadata_udf_processed", true);
1542-
String json = OBJECT_MAPPER.writeValueAsString(node);
1543-
c.output(FailsafeElement.of(json, json));
1544-
} catch (Exception e) {
1545-
throw new RuntimeException(e);
1546-
}
1547-
}
1548-
}
1549-
15501534
public static class ApplyUdfToDataField
15511535
extends PTransform<
15521536
PCollection<FailsafeElement<String, String>>,
@@ -1613,13 +1597,9 @@ public PCollection<FailsafeElement<String, String>> expand(
16131597
.get(UDF_SUCCESS_TAG)
16141598
.setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()));
16151599

1616-
// Mark as UDF processed to avoid re-processing on DLQ retries
1617-
PCollection<FailsafeElement<String, String>> markedOutput =
1618-
restoredOutput.apply("Mark UDF Processed", ParDo.of(new MarkUdfProcessedFn()));
1619-
16201600
// Merge the restored UDF output with the events that bypassed the UDF.
16211601
// Both streams now contain the full event JSON in the required format for downstream steps.
1622-
return PCollectionList.of(markedOutput)
1602+
return PCollectionList.of(restoredOutput)
16231603
.and(bypassedElements)
16241604
.apply("Merge Streams", Flatten.pCollections());
16251605
}

‎v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/templates/datastream/MongoDbChangeEventContext.java‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ public class MongoDbChangeEventContext implements Serializable {
5050
public static final String OID_FIELD_NAME = "$oid";
5151

5252
private final JsonNode changeEvent;
53+
private final JsonNode originalChangeEvent;
5354
private final String dataCollection;
5455
private final String shadowCollection;
5556
private final Object documentId;
@@ -87,7 +88,16 @@ private boolean isDlqReconsumed(JsonNode changeEvent) {
8788

8889
public MongoDbChangeEventContext(JsonNode payload, String shadowCollectionPrefix)
8990
throws JsonProcessingException {
91+
this(payload, payload, shadowCollectionPrefix);
92+
}
93+
94+
public MongoDbChangeEventContext(
95+
JsonNode payload, JsonNode originalPayload, String shadowCollectionPrefix)
96+
throws JsonProcessingException {
97+
// Extracts the actual change event. If wrapped in a DLQ structure like {"changeEvent": {...}},
98+
// it extracts the inner object.
9099
this.changeEvent = Utils.extractInnerEvent(payload);
100+
this.originalChangeEvent = Utils.extractInnerEvent(originalPayload);
91101

92102
this.retryCount =
93103
changeEvent.has(DatastreamConstants.RETRY_COUNT)
@@ -198,6 +208,10 @@ public JsonNode getChangeEvent() {
198208
return changeEvent;
199209
}
200210

211+
public JsonNode getOriginalChangeEvent() {
212+
return originalChangeEvent;
213+
}
214+
201215
public String getDataCollection() {
202216
return dataCollection;
203217
}

‎v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/CreateMongoDbChangeEventContextFn.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,9 +52,10 @@ public CreateMongoDbChangeEventContextFn(String shadowCollectionPrefix) {
5252
public void processElement(ProcessContext context, MultiOutputReceiver out) {
5353
FailsafeElement<String, String> element = context.element();
5454
try {
55-
JsonNode jsonNode = OBJECT_MAPPER.readTree(element.getOriginalPayload());
55+
JsonNode jsonNode = OBJECT_MAPPER.readTree(element.getPayload());
56+
JsonNode originalNode = OBJECT_MAPPER.readTree(element.getOriginalPayload());
5657
MongoDbChangeEventContext changeEventContext =
57-
new MongoDbChangeEventContext(jsonNode, shadowCollectionPrefix);
58+
new MongoDbChangeEventContext(jsonNode, originalNode, shadowCollectionPrefix);
5859
out.get(successfulCreationTag).output(changeEventContext);
5960
} catch (Exception e) {
6061
LOG.error("Error creating MongoDbChangeEventContext, exception: {}, element: {}", e, element);

‎v2/datastream-mongodb-to-firestore/src/main/java/com/google/cloud/teleport/v2/transforms/MongoDbEventDeadLetterQueueSanitizer.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,8 +44,10 @@ public String getJsonMessage(
4444
// Serialize the change event to JSON
4545
ObjectNode jsonNode = OBJECT_MAPPER.createObjectNode();
4646

47-
// Add the original change event JSON
48-
jsonNode.set("changeEvent", changeEvent.getChangeEvent());
47+
// Add the original change event JSON. This preserves the raw source data in the
48+
// DLQ, ensuring that retry attempts can re-apply updated UDF logic to the original
49+
// payload.
50+
jsonNode.set("changeEvent", changeEvent.getOriginalChangeEvent());
4951

5052
// Add other important fields
5153
jsonNode.put("dataCollection", changeEvent.getDataCollection());

‎v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/templates/DataStreamMongoDBToFirestoreTest.java‎

Lines changed: 116 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,60 @@ public void prepareUdfInputFn_validJsonWithData() {
206206
pipeline.run();
207207
}
208208

209+
@Test
210+
public void prepareUdfInputFn_wrappedPayload() {
211+
// Test that PrepareUdfInputFn correctly extracts the inner event from a wrapped DLQ payload.
212+
// This simulates an event that has failed before and is being retried from the Dead Letter
213+
// Queue.
214+
// The structure contains the event under "changeEvent" and metadata like "dataCollection".
215+
String fullEventJson =
216+
"{\"changeEvent\":{\"data\":\"{\\\"name\\\":\\\"John\\\"}\"},\"dataCollection\":\"test\"}";
217+
218+
// We pass the fullEventJson as both payload and originalPayload, which is typical for the start
219+
// of a retry flow.
220+
FailsafeElement<String, String> input = FailsafeElement.of(fullEventJson, fullEventJson);
221+
222+
PCollectionTuple result =
223+
pipeline
224+
.apply(
225+
"CreateInput",
226+
Create.of(input)
227+
.withCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))
228+
.apply(
229+
"PrepareUdfInput",
230+
ParDo.of(new DataStreamMongoDBToFirestore.PrepareUdfInputFn())
231+
.withOutputTags(
232+
SUCCESS_TAG,
233+
TupleTagList.of(DataStreamMongoDBToFirestore.PREPARE_FAILURE_TAG)));
234+
235+
result
236+
.get(SUCCESS_TAG)
237+
.setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()));
238+
result
239+
.get(DataStreamMongoDBToFirestore.PREPARE_FAILURE_TAG)
240+
.setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()));
241+
242+
PAssert.that(result.get(SUCCESS_TAG))
243+
.satisfies(
244+
iterable -> {
245+
FailsafeElement<String, String> element = iterable.iterator().next();
246+
// Verify that the original payload is preserved intact for DLQ purposes.
247+
assertEquals(fullEventJson, element.getOriginalPayload());
248+
try {
249+
ObjectMapper mapper = new ObjectMapper();
250+
JsonNode node = mapper.readTree(element.getPayload());
251+
// Verify that PrepareUdfInputFn successfully extracted the inner "data" field
252+
// from the wrapped "changeEvent", rather than passing the whole wrapper to UDF.
253+
assertEquals("John", node.get("name").asText());
254+
} catch (Exception e) {
255+
throw new RuntimeException(e);
256+
}
257+
return null;
258+
});
259+
260+
pipeline.run();
261+
}
262+
209263
@Test
210264
public void prepareUdfInputFn_missingData_routesToDlq() {
211265
String fullEventJson = "{\"op\":\"u\"}";
@@ -371,47 +425,6 @@ public void prepareUdfInputFn_wrappedEvent_succeeds() {
371425
pipeline.run();
372426
}
373427

374-
@Test
375-
public void prepareUdfInputFn_alreadyProcessed_bypassesUdf() {
376-
String processedJson = "{\"_metadata_udf_processed\": true, \"data\": \"{}\"}";
377-
FailsafeElement<String, String> input = FailsafeElement.of(processedJson, processedJson);
378-
379-
PCollectionTuple result =
380-
pipeline
381-
.apply(
382-
"CreateInput",
383-
Create.of(input)
384-
.withCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))
385-
.apply(
386-
"PrepareUdfInput",
387-
ParDo.of(new DataStreamMongoDBToFirestore.PrepareUdfInputFn())
388-
.withOutputTags(
389-
SUCCESS_TAG,
390-
TupleTagList.of(DataStreamMongoDBToFirestore.BYPASS_UDF_TAG)
391-
.and(DataStreamMongoDBToFirestore.PREPARE_FAILURE_TAG)));
392-
393-
result
394-
.get(SUCCESS_TAG)
395-
.setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()));
396-
result
397-
.get(DataStreamMongoDBToFirestore.BYPASS_UDF_TAG)
398-
.setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()));
399-
result
400-
.get(DataStreamMongoDBToFirestore.PREPARE_FAILURE_TAG)
401-
.setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()));
402-
403-
PAssert.that(result.get(DataStreamMongoDBToFirestore.BYPASS_UDF_TAG))
404-
.satisfies(
405-
iterable -> {
406-
FailsafeElement<String, String> element = iterable.iterator().next();
407-
assertEquals(processedJson, element.getOriginalPayload());
408-
assertEquals(processedJson, element.getPayload());
409-
return null;
410-
});
411-
412-
pipeline.run();
413-
}
414-
415428
@Test
416429
public void prepareUdfInputFn_invalidJson_routesToDlq() {
417430
String invalidJson = "{invalid}";
@@ -483,7 +496,6 @@ public void restoreUdfOutputFn_unchangedPayload() {
483496
try {
484497
ObjectMapper mapper = new ObjectMapper();
485498
JsonNode node = mapper.readTree(element.getPayload());
486-
assertFalse(node.has("_metadata_udf_processed"));
487499
assertEquals("{\"name\":\"John\"}", node.get("data").asText());
488500
} catch (Exception e) {
489501
throw new RuntimeException(e);
@@ -573,11 +585,70 @@ public void restoreUdfOutputFn_changedPayload() {
573585
JsonNode dataNode = mapper.readTree(dataStr);
574586
assertEquals("Jane", dataNode.get("name").asText());
575587

576-
// Verify that originalPayload is ALSO overwritten with the transformed data
588+
// Verify that originalPayload is PRESERVED and not overwritten
577589
JsonNode originalNode = mapper.readTree(element.getOriginalPayload());
578-
String originalDataStr = originalNode.get("data").asText();
590+
JsonNode originalDataNode = originalNode.get("data");
591+
assertEquals("John", originalDataNode.get("name").asText());
592+
} catch (Exception e) {
593+
throw new RuntimeException(e);
594+
}
595+
return null;
596+
});
597+
598+
pipeline.run();
599+
}
600+
601+
@Test
602+
public void restoreUdfOutputFn_wrappedPayload() {
603+
// Test that RestoreUdfOutputFn correctly merges transformed data back into the wrapped event
604+
// (DLQ format). This ensures that when we retry an event from DLQ and apply UDF, the result
605+
// is correctly placed back into the wrapped structure for further processing.
606+
String fullEventJson =
607+
"{\"changeEvent\":{\"data\":\"{\\\"name\\\":\\\"John\\\"}\"},\"dataCollection\":\"test\"}";
608+
String transformedData = "{\"name\":\"Jane\"}";
609+
FailsafeElement<String, String> input = FailsafeElement.of(fullEventJson, transformedData);
610+
611+
PCollectionTuple result =
612+
pipeline
613+
.apply(
614+
"CreateInput",
615+
Create.of(input)
616+
.withCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))
617+
.apply(
618+
"RestoreUdfOutput",
619+
ParDo.of(new DataStreamMongoDBToFirestore.RestoreUdfOutputFn())
620+
.withOutputTags(
621+
SUCCESS_TAG,
622+
TupleTagList.of(DataStreamMongoDBToFirestore.RESTORE_FAILURE_TAG)));
623+
624+
result
625+
.get(SUCCESS_TAG)
626+
.setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()));
627+
result
628+
.get(DataStreamMongoDBToFirestore.RESTORE_FAILURE_TAG)
629+
.setCoder(FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()));
630+
631+
PAssert.that(result.get(SUCCESS_TAG))
632+
.satisfies(
633+
iterable -> {
634+
FailsafeElement<String, String> element = iterable.iterator().next();
635+
try {
636+
ObjectMapper mapper = new ObjectMapper();
637+
JsonNode node = mapper.readTree(element.getPayload());
638+
639+
// Verify that the inner changeEvent has the updated data after UDF application.
640+
JsonNode innerEvent = node.get("changeEvent");
641+
String dataStr = innerEvent.get("data").asText();
642+
JsonNode dataNode = mapper.readTree(dataStr);
643+
assertEquals("Jane", dataNode.get("name").asText());
644+
645+
// Verify that originalPayload is PRESERVED as the original wrapped event,
646+
// ensuring auditability and safety for further retries.
647+
JsonNode originalNode = mapper.readTree(element.getOriginalPayload());
648+
JsonNode originalInnerEvent = originalNode.get("changeEvent");
649+
String originalDataStr = originalInnerEvent.get("data").asText();
579650
JsonNode originalDataNode = mapper.readTree(originalDataStr);
580-
assertEquals("Jane", originalDataNode.get("name").asText());
651+
assertEquals("John", originalDataNode.get("name").asText());
581652
} catch (Exception e) {
582653
throw new RuntimeException(e);
583654
}
@@ -975,12 +1046,6 @@ public void applyUdfToDataField_mixedEvents_correctlyProcessed() throws IOExcept
9751046
assertTrue(eventMap.containsKey("UPDATE"));
9761047
assertTrue(eventMap.containsKey("READ"));
9771048

978-
// Verify UDF processed flag
979-
assertTrue(eventMap.get("INSERT").has("_metadata_udf_processed"));
980-
assertTrue(eventMap.get("UPDATE").has("_metadata_udf_processed"));
981-
assertTrue(eventMap.get("READ").has("_metadata_udf_processed"));
982-
assertFalse(eventMap.get("DELETE").has("_metadata_udf_processed"));
983-
9841049
try {
9851050
// Verify UDF was applied to INSERT, UPDATE, READ
9861051
String insertDataStr = eventMap.get("INSERT").get("data").asText();

‎v2/datastream-mongodb-to-firestore/src/test/java/com/google/cloud/teleport/v2/templates/datastream/MongoDbChangeEventContextTest.java‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import com.fasterxml.jackson.core.JsonProcessingException;
3030
import com.fasterxml.jackson.databind.JsonNode;
3131
import com.fasterxml.jackson.databind.ObjectMapper;
32+
import com.fasterxml.jackson.databind.node.ObjectNode;
3233
import com.google.cloud.teleport.v2.transforms.Utils;
3334
import com.google.common.collect.ImmutableMap;
3435
import org.bson.Document;
@@ -572,4 +573,52 @@ public void testToString() throws JsonProcessingException {
572573

573574
assertEquals(jsonString, expectedJson);
574575
}
576+
577+
@Test
578+
public void testConstructorWithOriginalEvent() throws JsonProcessingException {
579+
// Test that the 3-argument constructor correctly stores separate instances for
580+
// the current change event and the original change event.
581+
JsonNode originalEvent = insertEvent.deepCopy();
582+
((ObjectNode) insertEvent).put("modified", true);
583+
584+
MongoDbChangeEventContext context =
585+
new MongoDbChangeEventContext(insertEvent, originalEvent, SHADOW_PREFIX);
586+
587+
// Verify that both current and original events are preserved as passed.
588+
assertEquals(insertEvent, context.getChangeEvent());
589+
assertEquals(originalEvent, context.getOriginalChangeEvent());
590+
}
591+
592+
@Test
593+
public void test2ArgumentConstructorSetsOriginalSameAsCurrent() throws JsonProcessingException {
594+
// Test that the 2-argument constructor falls back to setting the original event
595+
// to be the same as the current event when no separate original event is provided.
596+
MongoDbChangeEventContext context = new MongoDbChangeEventContext(insertEvent, SHADOW_PREFIX);
597+
598+
assertEquals(insertEvent, context.getChangeEvent());
599+
assertEquals(insertEvent, context.getOriginalChangeEvent());
600+
}
601+
602+
@Test
603+
public void testConstructorWithWrappedPayload() throws JsonProcessingException {
604+
// Test that the constructor correctly detects and unwraps events that are
605+
// wrapped in a DLQ structure (containing "changeEvent" and "dataCollection").
606+
ObjectMapper mapper = new ObjectMapper();
607+
JsonNode wrappedPayload =
608+
mapper.readTree(
609+
"{\"changeEvent\":{\"_id\":\"{\\\"$oid\\\": \\\"645c9a7e7b8b1a0e9c0f8b3a\\\"}\",\"op\":\"i\",\"data\":\"{\\\"name\\\":\\\"John\\\"}\",\"_metadata_source\":{\"collection\":\"test_collection\"},\"_metadata_timestamp_seconds\":1683782270,\"_metadata_timestamp_nanos\":123456789},\"dataCollection\":\"test\"}");
610+
611+
MongoDbChangeEventContext context =
612+
new MongoDbChangeEventContext(wrappedPayload, SHADOW_PREFIX);
613+
614+
// Verify that it successfully extracted and unwrapped the inner event.
615+
JsonNode extractedEvent = context.getChangeEvent();
616+
assertEquals("i", extractedEvent.get("op").asText());
617+
assertEquals(
618+
"test_collection", extractedEvent.get("_metadata_source").get("collection").asText());
619+
620+
// Verify that originalChangeEvent is also correctly extracted and unwrapped.
621+
JsonNode originalExtractedEvent = context.getOriginalChangeEvent();
622+
assertEquals("i", originalExtractedEvent.get("op").asText());
623+
}
575624
}

0 commit comments

Comments
 (0)