Add Transformation support for tuples

Signed-off-by: Andre Kurait <[email protected]>

TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/ResultsToLogsConsumer.java

@@ -14,6 +14,7 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 
 import org.opensearch.migrations.replay.datatypes.UniqueSourceRequestKey;
+import org.opensearch.migrations.transform.IJsonTransformer;
 
 import io.netty.buffer.ByteBuf;
 import lombok.Lombok;
@@ -27,26 +28,31 @@ public class ResultsToLogsConsumer implements BiConsumer<SourceTargetCaptureTupl
     public static final String TRANSACTION_SUMMARY_LOGGER = "TransactionSummaryLogger";
     private static final String MISSING_STR = "-";
     private static final ObjectMapper PLAIN_MAPPER = new ObjectMapper();
+    private static final IJsonTransformer NOOP_JSON_TRANSFORMER = new TransformationLoader().getTransformerFactoryLoader(
+        null, null, "NoopTransformerProvider");
 
     private final Logger tupleLogger;
     private final Logger progressLogger;
+    private final IJsonTransformer tupleTransformer;
+
     private final AtomicInteger tupleCounter;
 
     public ResultsToLogsConsumer() {
-        this(null, null);
+        this(null, null, null);
     }
 
-    public ResultsToLogsConsumer(Logger tupleLogger, Logger progressLogger) {
+    public ResultsToLogsConsumer(Logger tupleLogger, Logger progressLogger, IJsonTransformer tupleTransformer) {
         this.tupleLogger = tupleLogger != null ? tupleLogger : LoggerFactory.getLogger(OUTPUT_TUPLE_JSON_LOGGER);
         this.progressLogger = progressLogger != null ? progressLogger : makeTransactionSummaryLogger();
         tupleCounter = new AtomicInteger();
+        this.tupleTransformer = tupleTransformer != null ? tupleTransformer : NOOP_JSON_TRANSFORMER ;
     }
 
     // set this up so that the preamble prints out once, right after we have a logger
     // if it's configured to output at all
     private static Logger makeTransactionSummaryLogger() {
         var logger = LoggerFactory.getLogger(TRANSACTION_SUMMARY_LOGGER);
-        logger.atInfo().setMessage("{}").addArgument(() -> getTransactionSummaryStringPreamble()).log();
+        logger.atInfo().setMessage("{}").addArgument(ResultsToLogsConsumer::getTransactionSummaryStringPreamble).log();
         return logger;
     }
 
@@ -73,43 +79,48 @@ private Map<String, Object> toJSONObject(SourceTargetCaptureTuple tuple, ParsedH
     /**
      * Writes a tuple object to an output stream as a JSON object.
      * The JSON tuple is output on one line, and has several objects: "sourceRequest", "sourceResponse",
-     * "targetRequest", and "targetResponse". The "connectionId" is also included to aid in debugging.
+     * "targetRequest", and "targetResponses". The "connectionId", "numRequests", and "numErrors" are also included to aid in debugging.
      * An example of the format is below.
      * <p>
      * {
-     * "sourceRequest": {
-     * "Request-URI": XYZ,
-     * "Method": XYZ,
-     * "HTTP-Version": XYZ
-     * "body": XYZ,
-     * "header-1": XYZ,
-     * "header-2": XYZ
-     * },
-     * "targetRequest": {
-     * "Request-URI": XYZ,
-     * "Method": XYZ,
-     * "HTTP-Version": XYZ
-     * "body": XYZ,
-     * "header-1": XYZ,
-     * "header-2": XYZ
-     * },
-     * "sourceResponse": {
-     * "HTTP-Version": ABC,
-     * "Status-Code": ABC,
-     * "Reason-Phrase": ABC,
-     * "response_time_ms": 123,
-     * "body": ABC,
-     * "header-1": ABC
-     * },
-     * "targetResponse": {
-     * "HTTP-Version": ABC,
-     * "Status-Code": ABC,
-     * "Reason-Phrase": ABC,
-     * "response_time_ms": 123,
-     * "body": ABC,
-     * "header-2": ABC
-     * },
-     * "connectionId": "0242acfffe1d0008-0000000c-00000003-0745a19f7c3c5fc9-121001ff.0"
+     *   "sourceRequest": {
+     *     "Request-URI": "XYZ",
+     *     "Method": "XYZ",
+     *     "HTTP-Version": "XYZ",
+     *     "header-1": "XYZ",
+     *     "header-2": "XYZ",
+     *     "body": "BASE64ENCODEDSTRING"
+     *   },
+     *   "targetRequest": {
+     *     "Request-URI": "XYZ",
+     *     "Method": "XYZ",
+     *     "HTTP-Version": "XYZ",
+     *     "header-1": "XYZ",
+     *     "header-2": "XYZ",
+     *     "body": "BASE64ENCODEDSTRING"
+     *   },
+     *   "sourceResponse": {
+     *     "response_time_ms": 0,
+     *     "HTTP-Version": "XYZ",
+     *     "Status-Code": XYZ,
+     *     "Reason-Phrase": "XYZ",
+     *     "header-1": "XYZ",
+     *     "header-2": "XYZ",
+     *     "body": "BASE64ENCODEDSTRING"
+     *   },
+     *   "targetResponses": [{
+     *     "response_time_ms": 0,
+     *     "HTTP-Version": "XYZ",
+     *     "Status-Code": XYZ,
+     *     "Reason-Phrase": "XYZ",
+     *     "header-1": "XYZ",
+     *     "header-2": "XYZ",
+     *     "body": "BASE64ENCODEDSTRING"
+     *   }],
+     *   "connectionId": "XYZ",
+     *   "numRequests": XYZ,
+     *   "numErrors": XYZ
+     *   "error": "XYZ"
      * }
      *
      * @param tuple the RequestResponseResponseTriple object to be converted into json and written to the stream.
@@ -122,7 +133,9 @@ public void accept(SourceTargetCaptureTuple tuple, ParsedHttpMessagesAsDicts par
             .log();
         if (tupleLogger.isInfoEnabled()) {
             try {
-                var tupleString = PLAIN_MAPPER.writeValueAsString(toJSONObject(tuple, parsedMessages));
+                var originalTuple = toJSONObject(tuple, parsedMessages);
+                var transformedTuple = tupleTransformer.transformJson(originalTuple);
+                var tupleString = PLAIN_MAPPER.writeValueAsString(transformedTuple);
                 tupleLogger.atInfo().setMessage("{}").addArgument(() -> tupleString).log();
             } catch (Exception e) {
                 log.atError().setMessage("Exception converting tuple to string").setCause(e).log();

TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/TrafficReplayer.java

@@ -34,7 +34,9 @@
 import com.beust.jcommander.JCommander;
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.ParameterException;
+import com.beust.jcommander.ParametersDelegate;
 import io.netty.util.concurrent.DefaultThreadFactory;
+import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
 import org.slf4j.LoggerFactory;
 import org.slf4j.event.Level;
@@ -136,31 +138,12 @@ public static class Parameters {
                 + "(cannot be used with other auth arguments)")
         String useSigV4ServiceAndRegion;
 
-        @Parameter(
-            required = false,
-            names = "--transformer-config-base64",
-            arity = 1,
-            description = "Configuration of message transformers.  The same contents as --transformer-config but " +
-                "Base64 encoded so that the configuration is easier to pass as a command line parameter.")
-        String transformerConfigEncoded;
+        @ParametersDelegate
+        private RequestTransformationParams requestTransformationParams = new RequestTransformationParams();
 
-        @Parameter(
-            required = false,
-            names = "--transformer-config",
-            arity = 1,
-            description = "Configuration of message transformers.  Either as a string that identifies the "
-            + "transformer that should be run (with default settings) or as json to specify options "
-            + "as well as multiple transformers to run in sequence.  "
-            + "For json, keys are the (simple) names of the loaded transformers and values are the "
-            + "configuration passed to each of the transformers.")
-        String transformerConfig;
+        @ParametersDelegate
+        private TupleTransformationParams tupleTransformationParams = new TupleTransformationParams();
 
-        @Parameter(
-            required = false,
-            names = "--transformer-config-file",
-            arity = 1,
-            description = "Path to the JSON configuration file of message transformers.")
-        String transformerConfigFile;
         @Parameter(
             required = false,
             names = "--user-agent",
@@ -257,6 +240,68 @@ public static class Parameters {
         String otelCollectorEndpoint;
     }
 
+    public interface TransformerParams {
+        String getTransformerConfigEncoded();
+        String getTransformerConfig();
+        String getTransformerConfigFile();
+    }
+
+    @Getter
+    public static class RequestTransformationParams implements TransformerParams {
+
+        @Parameter(
+            required = false,
+            names = "--transformer-config-encoded",
+            arity = 1,
+            description = "Configuration of message transformers.  The same contents as --transformer-config but " +
+                "Base64 encoded so that the configuration is easier to pass as a command line parameter.")
+        private String transformerConfigEncoded;
+
+        @Parameter(
+            required = false,
+            names = "--transformer-config",
+            arity = 1,
+            description = "Configuration of message transformers.  Either as a string that identifies the "
+                + "transformer that should be run (with default settings) or as json to specify options "
+                + "as well as multiple transformers to run in sequence.  "
+                + "For json, keys are the (simple) names of the loaded transformers and values are the "
+                + "configuration passed to each of the transformers.")
+        private String transformerConfig;
+
+        @Parameter(
+            required = false,
+            names = "--transformer-config-file",
+            arity = 1,
+            description = "Path to the JSON configuration file of message transformers.")
+        private String transformerConfigFile;
+    }
+
+    @Getter
+    public static class TupleTransformationParams implements TransformerParams {
+
+        @Parameter(
+            required = false,
+            names = "--tuple-transformer-config-base64",
+            arity = 1,
+            description = "Base64 encoded transformer configuration for tuple transformation.")
+        private String transformerConfigEncoded;
+
+        @Parameter(
+            required = false,
+            names = "--tuple-transformer-config",
+            arity = 1,
+            description = "Transformer configuration string for tuple transformation.")
+        private String transformerConfig;
+
+        @Parameter(
+            required = false,
+            names = "--tuple-transformer-config-file",
+            arity = 1,
+            description = "Path to the JSON configuration file for tuple transformation.")
+        private String transformerConfigFile;
+    }
+
+
     private static Parameters parseArgs(String[] args) {
         Parameters p = new Parameters();
         JCommander jCommander = new JCommander(p);
@@ -276,31 +321,31 @@ private static int isConfigured(String s) {
         return (s == null || s.isBlank()) ? 0 : 1;
     }
 
-    private static String getTransformerConfig(Parameters params) {
-        var configuredCount = isConfigured(params.transformerConfigFile) +
-                isConfigured(params.transformerConfigEncoded) +
-                isConfigured(params.transformerConfig);
+    private static String getTransformerConfig(TransformerParams params) {
+        var configuredCount = isConfigured(params.getTransformerConfigFile()) +
+                isConfigured(params.getTransformerConfigEncoded()) +
+                isConfigured(params.getTransformerConfig());
         if (configuredCount > 1) {
             System.err.println("Specify only one of --transformer-config-base64, --transformer-config or " +
                 "--transformer-config-file.");
             System.exit(4);
         }
 
-        if (params.transformerConfigFile != null && !params.transformerConfigFile.isBlank()) {
+        if (params.getTransformerConfigFile() != null && !params.getTransformerConfigFile().isBlank()) {
             try {
-                return Files.readString(Paths.get(params.transformerConfigFile), StandardCharsets.UTF_8);
+                return Files.readString(Paths.get(params.getTransformerConfigFile()), StandardCharsets.UTF_8);
             } catch (IOException e) {
                 System.err.println("Error reading transformer configuration file: " + e.getMessage());
                 System.exit(5);
             }
         }
 
-        if (params.transformerConfig != null && !params.transformerConfig.isBlank()) {
-            return params.transformerConfig;
+        if (params.getTransformerConfig() != null && !params.getTransformerConfig().isBlank()) {
+            return params.getTransformerConfig();
         }
 
-        if (params.transformerConfigEncoded != null && !params.transformerConfigEncoded.isBlank()) {
-            return new String(Base64.getDecoder().decode(params.transformerConfigEncoded));
+        if (params.getTransformerConfigEncoded() != null && !params.getTransformerConfigEncoded().isBlank()) {
+            return new String(Base64.getDecoder().decode(params.getTransformerConfigEncoded()));
         }
 
         return null;
@@ -365,18 +410,26 @@ public static void main(String[] args) throws Exception {
             var timeShifter = new TimeShifter(params.speedupFactor);
             var serverTimeout = Duration.ofSeconds(params.targetServerResponseTimeoutSeconds);
 
-            String transformerConfig = getTransformerConfig(params);
-            if (transformerConfig != null) {
-                log.atInfo().setMessage(() -> "Transformations config string: " + transformerConfig).log();
+            String requestTransformerConfig = getTransformerConfig(params.requestTransformationParams);
+            if (requestTransformerConfig != null) {
+                log.atInfo().setMessage("Request Transformations config string: {}")
+                    .addArgument(requestTransformerConfig).log();
             }
+
+            String tupleTransformerConfig = getTransformerConfig(params.tupleTransformationParams);
+            if (requestTransformerConfig != null) {
+                log.atInfo().setMessage("Tuple Transformations config string: {}")
+                    .addArgument(tupleTransformerConfig).log();
+            }
+
             final var orderedRequestTracker = new OrderedWorkerTracker<Void>();
             final var hostname = uri.getHost();
 
             var tr = new TrafficReplayerTopLevel(
                 topContext,
                 uri,
                 authTransformer,
-                new TransformationLoader().getTransformerFactoryLoader(hostname, params.userAgent, transformerConfig),
+                new TransformationLoader().getTransformerFactoryLoader(hostname, params.userAgent, requestTransformerConfig),
                 TrafficReplayerTopLevel.makeNettyPacketConsumerConnectionPool(
                     uri,
                     params.allowInsecureConnections,
@@ -404,7 +457,8 @@ public static void main(String[] args) throws Exception {
             }, ACTIVE_WORK_MONITOR_CADENCE_MS, ACTIVE_WORK_MONITOR_CADENCE_MS, TimeUnit.MILLISECONDS);
 
             setupShutdownHookForReplayer(tr);
-            var tupleWriter = new TupleParserChainConsumer(new ResultsToLogsConsumer());
+            var tupleWriter = new TupleParserChainConsumer(new ResultsToLogsConsumer(null, null,
+                new TransformationLoader().getTransformerFactoryLoader(hostname, params.userAgent, requestTransformerConfig)));
             tr.setupRunAndWaitForReplayWithShutdownChecks(
                 Duration.ofSeconds(params.observedPacketConnectionTimeout),
                 serverTimeout,