org.opensearch.migrations.replay.TransformationLoader Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of trafficReplayer Show documentation
Everything opensearch migrations
There is a newer version: 0.2.0.4
package org.opensearch.migrations.replay;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.ServiceLoader;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;

import org.opensearch.migrations.transform.IJsonTransformer;
import org.opensearch.migrations.transform.IJsonTransformerProvider;
import org.opensearch.migrations.transform.JsonCompositeTransformer;
import org.opensearch.migrations.transform.JsonKeysForHttpMessage;

import lombok.AllArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;

@Slf4j
public class TransformationLoader {
    public static final String WRONG_JSON_STRUCTURE_MESSAGE =
        "Must specify the top-level configuration list with a sequence of "
            + "maps that have only one key each, where the key is the name of the transformer to be configured.";
    public static final Pattern CLASS_NAME_PATTERN = Pattern.compile("^[^{}]*$");
    private final List providers;
    ObjectMapper objMapper = new ObjectMapper();

    public TransformationLoader() {
        ServiceLoader transformerProviders = ServiceLoader.load(
            IJsonTransformerProvider.class
        );
        var inProgressProviders = new ArrayList();
        for (var provider : transformerProviders) {
            log.info("Adding IJsonTransfomerProvider: " + provider);
            inProgressProviders.add(provider);
        }
        providers = Collections.unmodifiableList(inProgressProviders);
        log.atInfo()
            .setMessage(
                () -> "IJsonTransformerProviders loaded: "
                    + providers.stream().map(p -> p.getClass().toString()).collect(Collectors.joining("; "))
            )
            .log();
    }

    List> parseFullConfig(String fullConfig) throws JsonProcessingException {
        if (CLASS_NAME_PATTERN.matcher(fullConfig).matches()) {
            return List.of(Collections.singletonMap(fullConfig, null));
        } else {
            return objMapper.readValue(fullConfig, new TypeReference<>() {
            });
        }
    }

    protected Stream getTransformerFactoryFromServiceLoader(String fullConfig)
        throws JsonProcessingException {
        var configList = fullConfig == null ? List.of() : parseFullConfig(fullConfig);
        if (configList.isEmpty() || providers.isEmpty()) {
            log.warn("No transformer configuration specified.  No custom transformations will be performed");
            return Stream.of();
        } else {
            return configList.stream().map(c -> configureTransformerFromConfig((Map) c));
        }
    }

    @SneakyThrows // JsonProcessingException should be impossible since the contents are those that were just parsed
    private IJsonTransformer configureTransformerFromConfig(Map c) {
        var keys = c.keySet();
        if (keys.size() != 1) {
            throw new IllegalArgumentException(WRONG_JSON_STRUCTURE_MESSAGE);
        }
        var key = keys.stream()
            .findFirst()
            .orElseThrow(() -> new IllegalArgumentException(WRONG_JSON_STRUCTURE_MESSAGE));
        for (var p : providers) {
            var transformerName = p.getName();
            if (transformerName.equals(key)) {
                var configuration = c.get(key);
                log.atInfo()
                    .setMessage(
                        () -> "Creating a transformer through provider=" + p + " with configuration=" + configuration
                    )
                    .log();
                return p.createTransformer(configuration);
            }
        }
        throw new IllegalArgumentException("Could not find a provider named: " + key);
    }

    public IJsonTransformer getTransformerFactoryLoader(String newHostName) {
        return getTransformerFactoryLoader(newHostName, null, null);
    }

    public IJsonTransformer getTransformerFactoryLoader(String newHostName, String userAgent, String fullConfig) {
        try {
            var loadedTransformers = getTransformerFactoryFromServiceLoader(fullConfig);
            return new JsonCompositeTransformer(
                Stream.concat(
                    loadedTransformers,
                    Stream.concat(
                        Optional.ofNullable(userAgent).stream().map(UserAgentTransformer::new),
                        Optional.ofNullable(newHostName).stream().map(HostTransformer::new)
                    )
                ).toArray(IJsonTransformer[]::new)
            );
        } catch (JsonProcessingException e) {
            throw new IllegalArgumentException("Could not parse the transformer configuration as a json list", e);
        }
    }

    @AllArgsConstructor
    private static class UserAgentTransformer implements IJsonTransformer {
        public static final String USER_AGENT = "user-agent";
        private final String userAgent;

        @Override
        public Map transformJson(Map incomingJson) {
            var headers = (Map) incomingJson.get(JsonKeysForHttpMessage.HEADERS_KEY);
            var oldVal = headers.get(USER_AGENT);
            if (oldVal != null) {
                if (oldVal instanceof List) {
                    // see https://www.rfc-editor.org/rfc/rfc9110.html#name-field-lines-and-combined-fi
                    oldVal = String.join(", ", (List) oldVal);
                }
                headers.replace(USER_AGENT, oldVal + "; " + userAgent);
            } else {
                headers.put(USER_AGENT, userAgent);
            }
            return incomingJson;
        }
    }

    @AllArgsConstructor
    private static class HostTransformer implements IJsonTransformer {
        private final String newHostName;

        @Override
        public Map transformJson(Map incomingJson) {
            var headers = (Map) incomingJson.get(JsonKeysForHttpMessage.HEADERS_KEY);
            headers.replace("host", newHostName);
            return incomingJson;
        }
    }
}