Add option to unwrap messages replayed from DLQ to avoid nesting loops. (#2888)

* Draft

* Add option to unwrap replayed messages.

* Comments.

Author: claudevdm

v1/src/main/java/com/google/cloud/teleport/templates/PubSubToSplunk.java

@@ -18,6 +18,7 @@
 import com.google.cloud.teleport.coders.FailsafeElementCoder;
 import com.google.cloud.teleport.metadata.Template;
 import com.google.cloud.teleport.metadata.TemplateCategory;
+import com.google.cloud.teleport.metadata.TemplateParameter;
 import com.google.cloud.teleport.splunk.SplunkEvent;
 import com.google.cloud.teleport.splunk.SplunkEventCoder;
 import com.google.cloud.teleport.splunk.SplunkIO;
@@ -60,6 +61,8 @@
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.json.JSONException;
+import org.json.JSONObject;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -254,18 +257,38 @@ public static PipelineResult run(PubSubToSplunkOptions options) {
                     .withEnableGzipHttpCompression(options.getEnableGzipHttpCompression())
                     .build());
 
+    final ValueProvider<Boolean> unwrapHecProvider = options.getUnwrapHecForDeadletter();
+
     // 5a) Wrap write failures into a FailsafeElement.
     PCollection<FailsafeElement<String, String>> wrappedSplunkWriteErrors =
         writeErrors.apply(
             "WrapSplunkWriteErrors",
             ParDo.of(
                 new DoFn<SplunkWriteError, FailsafeElement<String, String>>() {
 
+                  private boolean unwrapHecForDeadletter = false;
+
+                  @Setup
+                  public void setup() {
+                    if (unwrapHecProvider != null) {
+                      unwrapHecForDeadletter =
+                          MoreObjects.firstNonNull(unwrapHecProvider.get(), false);
+                    }
+                    LOG.info("unwrapHecForDeadletter set to: {}", unwrapHecForDeadletter);
+                  }
+
                   @ProcessElement
                   public void processElement(ProcessContext context) {
                     SplunkWriteError error = context.element();
+
+                    // Extract original payload if unwrap is enabled
+                    String payload = error.payload();
+                    if (unwrapHecForDeadletter) {
+                      payload = extractOriginalPayloadFromHec(payload);
+                    }
+
                     FailsafeElement<String, String> failsafeElement =
-                        FailsafeElement.of(error.payload(), error.payload());
+                        FailsafeElement.of(payload, payload);
 
                     if (error.statusMessage() != null) {
                       failsafeElement.setErrorMessage(error.statusMessage());
@@ -304,7 +327,20 @@ public interface PubSubToSplunkOptions
       extends SplunkOptions,
           PubsubReadSubscriptionOptions,
           PubsubWriteDeadletterTopicOptions,
-          JavascriptTextTransformerOptions {}
+          JavascriptTextTransformerOptions {
+
+    @TemplateParameter.Boolean(
+        optional = true,
+        description = "Unwrap Splunk HEC format from deadletter messages.",
+        helpText =
+            "When enabled, if a message fails to write to Splunk and is sent to the deadletter "
+                + "queue, the original payload will be extracted from the Splunk HEC format before "
+                + "writing to the deadletter topic. This prevents event nesting when messages are "
+                + "replayed. Default: `false`.")
+    ValueProvider<Boolean> getUnwrapHecForDeadletter();
+
+    void setUnwrapHecForDeadletter(ValueProvider<Boolean> unwrapHecForDeadletter);
+  }
 
   /**
    * A {@link PTransform} that reads messages from a Pub/Sub subscription, increments a counter and
@@ -411,4 +447,43 @@ private static JsonObject getAttributesJson(Map<String, String> attributesMap) {
 
     return attributesJson;
   }
+
+  /**
+   * Extracts the original payload from a Splunk HEC-formatted JSON string.
+   *
+   * <p>When a SplunkEvent fails to write to Splunk, it's stored in HEC format like: {"event":
+   * "original payload", "time": ..., "host": ..., etc.}
+   *
+   * <p>This method extracts just the "event" field to prevent nesting when replaying messages from
+   * the deadletter queue.
+   *
+   * @param hecPayload The HEC-formatted JSON string
+   * @return The original event payload, or the input if extraction fails
+   */
+  @VisibleForTesting
+  protected static String extractOriginalPayloadFromHec(String hecPayload) {
+    try {
+      JSONObject json = new JSONObject(hecPayload);
+
+      // Check if this looks like HEC format (has "event" field)
+      if (json.has("event")) {
+        Object eventObj = json.get("event");
+
+        // The event field could be a JSON object/array or a string
+        if (eventObj instanceof String) {
+          return (String) eventObj;
+        } else {
+          // If it's a JSONObject or JSONArray, return its string
+          // representation
+          return eventObj.toString();
+        }
+      }
+    } catch (JSONException e) {
+      // If parsing fails, log and return the original
+      LOG.debug("Failed to parse payload as HEC JSON, returning as-is. Error: {}", e.getMessage());
+    }
+
+    // If no "event" field found or parsing failed, return original
+    return hecPayload;
+  }
 }

v1/src/test/java/com/google/cloud/teleport/templates/PubsubToSplunkTest.java

@@ -16,6 +16,7 @@
 package com.google.cloud.teleport.templates;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertEquals;
 
 import com.google.gson.Gson;
 import com.google.gson.JsonObject;
@@ -88,4 +89,74 @@ public void testPubsubMessageWithComplexAttributes() {
       assertThat(attributesJson.get(key).getAsString()).isEqualTo(attributes.get(key));
     }
   }
+
+  /** Tests extraction of original payload from HEC-formatted JSON. */
+  @Test
+  public void testExtractOriginalPayloadFromHec_withStringEvent() {
+    String hecPayload = "{\"event\":\"simple text message\",\"time\":1634567890}";
+
+    String extracted = PubSubToSplunk.extractOriginalPayloadFromHec(hecPayload);
+
+    assertEquals("simple text message", extracted);
+  }
+
+  /** Tests extraction when event is a JSON object. */
+  @Test
+  public void testExtractOriginalPayloadFromHec_withJsonEvent() {
+    String hecPayload =
+        "{\"event\":{\"userId\":123,\"action\":\"login\"},\"time\":1634567890,\"host\":\"web-01\"}";
+
+    String extracted = PubSubToSplunk.extractOriginalPayloadFromHec(hecPayload);
+
+    // The extracted event should be the JSON object as a string
+    JsonObject extractedJson = GSON.fromJson(extracted, JsonObject.class);
+    assertThat(extractedJson.get("userId").getAsInt()).isEqualTo(123);
+    assertThat(extractedJson.get("action").getAsString()).isEqualTo("login");
+  }
+
+  /** Tests extraction when payload is already unwrapped (not HEC format). */
+  @Test
+  public void testExtractOriginalPayloadFromHec_alreadyUnwrapped() {
+    String originalPayload = "{\"userId\":123,\"action\":\"login\"}";
+
+    String extracted = PubSubToSplunk.extractOriginalPayloadFromHec(originalPayload);
+
+    // Should return as-is since no "event" field
+    assertEquals(originalPayload, extracted);
+  }
+
+  /** Tests extraction with plain text (not JSON). */
+  @Test
+  public void testExtractOriginalPayloadFromHec_plainText() {
+    String plainText = "This is just plain text, not JSON";
+
+    String extracted = PubSubToSplunk.extractOriginalPayloadFromHec(plainText);
+
+    // Should return as-is when parsing fails
+    assertEquals(plainText, extracted);
+  }
+
+  /** Tests extraction with complex HEC format including all metadata fields. */
+  @Test
+  public void testExtractOriginalPayloadFromHec_fullHecFormat() {
+    String fullHecPayload =
+        "{\"event\":{\"log\":\"Application started\",\"level\":\"INFO\"},"
+            + "\"time\":1634567890,"
+            + "\"host\":\"app-server-01\","
+            + "\"source\":\"application\","
+            + "\"sourcetype\":\"json\","
+            + "\"index\":\"main\","
+            + "\"fields\":{\"environment\":\"production\"}}";
+
+    String extracted = PubSubToSplunk.extractOriginalPayloadFromHec(fullHecPayload);
+
+    // Should extract just the event field
+    JsonObject extractedJson = GSON.fromJson(extracted, JsonObject.class);
+    assertThat(extractedJson.get("log").getAsString()).isEqualTo("Application started");
+    assertThat(extractedJson.get("level").getAsString()).isEqualTo("INFO");
+
+    // Verify other HEC fields are NOT in the extracted payload
+    assertThat(extractedJson.has("time")).isFalse();
+    assertThat(extractedJson.has("host")).isFalse();
+  }
 }