diff --git a/.github/workflows/kafka-plugin-integration-tests.yml b/.github/workflows/kafka-plugin-integration-tests.yml index 28fb5bdc05..507ca1d435 100644 --- a/.github/workflows/kafka-plugin-integration-tests.yml +++ b/.github/workflows/kafka-plugin-integration-tests.yml @@ -8,12 +8,16 @@ on: paths: - 'data-prepper-plugins/kafka-plugins/**' - '*gradle*' - pull_request: + pull_request_target: + types: [ opened, synchronize, reopened ] paths: - 'data-prepper-plugins/kafka-plugins/**' - '*gradle*' workflow_dispatch: +permissions: + id-token: write + contents: read jobs: integration-tests: @@ -41,9 +45,28 @@ jobs: - name: Wait for Kafka run: | ./gradlew data-prepper-plugins:kafka-plugins:integrationTest -Dtests.kafka.bootstrap_servers=localhost:9092 -Dtests.kafka.authconfig.username=admin -Dtests.kafka.authconfig.password=admin --tests KafkaStartIT + + - name: Configure AWS credentials + id: aws-credentials + uses: aws-actions/configure-aws-credentials@v4 + with: + role-to-assume: ${{ secrets.TEST_IAM_ROLE_ARN }} + aws-region: ${{ secrets.TEST_REGION }} + output-credentials: true + - name: Configure AWS default credentials + run: | + aws configure set default.region ${{ secrets.TEST_REGION }} + aws configure set default.aws_access_key_id ${{ steps.aws-credentials.outputs.aws-access-key-id }} + aws configure set default.aws_secret_access_key ${{ steps.aws-credentials.outputs.aws-secret-access-key }} + aws configure set default.aws_session_token ${{ steps.aws-credentials.outputs.aws-session-token }} + - name: Run Kafka integration tests run: | - ./gradlew data-prepper-plugins:kafka-plugins:integrationTest -Dtests.kafka.bootstrap_servers=localhost:9092 -Dtests.kafka.authconfig.username=admin -Dtests.kafka.authconfig.password=admin --tests KafkaSourceJsonTypeIT --tests KafkaBufferIT --tests KafkaBufferOTelIT + ./gradlew data-prepper-plugins:kafka-plugins:integrationTest \ + -Dtests.kafka.bootstrap_servers=localhost:9092 \ + -Dtests.kafka.authconfig.username=admin -Dtests.kafka.authconfig.password=admin \ + -Dtests.kafka.kms_key=alias/DataPrepperTesting \ + --tests '*kafka.buffer*' --tests KafkaSourceJsonTypeIT --tests KafkaBufferOTelIT - name: Upload Unit Test Results if: always() diff --git a/build.gradle b/build.gradle index dfbb10d3e4..cc6766680e 100644 --- a/build.gradle +++ b/build.gradle @@ -5,21 +5,6 @@ import com.github.jk1.license.render.TextReportRenderer buildscript { - repositories { - mavenCentral() { - metadataSources { - mavenPom() - ignoreGradleMetadataRedirection() - } - } - maven { - url 'https://plugins.gradle.org/m2/' - metadataSources { - mavenPom() - ignoreGradleMetadataRedirection() - } - } - } dependencies { classpath 'com.github.jk1:gradle-license-report:2.1' } @@ -41,14 +26,6 @@ allprojects { ext { mavenPublicationRootFile = file("${rootProject.buildDir}/m2") } - - repositories { - mavenCentral() - maven { url 'https://jitpack.io' } - maven { - url 'https://packages.confluent.io/maven/' - } - } spotless { format 'markdown', { @@ -88,7 +65,7 @@ subprojects { } } dependencies { - implementation platform('com.fasterxml.jackson:jackson-bom:2.15.3') + implementation platform('com.fasterxml.jackson:jackson-bom:2.16.1') implementation platform('org.eclipse.jetty:jetty-bom:9.4.53.v20231009') implementation platform('io.micrometer:micrometer-bom:1.10.5') implementation libs.guava.core @@ -194,6 +171,12 @@ subprojects { } because 'CVE from transitive dependencies' } + implementation('com.jayway.jsonpath:json-path') { + version { + require '2.9.0' + } + because 'Fixes CVE-2023-51074 from transitive dependencies' + } implementation('org.bitbucket.b_c:jose4j') { version { require '0.9.3' @@ -218,6 +201,12 @@ subprojects { } because 'CVE-2023-5072, CVE from transitive dependencies' } + implementation('org.apache.commons:commons-compress') { + version { + require '1.26.0' + } + because 'CVE-2024-25710, CVE-2024-26308' + } } } diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/AbstractBuffer.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/AbstractBuffer.java index f8ad3479a8..09d6d6004c 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/AbstractBuffer.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/AbstractBuffer.java @@ -145,12 +145,14 @@ public Map.Entry, CheckpointState> read(int timeoutInMillis) { @Override public void checkpoint(final CheckpointState checkpointState) { checkpointTimer.record(() -> doCheckpoint(checkpointState)); - final int numRecordsToBeChecked = checkpointState.getNumRecordsToBeChecked(); - recordsInFlight.addAndGet(-numRecordsToBeChecked); - recordsProcessedCounter.increment(numRecordsToBeChecked); + if (!isByteBuffer()) { + final int numRecordsToBeChecked = checkpointState.getNumRecordsToBeChecked(); + recordsInFlight.addAndGet(-numRecordsToBeChecked); + recordsProcessedCounter.increment(numRecordsToBeChecked); + } } - protected int getRecordsInFlight() { + public int getRecordsInFlight() { return recordsInFlight.intValue(); } diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/codec/ByteDecoder.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/codec/ByteDecoder.java index 46420ca7cc..eb43f72489 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/codec/ByteDecoder.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/codec/ByteDecoder.java @@ -11,6 +11,7 @@ import java.io.InputStream; import java.io.Serializable; import java.util.function.Consumer; +import java.time.Instant; public interface ByteDecoder extends Serializable { /** @@ -18,9 +19,10 @@ public interface ByteDecoder extends Serializable { * {@link Record} loaded from the {@link InputStream}. * * @param inputStream The input stream for code to process + * @param timeReceived The time received value to be populated in the Record * @param eventConsumer The consumer which handles each event from the stream * @throws IOException throws IOException when invalid input is received or incorrect codec name is provided */ - void parse(InputStream inputStream, Consumer> eventConsumer) throws IOException; + void parse(InputStream inputStream, Instant timeReceived, Consumer> eventConsumer) throws IOException; } diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/codec/JsonDecoder.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/codec/JsonDecoder.java index 1aba7e56ee..f0793aa65f 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/codec/JsonDecoder.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/codec/JsonDecoder.java @@ -16,6 +16,7 @@ import org.opensearch.dataprepper.model.record.Record; import java.io.IOException; import java.io.InputStream; +import java.time.Instant; import java.util.Map; import java.util.Objects; import java.util.function.Consumer; @@ -24,7 +25,7 @@ public class JsonDecoder implements ByteDecoder { private final ObjectMapper objectMapper = new ObjectMapper(); private final JsonFactory jsonFactory = new JsonFactory(); - public void parse(InputStream inputStream, Consumer> eventConsumer) throws IOException { + public void parse(InputStream inputStream, Instant timeReceived, Consumer> eventConsumer) throws IOException { Objects.requireNonNull(inputStream); Objects.requireNonNull(eventConsumer); @@ -32,25 +33,28 @@ public void parse(InputStream inputStream, Consumer> eventConsumer while (!jsonParser.isClosed() && jsonParser.nextToken() != JsonToken.END_OBJECT) { if (jsonParser.getCurrentToken() == JsonToken.START_ARRAY) { - parseRecordsArray(jsonParser, eventConsumer); + parseRecordsArray(jsonParser, timeReceived, eventConsumer); } } } - private void parseRecordsArray(final JsonParser jsonParser, final Consumer> eventConsumer) throws IOException { + private void parseRecordsArray(final JsonParser jsonParser, final Instant timeReceived, final Consumer> eventConsumer) throws IOException { while (jsonParser.nextToken() != JsonToken.END_ARRAY) { final Map innerJson = objectMapper.readValue(jsonParser, Map.class); - final Record record = createRecord(innerJson); + final Record record = createRecord(innerJson, timeReceived); eventConsumer.accept(record); } } - private Record createRecord(final Map json) { - final JacksonEvent event = (JacksonEvent)JacksonLog.builder() + private Record createRecord(final Map json, final Instant timeReceived) { + final JacksonLog.Builder logBuilder = JacksonLog.builder() .withData(json) - .getThis() - .build(); + .getThis(); + if (timeReceived != null) { + logBuilder.withTimeReceived(timeReceived); + } + final JacksonEvent event = (JacksonEvent)logBuilder.build(); return new Record<>(event); } diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/Event.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/Event.java index a27906b2cc..04139a1669 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/Event.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/Event.java @@ -65,6 +65,12 @@ public interface Event extends Serializable { */ void delete(String key); + /** + * Delete all keys from the Event + * @since 2.8 + */ + void clear(); + /** * Generates a serialized Json string of the entire Event * diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/JacksonEvent.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/JacksonEvent.java index cd604a1b9d..ff60157b7e 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/JacksonEvent.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/JacksonEvent.java @@ -29,6 +29,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.HashMap; +import java.util.Iterator; import java.util.LinkedList; import java.util.List; import java.util.Map; @@ -136,10 +137,11 @@ public JsonNode getJsonNode() { */ @Override public void put(final String key, final Object value) { + checkArgument(!key.isEmpty(), "key cannot be an empty string for put method"); final String trimmedKey = checkAndTrimKey(key); - final LinkedList keys = new LinkedList<>(Arrays.asList(trimmedKey.split(SEPARATOR))); + final LinkedList keys = new LinkedList<>(Arrays.asList(trimmedKey.split(SEPARATOR, -1))); JsonNode parentNode = jsonNode; @@ -247,7 +249,12 @@ private List mapNodeToList(final String key, final JsonNode node, final C } private JsonPointer toJsonPointer(final String key) { - String jsonPointerExpression = SEPARATOR + key; + final String jsonPointerExpression; + if (key.isEmpty() || key.startsWith("/")) { + jsonPointerExpression = key; + } else { + jsonPointerExpression = SEPARATOR + key; + } return JsonPointer.compile(jsonPointerExpression); } @@ -259,6 +266,7 @@ private JsonPointer toJsonPointer(final String key) { @Override public void delete(final String key) { + checkArgument(!key.isEmpty(), "key cannot be an empty string for delete method"); final String trimmedKey = checkAndTrimKey(key); final int index = trimmedKey.lastIndexOf(SEPARATOR); @@ -276,6 +284,16 @@ public void delete(final String key) { } } + @Override + public void clear() { + // Delete all entries from the event + Iterator iter = toMap().keySet().iterator(); + JsonNode baseNode = jsonNode; + while (iter.hasNext()) { + ((ObjectNode) baseNode).remove((String)iter.next()); + } + } + @Override public String toJsonString() { return jsonNode.toString(); @@ -399,24 +417,31 @@ public static boolean isValidEventKey(final String key) { } private String checkAndTrimKey(final String key) { checkKey(key); - return trimKey(key); + return trimTrailingSlashInKey(key); } private static void checkKey(final String key) { checkNotNull(key, "key cannot be null"); - checkArgument(!key.isEmpty(), "key cannot be an empty string"); + if (key.isEmpty()) { + // Empty string key is valid + return; + } if (key.length() > MAX_KEY_LENGTH) { throw new IllegalArgumentException("key cannot be longer than " + MAX_KEY_LENGTH + " characters"); } if (!isValidKey(key)) { - throw new IllegalArgumentException("key " + key + " must contain only alphanumeric chars with .-_ and must follow JsonPointer (ie. 'field/to/key')"); + throw new IllegalArgumentException("key " + key + " must contain only alphanumeric chars with .-_@/ and must follow JsonPointer (ie. 'field/to/key')"); } } private String trimKey(final String key) { final String trimmedLeadingSlash = key.startsWith(SEPARATOR) ? key.substring(1) : key; - return trimmedLeadingSlash.endsWith(SEPARATOR) ? trimmedLeadingSlash.substring(0, trimmedLeadingSlash.length() - 2) : trimmedLeadingSlash; + return trimTrailingSlashInKey(trimmedLeadingSlash); + } + + private String trimTrailingSlashInKey(final String key) { + return key.length() > 1 && key.endsWith(SEPARATOR) ? key.substring(0, key.length() - 1) : key; } private static boolean isValidKey(final String key) { @@ -430,7 +455,9 @@ private static boolean isValidKey(final String key) { || c == '-' || c == '_' || c == '@' - || c == '/')) { + || c == '/' + || c == '[' + || c == ']')) { return false; } diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/log/JacksonOtelLog.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/log/JacksonOtelLog.java index bb0c02be6d..b18013f637 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/log/JacksonOtelLog.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/log/JacksonOtelLog.java @@ -10,6 +10,7 @@ import org.opensearch.dataprepper.model.event.EventType; import org.opensearch.dataprepper.model.event.JacksonEvent; +import java.time.Instant; import java.util.HashMap; import java.util.Iterator; import java.util.Map; @@ -161,6 +162,18 @@ public Builder withAttributes(final Map attributes) { return getThis(); } + /** + * Sets the time received for populating event origination time in event handle + * + * @param timeReceived time received + * @return the builder + * @since 2.7 + */ + @Override + public Builder withTimeReceived(final Instant timeReceived) { + return (Builder)super.withTimeReceived(timeReceived); + } + /** * Sets the observed time of the log event * diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonExponentialHistogram.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonExponentialHistogram.java index b865ce0eb5..a1ba387ee2 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonExponentialHistogram.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonExponentialHistogram.java @@ -7,6 +7,7 @@ import org.opensearch.dataprepper.model.event.EventType; +import java.time.Instant; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -237,6 +238,18 @@ public JacksonExponentialHistogram.Builder withPositiveOffset(int offset) { return this; } + /** + * Sets the time received for populating event origination time in event handle + * + * @param timeReceived time received + * @return the builder + * @since 2.7 + */ + @Override + public JacksonExponentialHistogram.Builder withTimeReceived(final Instant timeReceived) { + return (JacksonExponentialHistogram.Builder)super.withTimeReceived(timeReceived); + } + /** * Sets the offset for the negative buckets * diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonGauge.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonGauge.java index 4df5bf4793..33c633d951 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonGauge.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonGauge.java @@ -7,6 +7,7 @@ import org.opensearch.dataprepper.model.event.EventType; +import java.time.Instant; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -67,6 +68,18 @@ public Builder withValue(final Double value) { return this; } + /** + * Sets the time received for populating event origination time in event handle + * + * @param timeReceived time received + * @return the builder + * @since 2.7 + */ + @Override + public Builder withTimeReceived(final Instant timeReceived) { + return (Builder)super.withTimeReceived(timeReceived); + } + /** * Returns a newly created {@link JacksonGauge} * @return a JacksonGauge diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonHistogram.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonHistogram.java index f9e066875d..0a325bf7fd 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonHistogram.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonHistogram.java @@ -7,6 +7,7 @@ import org.opensearch.dataprepper.model.event.EventType; +import java.time.Instant; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -190,6 +191,18 @@ public JacksonHistogram.Builder withAggregationTemporality(String aggregationTe return this; } + /** + * Sets the time received for populating event origination time in event handle + * + * @param timeReceived time received + * @return the builder + * @since 2.7 + */ + @Override + public JacksonHistogram.Builder withTimeReceived(final Instant timeReceived) { + return (JacksonHistogram.Builder)super.withTimeReceived(timeReceived); + } + /** * Sets the buckets for this histogram * @param buckets a list of buckets diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonMetric.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonMetric.java index 8d8ebf0f87..4a30f3a637 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonMetric.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonMetric.java @@ -9,6 +9,7 @@ import com.fasterxml.jackson.databind.node.ObjectNode; import org.opensearch.dataprepper.model.event.JacksonEvent; +import java.time.Instant; import java.util.HashMap; import java.util.Iterator; import java.util.List; @@ -32,13 +33,21 @@ public abstract class JacksonMetric extends JacksonEvent implements Metric { protected static final String SCHEMA_URL_KEY = "schemaUrl"; protected static final String EXEMPLARS_KEY = "exemplars"; protected static final String FLAGS_KEY = "flags"; - private final boolean flattenAttributes; + private boolean flattenAttributes; protected JacksonMetric(Builder builder, boolean flattenAttributes) { super(builder); this.flattenAttributes = flattenAttributes; } + public void setFlattenAttributes(boolean flattenAttributes) { + this.flattenAttributes = flattenAttributes; + } + + boolean getFlattenAttributes() { + return flattenAttributes; + } + @Override public String toJsonString() { if (!flattenAttributes) { @@ -226,6 +235,19 @@ public T withSchemaUrl(final String schemaUrl) { return getThis(); } + /** + * Sets the time received for populating event origination time in event handle + * + * @param timeReceived time received + * @return the builder + * @since 2.7 + */ + @Override + public T withTimeReceived(final Instant timeReceived) { + return (T)super.withTimeReceived(timeReceived); + } + + /** * Sets the exemplars that are associated with this metric event * @param exemplars sets the exemplars for this metric diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonSum.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonSum.java index 9835650dd0..a5b9c4e1a0 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonSum.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonSum.java @@ -7,6 +7,7 @@ import org.opensearch.dataprepper.model.event.EventType; +import java.time.Instant; import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; @@ -101,6 +102,18 @@ public Builder withIsMonotonic(final boolean isMonotonic) { return this; } + /** + * Sets the time received for populating event origination time in event handle + * + * @param timeReceived time received + * @return the builder + * @since 2.7 + */ + @Override + public Builder withTimeReceived(final Instant timeReceived) { + return (Builder)super.withTimeReceived(timeReceived); + } + /** * Returns a newly created {@link JacksonSum} * @return a JacksonSum diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonSummary.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonSummary.java index 2196e59075..01425d0c8c 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonSummary.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/metric/JacksonSummary.java @@ -7,6 +7,7 @@ import org.opensearch.dataprepper.model.event.EventType; +import java.time.Instant; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; @@ -106,6 +107,18 @@ public Builder withSum(double sum) { return this; } + /** + * Sets the time received for populating event origination time in event handle + * + * @param timeReceived time received + * @return the builder + * @since 2.7 + */ + @Override + public Builder withTimeReceived(final Instant timeReceived) { + return (Builder)super.withTimeReceived(timeReceived); + } + /** * Sets the count * @param count the count of this summary diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/trace/JacksonSpan.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/trace/JacksonSpan.java index d6bb211ca2..08a7fdccac 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/trace/JacksonSpan.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/trace/JacksonSpan.java @@ -15,6 +15,7 @@ import org.opensearch.dataprepper.model.event.EventType; import org.opensearch.dataprepper.model.event.JacksonEvent; +import java.time.Instant; import java.util.Arrays; import java.util.HashMap; import java.util.Iterator; @@ -444,6 +445,18 @@ public Builder withTraceGroup(final String traceGroup) { return this; } + /** + * Sets the time received for populating event origination time in event handle + * + * @param timeReceived time received + * @return the builder + * @since 2.7 + */ + @Override + public Builder withTimeReceived(final Instant timeReceived) { + return (Builder)super.withTimeReceived(timeReceived); + } + /** * Sets the duration of the span * diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/codec/JsonDecoderTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/codec/JsonDecoderTest.java index 1c3c789c79..d2c7287313 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/codec/JsonDecoderTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/codec/JsonDecoderTest.java @@ -2,9 +2,11 @@ import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; import org.opensearch.dataprepper.model.record.Record; import java.io.ByteArrayInputStream; +import java.time.Instant; import java.util.Map; import java.util.Random; import java.util.UUID; @@ -18,12 +20,13 @@ public class JsonDecoderTest { private JsonDecoder jsonDecoder; private Record receivedRecord; + private Instant receivedTime; private JsonDecoder createObjectUnderTest() { return new JsonDecoder(); } - @BeforeEach +@BeforeEach void setup() { jsonDecoder = createObjectUnderTest(); receivedRecord = null; @@ -36,7 +39,7 @@ void test_basicJsonDecoder() { int intValue = r.nextInt(); String inputString = "[{\"key1\":\""+stringValue+"\", \"key2\":"+intValue+"}]"; try { - jsonDecoder.parse(new ByteArrayInputStream(inputString.getBytes()), (record) -> { + jsonDecoder.parse(new ByteArrayInputStream(inputString.getBytes()), null, (record) -> { receivedRecord = record; }); } catch (Exception e){} @@ -47,4 +50,25 @@ void test_basicJsonDecoder() { assertThat(map.get("key2"), equalTo(intValue)); } + @Test + void test_basicJsonDecoder_withTimeReceived() { + String stringValue = UUID.randomUUID().toString(); + Random r = new Random(); + int intValue = r.nextInt(); + String inputString = "[{\"key1\":\""+stringValue+"\", \"key2\":"+intValue+"}]"; + final Instant now = Instant.now(); + try { + jsonDecoder.parse(new ByteArrayInputStream(inputString.getBytes()), now, (record) -> { + receivedRecord = record; + receivedTime = ((DefaultEventHandle)(((Event)record.getData()).getEventHandle())).getInternalOriginationTime(); + }); + } catch (Exception e){} + + assertNotEquals(receivedRecord, null); + Map map = receivedRecord.getData().toMap(); + assertThat(map.get("key1"), equalTo(stringValue)); + assertThat(map.get("key2"), equalTo(intValue)); + assertThat(receivedTime, equalTo(now)); + } + } diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/JacksonEventTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/JacksonEventTest.java index bf3a320728..b51f947f0d 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/JacksonEventTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/JacksonEventTest.java @@ -23,6 +23,7 @@ import java.util.Random; import java.util.UUID; +import static org.hamcrest.CoreMatchers.containsStringIgnoringCase; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.not; @@ -70,7 +71,7 @@ public void testPutAndGet_withRandomString() { } @ParameterizedTest - @ValueSource(strings = {"foo", "foo-bar", "foo_bar", "foo.bar", "/foo", "/foo/", "a1K.k3-01_02"}) + @ValueSource(strings = {"/", "foo", "foo-bar", "foo_bar", "foo.bar", "/foo", "/foo/", "a1K.k3-01_02", "keyWithBrackets[]"}) void testPutAndGet_withStrings(final String key) { final UUID value = UUID.randomUUID(); @@ -81,6 +82,12 @@ void testPutAndGet_withStrings(final String key) { assertThat(result, is(equalTo(value))); } + @Test + public void testPutKeyCannotBeEmptyString() { + Throwable exception = assertThrows(IllegalArgumentException.class, () -> event.put("", "value")); + assertThat(exception.getMessage(), containsStringIgnoringCase("key cannot be an empty string")); + } + @Test public void testPutAndGet_withMultLevelKey() { final String key = "foo/bar"; @@ -126,6 +133,28 @@ public void testPutAndGet_withMultiLevelKeyWithADash() { assertThat(result, is(equalTo(value))); } + @ParameterizedTest + @ValueSource(strings = {"foo", "/foo", "/foo/", "foo/"}) + void testGetAtRootLevel(final String key) { + final String value = UUID.randomUUID().toString(); + + event.put(key, value); + final Map result = event.get("", Map.class); + + assertThat(result, is(Map.of("foo", value))); + } + + @ParameterizedTest + @ValueSource(strings = {"/foo/bar", "foo/bar", "foo/bar/"}) + void testGetAtRootLevelWithMultiLevelKey(final String key) { + final String value = UUID.randomUUID().toString(); + + event.put(key, value); + final Map result = event.get("", Map.class); + + assertThat(result, is(Map.of("foo", Map.of("bar", value)))); + } + @Test public void testPutUpdateAndGet_withPojo() { final String key = "foo/bar"; @@ -250,10 +279,9 @@ public void testOverwritingExistingKey() { assertThat(result, is(equalTo(value))); } - @Test - public void testDeletingKey() { - final String key = "foo"; - + @ParameterizedTest + @ValueSource(strings = {"/", "foo", "/foo", "/foo/bar", "foo/bar", "foo/bar/", "/foo/bar/leaf/key", "keyWithBrackets[]"}) + public void testDeleteKey(final String key) { event.put(key, UUID.randomUUID()); event.delete(key); final UUID result = event.get(key, UUID.class); @@ -262,19 +290,25 @@ public void testDeletingKey() { } @Test - public void testDelete_withNestedKey() { - final String key = "foo/bar"; - - event.put(key, UUID.randomUUID()); - event.delete(key); - final UUID result = event.get(key, UUID.class); - + public void testClear() { + event.put("key1", UUID.randomUUID()); + event.put("key2", UUID.randomUUID()); + event.put("key3/key4", UUID.randomUUID()); + event.clear(); + UUID result = event.get("key1", UUID.class); + assertThat(result, is(nullValue())); + result = event.get("key2", UUID.class); assertThat(result, is(nullValue())); + result = event.get("key3", UUID.class); + assertThat(result, is(nullValue())); + result = event.get("key3/key4", UUID.class); + assertThat(result, is(nullValue())); + assertThat(event.toMap().size(), equalTo(0)); } - @Test - public void testDelete_withNonexistentKey() { - final String key = "foo/bar"; + @ParameterizedTest + @ValueSource(strings = {"/", "foo", "/foo", "/foo/bar", "foo/bar", "foo/bar/", "/foo/bar/leaf/key"}) + public void testDelete_withNonexistentKey(final String key) { UUID result = event.get(key, UUID.class); assertThat(result, is(nullValue())); @@ -285,19 +319,27 @@ public void testDelete_withNonexistentKey() { } @Test - public void testContainsKey_withKey() { - final String key = "foo"; - - event.put(key, UUID.randomUUID()); - assertThat(event.containsKey(key), is(true)); + public void testDeleteKeyCannotBeEmptyString() { + Throwable exception = assertThrows(IllegalArgumentException.class, () -> event.delete("")); + assertThat(exception.getMessage(), containsStringIgnoringCase("key cannot be an empty string")); } @Test - public void testContainsKey_withouthKey() { - final String key = "foo"; + public void testContainsKeyReturnsTrueForEmptyStringKey() { + assertThat(event.containsKey(""), is(true)); + } + @ParameterizedTest + @ValueSource(strings = {"/", "foo", "/foo", "/foo/bar", "foo/bar", "foo/bar/", "/foo/bar/leaf/key"}) + public void testContainsKey_withKey(final String key) { event.put(key, UUID.randomUUID()); - assertThat(event.containsKey("bar"), is(false)); + assertThat(event.containsKey(key), is(true)); + } + + @ParameterizedTest + @ValueSource(strings = {"/", "foo", "/foo", "/foo/bar", "foo/bar", "foo/bar/", "/foo/bar/leaf/key"}) + public void testContainsKey_withouthKey(final String key) { + assertThat(event.containsKey(key), is(false)); } @Test @@ -324,8 +366,8 @@ public void testIsValueAList_withNull() { } @ParameterizedTest - @ValueSource(strings = {"", "withSpecialChars*$%", "\\-withEscapeChars", "\\\\/withMultipleEscapeChars", - "with,Comma", "with:Colon", "with[Bracket", "with|Brace"}) + @ValueSource(strings = {"withSpecialChars*$%", "\\-withEscapeChars", "\\\\/withMultipleEscapeChars", + "with,Comma", "with:Colon", "with|Brace"}) void testKey_withInvalidKey_throwsIllegalArgumentException(final String invalidKey) { assertThrowsForKeyCheck(IllegalArgumentException.class, invalidKey); } diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/log/JacksonOtelLogTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/log/JacksonOtelLogTest.java index fc8213323b..0fe7cebd9c 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/log/JacksonOtelLogTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/log/JacksonOtelLogTest.java @@ -12,7 +12,9 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.skyscreamer.jsonassert.JSONAssert; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; +import java.time.Instant; import java.util.Date; import java.util.Map; import java.util.UUID; @@ -83,6 +85,14 @@ public void testGetServiceName() { assertThat(name, is(equalTo(TEST_SERVICE_NAME))); } + @Test + public void testGetTimeReceived() { + Instant now = Instant.now(); + builder.withTimeReceived(now); + log = builder.build(); + assertThat(((DefaultEventHandle)log.getEventHandle()).getInternalOriginationTime(), is(now)); + } + @Test public void testGetSchemaUrl() { final String schemaUrl = log.getSchemaUrl(); diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonExponentialHistogramTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonExponentialHistogramTest.java index 31ab4692f5..724d3ee130 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonExponentialHistogramTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonExponentialHistogramTest.java @@ -11,8 +11,10 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.model.event.TestObject; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; import org.skyscreamer.jsonassert.JSONAssert; +import java.time.Instant; import java.time.LocalDateTime; import java.time.ZoneOffset; import java.time.ZonedDateTime; @@ -156,6 +158,9 @@ public void testGetScale() { public void testZeroCount() { Long zeroCount = histogram.getZeroCount(); assertThat(zeroCount, is(equalTo(TEST_ZERO_COUNT))); + assertThat(((JacksonMetric)histogram).getFlattenAttributes(), equalTo(true)); + ((JacksonMetric)histogram).setFlattenAttributes(false); + assertThat(((JacksonMetric)histogram).getFlattenAttributes(), equalTo(false)); } @Test @@ -243,6 +248,14 @@ public void testGetAttributes_withNull_mustBeEmpty() { assertThat(histogram.getAttributes(), is(anEmptyMap())); } + @Test + public void testGetTimeReceived() { + Instant now = Instant.now(); + builder.withTimeReceived(now); + JacksonExponentialHistogram histogram = builder.build(); + assertThat(((DefaultEventHandle)histogram.getEventHandle()).getInternalOriginationTime(), is(now)); + } + @Test public void testHistogramToJsonString() throws JSONException { histogram.put("foo", "bar"); diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonGaugeTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonGaugeTest.java index 842f35bc4f..55fda1085c 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonGaugeTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonGaugeTest.java @@ -11,8 +11,10 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.model.event.TestObject; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; import org.skyscreamer.jsonassert.JSONAssert; +import java.time.Instant; import java.util.Arrays; import java.util.Date; import java.util.List; @@ -89,6 +91,14 @@ public void testGetName() { assertThat(name, is(equalTo(TEST_NAME))); } + @Test + public void testGetTimeReceived() { + Instant now = Instant.now(); + builder.withTimeReceived(now); + gauge = builder.build(); + assertThat(((DefaultEventHandle)gauge.getEventHandle()).getInternalOriginationTime(), is(now)); + } + @Test public void testGetDescription() { final String description = gauge.getDescription(); diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonHistogramTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonHistogramTest.java index 15b22802a2..f9d6cffc11 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonHistogramTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonHistogramTest.java @@ -11,8 +11,10 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.model.event.TestObject; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; import org.skyscreamer.jsonassert.JSONAssert; +import java.time.Instant; import java.util.Arrays; import java.util.Date; import java.util.List; @@ -140,6 +142,14 @@ public void testGetCount() { assertThat(count, is(equalTo(TEST_COUNT))); } + @Test + public void testGetTimeReceived() { + Instant now = Instant.now(); + builder.withTimeReceived(now); + histogram = builder.build(); + assertThat(((DefaultEventHandle)histogram.getEventHandle()).getInternalOriginationTime(), is(now)); + } + @Test public void testGetServiceName() { final String name = histogram.getServiceName(); diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonSumTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonSumTest.java index 58cab11f11..4dd47541e9 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonSumTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonSumTest.java @@ -9,6 +9,9 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; + +import java.time.Instant; import java.util.Date; import java.util.Map; import java.util.UUID; @@ -96,6 +99,14 @@ public void testGetServiceName() { assertThat(name, is(equalTo(TEST_SERVICE_NAME))); } + @Test + public void testGetTimeReceived() { + Instant now = Instant.now(); + builder.withTimeReceived(now); + sum = builder.build(); + assertThat(((DefaultEventHandle)sum.getEventHandle()).getInternalOriginationTime(), is(now)); + } + @Test public void testGetAggregationTemporality() { final String aggregationTemporality = sum.getAggregationTemporality(); diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonSummaryTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonSummaryTest.java index 31820b86c6..6bc865169e 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonSummaryTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/metric/JacksonSummaryTest.java @@ -10,6 +10,9 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; + +import java.time.Instant; import java.util.Arrays; import java.util.Date; import java.util.List; @@ -116,6 +119,14 @@ public void testGetSum() { assertThat(sum, is(equalTo(TEST_SUM))); } + @Test + public void testGetTimeReceived() { + Instant now = Instant.now(); + builder.withTimeReceived(now); + summary = builder.build(); + assertThat(((DefaultEventHandle)summary.getEventHandle()).getInternalOriginationTime(), is(now)); + } + @Test public void testGetCount() { diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/trace/JacksonSpanTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/trace/JacksonSpanTest.java index 2648b172aa..b9b7e6e959 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/trace/JacksonSpanTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/trace/JacksonSpanTest.java @@ -15,6 +15,7 @@ import org.opensearch.dataprepper.model.event.EventMetadata; import org.opensearch.dataprepper.model.event.EventType; import org.opensearch.dataprepper.model.event.JacksonEvent; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; import java.time.Instant; import java.util.Arrays; @@ -440,6 +441,13 @@ public void testBuilder_withNullDroppedAttributesCount_createsSpanWithDefaultVal assertThat(span.getDroppedAttributesCount(), is(equalTo(0))); } + @Test + public void testGetTimeReceived() { + Instant now = Instant.now(); + final JacksonSpan span = builder.withTimeReceived(now).build(); + assertThat(((DefaultEventHandle)span.getEventHandle()).getInternalOriginationTime(), is(now)); + } + @Test public void testBuilder_withNullEvents_createsSpanWithDefaultValue() { final JacksonSpan span = builder.withEvents(null).build(); diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/DataPrepper.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/DataPrepper.java index 1c93f41ec8..19c0822ce9 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/DataPrepper.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/DataPrepper.java @@ -19,6 +19,8 @@ import javax.inject.Inject; import javax.inject.Named; +import java.util.LinkedList; +import java.util.List; import java.util.Map; import java.util.function.Predicate; @@ -45,7 +47,7 @@ public class DataPrepper implements PipelinesProvider { @Inject @Lazy private DataPrepperServer dataPrepperServer; - private DataPrepperShutdownListener shutdownListener; + private List shutdownListeners = new LinkedList<>(); /** * returns serviceName if exists or default serviceName @@ -115,8 +117,8 @@ public void shutdownPipelines() { public void shutdownServers() { dataPrepperServer.stop(); peerForwarderServer.stop(); - if(shutdownListener != null) { - shutdownListener.handleShutdown(); + if(shutdownListeners != null) { + shutdownListeners.forEach(DataPrepperShutdownListener::handleShutdown); } } @@ -139,7 +141,7 @@ public Map getTransformationPipelines() { } public void registerShutdownHandler(final DataPrepperShutdownListener shutdownListener) { - this.shutdownListener = shutdownListener; + this.shutdownListeners.add(shutdownListener); } private class PipelinesObserver implements PipelineObserver { diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/DataPrepperShutdownListener.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/DataPrepperShutdownListener.java index d33fcef363..44b5b19f29 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/DataPrepperShutdownListener.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/DataPrepperShutdownListener.java @@ -5,6 +5,7 @@ package org.opensearch.dataprepper; -interface DataPrepperShutdownListener { +@FunctionalInterface +public interface DataPrepperShutdownListener { void handleShutdown(); } diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/ProcessWorker.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/ProcessWorker.java index 8cfd1ddd43..7e27db0afd 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/ProcessWorker.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/ProcessWorker.java @@ -21,6 +21,7 @@ import org.slf4j.LoggerFactory; import java.util.Collection; +import java.util.Collections; import java.util.List; import java.util.ArrayList; import java.util.Map; @@ -102,10 +103,9 @@ private void processAcknowledgements(List inputEvents, Collection inputEvents = null; if (acknowledgementsEnabled) { - inputEvents = ((ArrayList>)records).stream().map(Record::getData).collect(Collectors.toList()); + inputEvents = ((ArrayList>) records).stream().map(Record::getData).collect(Collectors.toList()); } - records = processor.execute(records); - if (inputEvents != null) { - processAcknowledgements(inputEvents, records); + + try { + records = processor.execute(records); + if (inputEvents != null) { + processAcknowledgements(inputEvents, records); + } + } catch (final Exception e) { + LOG.error("A processor threw an exception. This batch of Events will be dropped, and their EventHandles will be released: ", e); + if (inputEvents != null) { + processAcknowledgements(inputEvents, Collections.emptyList()); + } + + records = Collections.emptyList(); + break; } } diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/AbstractContextManagerTest.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/AbstractContextManagerTest.java index b5eb163a71..38fbfdcf2e 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/AbstractContextManagerTest.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/AbstractContextManagerTest.java @@ -92,7 +92,7 @@ void getDataPrepperBean_closes_application_contexts_when_DataPrepper_shutsdown() objectUnderTest.getDataPrepperBean(); - ArgumentCaptor dataPrepperShutdownListenerArgumentCaptor = ArgumentCaptor.forClass(DataPrepperShutdownListener.class); + final ArgumentCaptor dataPrepperShutdownListenerArgumentCaptor = ArgumentCaptor.forClass(DataPrepperShutdownListener.class); verify(dataPrepper).registerShutdownHandler(dataPrepperShutdownListenerArgumentCaptor.capture()); final DataPrepperShutdownListener shutdownListener = dataPrepperShutdownListenerArgumentCaptor.getValue(); diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/ProcessWorkerTest.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/ProcessWorkerTest.java new file mode 100644 index 0000000000..9b31b20691 --- /dev/null +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/ProcessWorkerTest.java @@ -0,0 +1,211 @@ +package org.opensearch.dataprepper.pipeline; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.MockedStatic; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.model.CheckpointState; +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; +import org.opensearch.dataprepper.model.buffer.Buffer; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.event.EventHandle; +import org.opensearch.dataprepper.model.processor.Processor; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.source.Source; +import org.opensearch.dataprepper.pipeline.common.FutureHelper; +import org.opensearch.dataprepper.pipeline.common.FutureHelperResult; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.Future; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +public class ProcessWorkerTest { + + @Mock + private Pipeline pipeline; + + @Mock + private Buffer buffer; + + @Mock + private Source source; + + private List> sinkFutures; + + private List processors; + + @BeforeEach + void setup() { + when(pipeline.isStopRequested()).thenReturn(false).thenReturn(true); + when(source.areAcknowledgementsEnabled()).thenReturn(false); + when(pipeline.getSource()).thenReturn(source); + when(buffer.isEmpty()).thenReturn(true); + when(pipeline.getPeerForwarderDrainTimeout()).thenReturn(Duration.ofMillis(100)); + when(pipeline.getReadBatchTimeoutInMillis()).thenReturn(500); + + final Future sinkFuture = mock(Future.class); + sinkFutures = List.of(sinkFuture); + when(pipeline.publishToSinks(any())).thenReturn(sinkFutures); + } + + private ProcessWorker createObjectUnderTest() { + return new ProcessWorker(buffer, processors, pipeline); + } + + @Test + void testProcessWorkerHappyPath() { + + final List records = List.of(mock(Record.class)); + final CheckpointState checkpointState = mock(CheckpointState.class); + final Map.Entry readResult = Map.entry(records, checkpointState); + when(buffer.read(pipeline.getReadBatchTimeoutInMillis())).thenReturn(readResult); + + final Processor processor = mock(Processor.class); + when(processor.execute(records)).thenReturn(records); + when(processor.isReadyForShutdown()).thenReturn(true); + processors = List.of(processor); + + final FutureHelperResult futureHelperResult = mock(FutureHelperResult.class); + when(futureHelperResult.getFailedReasons()).thenReturn(Collections.emptyList()); + + + try (final MockedStatic futureHelperMockedStatic = mockStatic(FutureHelper.class)) { + futureHelperMockedStatic.when(() -> FutureHelper.awaitFuturesIndefinitely(sinkFutures)) + .thenReturn(futureHelperResult); + + final ProcessWorker processWorker = createObjectUnderTest(); + + processWorker.run(); + } + } + + @Test + void testProcessWorkerHappyPathWithAcknowledgments() { + + when(source.areAcknowledgementsEnabled()).thenReturn(true); + + final List> records = new ArrayList<>(); + final Record mockRecord = mock(Record.class); + final Event mockEvent = mock(Event.class); + final EventHandle eventHandle = mock(DefaultEventHandle.class); + when(((DefaultEventHandle) eventHandle).getAcknowledgementSet()).thenReturn(mock(AcknowledgementSet.class)); + when(mockRecord.getData()).thenReturn(mockEvent); + when(mockEvent.getEventHandle()).thenReturn(eventHandle); + + records.add(mockRecord); + + final CheckpointState checkpointState = mock(CheckpointState.class); + final Map.Entry readResult = Map.entry(records, checkpointState); + when(buffer.read(pipeline.getReadBatchTimeoutInMillis())).thenReturn(readResult); + + final Processor processor = mock(Processor.class); + when(processor.execute(records)).thenReturn(records); + when(processor.isReadyForShutdown()).thenReturn(true); + processors = List.of(processor); + + final FutureHelperResult futureHelperResult = mock(FutureHelperResult.class); + when(futureHelperResult.getFailedReasons()).thenReturn(Collections.emptyList()); + + + try (final MockedStatic futureHelperMockedStatic = mockStatic(FutureHelper.class)) { + futureHelperMockedStatic.when(() -> FutureHelper.awaitFuturesIndefinitely(sinkFutures)) + .thenReturn(futureHelperResult); + + final ProcessWorker processWorker = createObjectUnderTest(); + + processWorker.run(); + } + } + + @Test + void testProcessWorkerWithProcessorThrowingExceptionIsCaughtProperly() { + + final List records = List.of(mock(Record.class)); + final CheckpointState checkpointState = mock(CheckpointState.class); + final Map.Entry readResult = Map.entry(records, checkpointState); + when(buffer.read(pipeline.getReadBatchTimeoutInMillis())).thenReturn(readResult); + + final Processor processor = mock(Processor.class); + when(processor.execute(records)).thenThrow(RuntimeException.class); + when(processor.isReadyForShutdown()).thenReturn(true); + + final Processor skippedProcessor = mock(Processor.class); + when(skippedProcessor.isReadyForShutdown()).thenReturn(true); + processors = List.of(processor, skippedProcessor); + + final FutureHelperResult futureHelperResult = mock(FutureHelperResult.class); + when(futureHelperResult.getFailedReasons()).thenReturn(Collections.emptyList()); + + + try (final MockedStatic futureHelperMockedStatic = mockStatic(FutureHelper.class)) { + futureHelperMockedStatic.when(() -> FutureHelper.awaitFuturesIndefinitely(sinkFutures)) + .thenReturn(futureHelperResult); + + final ProcessWorker processWorker = createObjectUnderTest(); + + processWorker.run(); + } + + verify(skippedProcessor, never()).execute(any()); + } + + @Test + void testProcessWorkerWithProcessorThrowingExceptionAndAcknowledgmentsEnabledIsHandledProperly() { + + when(source.areAcknowledgementsEnabled()).thenReturn(true); + + final List> records = new ArrayList<>(); + final Record mockRecord = mock(Record.class); + final Event mockEvent = mock(Event.class); + final EventHandle eventHandle = mock(DefaultEventHandle.class); + when(((DefaultEventHandle) eventHandle).getAcknowledgementSet()).thenReturn(mock(AcknowledgementSet.class)); + doNothing().when(eventHandle).release(true); + when(mockRecord.getData()).thenReturn(mockEvent); + when(mockEvent.getEventHandle()).thenReturn(eventHandle); + + records.add(mockRecord); + + final CheckpointState checkpointState = mock(CheckpointState.class); + final Map.Entry readResult = Map.entry(records, checkpointState); + when(buffer.read(pipeline.getReadBatchTimeoutInMillis())).thenReturn(readResult); + + final Processor processor = mock(Processor.class); + when(processor.execute(records)).thenThrow(RuntimeException.class); + when(processor.isReadyForShutdown()).thenReturn(true); + + final Processor skippedProcessor = mock(Processor.class); + when(skippedProcessor.isReadyForShutdown()).thenReturn(true); + processors = List.of(processor, skippedProcessor); + + final FutureHelperResult futureHelperResult = mock(FutureHelperResult.class); + when(futureHelperResult.getFailedReasons()).thenReturn(Collections.emptyList()); + + + try (final MockedStatic futureHelperMockedStatic = mockStatic(FutureHelper.class)) { + futureHelperMockedStatic.when(() -> FutureHelper.awaitFuturesIndefinitely(sinkFutures)) + .thenReturn(futureHelperResult); + + final ProcessWorker processWorker = createObjectUnderTest(); + + processWorker.run(); + } + + verify(skippedProcessor, never()).execute(any()); + } +} diff --git a/data-prepper-expression/build.gradle b/data-prepper-expression/build.gradle index e4e6a65a21..0b1b9e0520 100644 --- a/data-prepper-expression/build.gradle +++ b/data-prepper-expression/build.gradle @@ -42,7 +42,7 @@ jacocoTestCoverageVerification { violationRules { rule { //in addition to core projects rule - this one checks for 100% code coverage for this project limit { - minimum = 1.0 //keep this at 100% + minimum = 1.0 // keep at 100% } } } diff --git a/data-prepper-expression/src/main/antlr/DataPrepperExpression.g4 b/data-prepper-expression/src/main/antlr/DataPrepperExpression.g4 index e2946cad8f..fbe75888bb 100644 --- a/data-prepper-expression/src/main/antlr/DataPrepperExpression.g4 +++ b/data-prepper-expression/src/main/antlr/DataPrepperExpression.g4 @@ -175,7 +175,6 @@ literal | Null ; - Integer : ZERO | NonZeroDigit Digit* @@ -217,7 +216,7 @@ JsonPointerCharacters fragment JsonPointerCharacter - : [A-Za-z0-9_] + : [A-Za-z0-9_.@] ; EscapedJsonPointer diff --git a/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/JoinExpressionFunction.java b/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/JoinExpressionFunction.java index 6ca35782af..d52729757b 100644 --- a/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/JoinExpressionFunction.java +++ b/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/JoinExpressionFunction.java @@ -54,15 +54,15 @@ public Object evaluate(final List args, Event event, Function sourceList = event.get(sourceKey, List.class); - return joinList(sourceList, delimiter); - } catch (Exception e) { - try { + if (event.isValueAList(sourceKey)) { + final List sourceList = event.get(sourceKey, List.class); + return joinList(sourceList, delimiter); + } else { final Map sourceMap = event.get(sourceKey, Map.class); return joinListsInMap(sourceMap, delimiter); - } catch (Exception ex) { - throw new RuntimeException("Unable to perform join function on " + sourceKey, ex); } + } catch (Exception ex) { + throw new RuntimeException("Unable to perform join function on " + sourceKey, ex); } } diff --git a/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/LiteralTypeConversionsConfiguration.java b/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/LiteralTypeConversionsConfiguration.java index f3a9051697..cba74c994e 100644 --- a/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/LiteralTypeConversionsConfiguration.java +++ b/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/LiteralTypeConversionsConfiguration.java @@ -9,6 +9,8 @@ import javax.inject.Named; import java.io.Serializable; +import java.util.ArrayList; +import java.util.LinkedHashMap; import java.util.Map; import java.util.function.Function; @@ -22,6 +24,8 @@ public Map, Function> literalTypeC Integer.class, Function.identity(), Float.class, Function.identity(), Long.class, Function.identity(), + ArrayList.class, Function.identity(), + LinkedHashMap.class, Function.identity(), Double.class, o -> ((Double) o).floatValue() ); } diff --git a/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluator_ConditionalIT.java b/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluator_ConditionalIT.java index 55ceb338a8..0bef1a65a0 100644 --- a/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluator_ConditionalIT.java +++ b/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluator_ConditionalIT.java @@ -152,6 +152,8 @@ private static Stream validExpressionArguments() { Arguments.of("/status_code == 200", event("{}"), false), Arguments.of("/success == /status_code", event("{\"success\": true, \"status_code\": 200}"), false), Arguments.of("/success != /status_code", event("{\"success\": true, \"status_code\": 200}"), true), + Arguments.of("/part1@part2.part3 != 111", event("{\"success\": true, \"part1@part2.part3\":111, \"status_code\": 200}"), false), + Arguments.of("/part1.part2@part3 != 111", event("{\"success\": true, \"part1.part2@part3\":222, \"status_code\": 200}"), true), Arguments.of("/pi == 3.14159", event("{\"pi\": 3.14159}"), true), Arguments.of("/value == 12345.678", event("{\"value\": 12345.678}"), true), Arguments.of("/value == 12345.678E12", event("{\"value\": 12345.678E12}"), true), @@ -208,7 +210,11 @@ private static Stream validExpressionArguments() { Arguments.of("cidrContains(/sourceIp,\"192.0.2.0/24\",\"192.1.1.0/24\")", event("{\"sourceIp\": \"192.2.2.3\"}"), false), Arguments.of("cidrContains(/sourceIp,\"2001:0db8::/32\")", event("{\"sourceIp\": \"2001:0db8:aaaa:bbbb::\"}"), true), Arguments.of("cidrContains(/sourceIp,\"2001:0db8::/32\",\"2001:aaaa::/32\")", event("{\"sourceIp\": \"2001:0db8:aaaa:bbbb::\"}"), true), - Arguments.of("cidrContains(/sourceIp,\"2001:0db8::/32\",\"2001:aaaa::/32\")", event("{\"sourceIp\": \"2001:abcd:aaaa:bbbb::\"}"), false) + Arguments.of("cidrContains(/sourceIp,\"2001:0db8::/32\",\"2001:aaaa::/32\")", event("{\"sourceIp\": \"2001:abcd:aaaa:bbbb::\"}"), false), + Arguments.of("/sourceIp != null", event("{\"sourceIp\": [10, 20]}"), true), + Arguments.of("/sourceIp == null", event("{\"sourceIp\": [\"test\", \"test_two\"]}"), false), + Arguments.of("/sourceIp == null", event("{\"sourceIp\": {\"test\": \"test_two\"}}"), false), + Arguments.of("/sourceIp != null", event("{\"sourceIp\": {\"test\": \"test_two\"}}"), true) ); } diff --git a/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/JoinExpressionFunctionTest.java b/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/JoinExpressionFunctionTest.java index 48da627e5e..3ec60a73c5 100644 --- a/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/JoinExpressionFunctionTest.java +++ b/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/JoinExpressionFunctionTest.java @@ -94,6 +94,13 @@ void testSourceFieldNotExistsInEventThrowsException() { () -> joinExpressionFunction.evaluate(List.of("/missingKey"), testEvent, testFunction)); } + @Test + void testSourceFieldNotListOrMapThrowsException() { + testEvent = createTestEvent(Map.of("key", "value")); + assertThrows(RuntimeException.class, + () -> joinExpressionFunction.evaluate(List.of("/key"), testEvent, testFunction)); + } + private static Stream joinSingleList() { final String inputData = "{\"list\":[\"string\", 1, true]}"; return Stream.of( diff --git a/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/ParseTreeCoercionServiceTest.java b/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/ParseTreeCoercionServiceTest.java index 56b46406a4..4dca1dc812 100644 --- a/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/ParseTreeCoercionServiceTest.java +++ b/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/ParseTreeCoercionServiceTest.java @@ -5,11 +5,11 @@ package org.opensearch.dataprepper.expression; -import org.opensearch.dataprepper.model.event.Event; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; import org.antlr.v4.runtime.Token; import org.antlr.v4.runtime.tree.TerminalNode; +import org.apache.commons.lang3.RandomStringUtils; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.junit.jupiter.params.ParameterizedTest; @@ -19,12 +19,13 @@ import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.dataprepper.expression.antlr.DataPrepperExpressionParser; import org.opensearch.dataprepper.expression.util.TestObject; -import org.apache.commons.lang3.RandomStringUtils; +import org.opensearch.dataprepper.model.event.Event; import java.util.HashMap; -import java.util.Map; import java.util.List; +import java.util.Map; import java.util.Random; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Function; import java.util.stream.Stream; @@ -35,8 +36,8 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -153,20 +154,6 @@ void testCoerceTerminalNodeJsonPointerTypeSupportedValues(final Object testValue } } - @ParameterizedTest - @MethodSource("provideUnSupportedJsonPointerValues") - void testCoerceTerminalNodeJsonPointerTypeUnSupportedValues(final Object testValue) { - final String testKey1 = "key1"; - final String testKey2 = "key2"; - final String testJsonPointerKey = String.format("/%s/%s", testKey1, testKey2); - final Event testEvent = testValue == null ? createTestEvent(new HashMap<>()) : - createTestEvent(Map.of(testKey1, Map.of(testKey2, testValue))); - when(token.getType()).thenReturn(DataPrepperExpressionParser.JsonPointer); - when(terminalNode.getSymbol()).thenReturn(token); - when(terminalNode.getText()).thenReturn(testJsonPointerKey); - assertThrows(ExpressionCoercionException.class, () -> objectUnderTest.coercePrimaryTerminalNode(terminalNode, testEvent)); - } - @ParameterizedTest @MethodSource("provideKeys") void testCoerceTerminalNodeEscapeJsonPointerTypeWithSpecialCharacters(final String testKey, final String testEscapeJsonPointer) @@ -200,19 +187,32 @@ void testCoerceTerminalNodeEscapeJsonPointerTypeSupportedValues(final Object tes } } - @ParameterizedTest - @MethodSource("provideUnSupportedJsonPointerValues") - void testCoerceTerminalNodeEscapeJsonPointerTypeUnSupportedValues(final Object testValue) { + @Test + void testCoerceTerminalNodeEscapeJsonPointerTypeUnSupportedValues() { final String testKey = "testKey"; + final AtomicBoolean testValue = new AtomicBoolean(); final String testEscapeJsonPointerKey = String.format("\"/%s\"", testKey); - final Event testEvent = testValue == null ? createTestEvent(new HashMap<>()) : - createTestEvent(Map.of(testKey, testValue)); + final Event testEvent = createInvalidTestEvent(Map.of(testKey, testValue)); when(token.getType()).thenReturn(DataPrepperExpressionParser.EscapedJsonPointer); when(terminalNode.getSymbol()).thenReturn(token); when(terminalNode.getText()).thenReturn(testEscapeJsonPointerKey); assertThrows(ExpressionCoercionException.class, () -> objectUnderTest.coercePrimaryTerminalNode(terminalNode, testEvent)); } + @Test + void testCoerceTerminalNodeJsonPointerTypeUnSupportedValues() { + final String testKey1 = "key1"; + final String testKey2 = "key2"; + final AtomicBoolean testValue = new AtomicBoolean(); + final String testJsonPointerKey = String.format("/%s/%s", testKey1, testKey2); + + final Event testEvent = createInvalidTestEvent(Map.of(testKey1, testValue)); + when(token.getType()).thenReturn(DataPrepperExpressionParser.JsonPointer); + when(terminalNode.getSymbol()).thenReturn(token); + when(terminalNode.getText()).thenReturn(testJsonPointerKey); + assertThrows(ExpressionCoercionException.class, () -> objectUnderTest.coercePrimaryTerminalNode(terminalNode, testEvent)); + } + @Test void testCoerceTerminalNodeUnsupportedType() { final Event testEvent = createTestEvent(new HashMap<>()); @@ -333,6 +333,12 @@ private Event createTestEvent(final Object data) { return event; } + private Event createInvalidTestEvent(final Object data) { + final Event event = mock(Event.class); + lenient().when(event.get(anyString(), any())).thenReturn(new AtomicBoolean()); + return event; + } + private static Stream provideKeys() { return Stream.of( Arguments.of("test key", "\"/test key\""), @@ -350,11 +356,9 @@ private static Stream provideSupportedJsonPointerValues() { Arguments.of("test value"), Arguments.of(1.1f), Arguments.of(1.1), + Arguments.of(List.of("test_value")), + Arguments.of(Map.of("test_key", "test_value")), Arguments.of((Object) null) ); } - - private static Stream provideUnSupportedJsonPointerValues() { - return Stream.of(Arguments.of(new HashMap<>())); - } } diff --git a/data-prepper-plugins/avro-codecs/src/test/java/org/opensearch/dataprepper/plugins/codec/avro/AvroOutputCodecTest.java b/data-prepper-plugins/avro-codecs/src/test/java/org/opensearch/dataprepper/plugins/codec/avro/AvroOutputCodecTest.java index e3b030c37f..7185e3ad7e 100644 --- a/data-prepper-plugins/avro-codecs/src/test/java/org/opensearch/dataprepper/plugins/codec/avro/AvroOutputCodecTest.java +++ b/data-prepper-plugins/avro-codecs/src/test/java/org/opensearch/dataprepper/plugins/codec/avro/AvroOutputCodecTest.java @@ -75,7 +75,6 @@ void constructor_throws_if_schema_is_invalid() { RuntimeException actualException = assertThrows(RuntimeException.class, this::createObjectUnderTest); assertThat(actualException.getMessage(), notNullValue()); - assertThat(actualException.getMessage(), containsString(invalidSchema)); assertThat(actualException.getMessage(), containsString("was expecting comma")); } diff --git a/data-prepper-plugins/cloudwatch-logs/build.gradle b/data-prepper-plugins/cloudwatch-logs/build.gradle index 983db202fa..dc374997f0 100644 --- a/data-prepper-plugins/cloudwatch-logs/build.gradle +++ b/data-prepper-plugins/cloudwatch-logs/build.gradle @@ -3,10 +3,6 @@ plugins { id 'java-library' } -repositories { - mavenCentral() -} - dependencies { implementation project(':data-prepper-plugins:aws-plugin-api') implementation project(path: ':data-prepper-plugins:common') diff --git a/data-prepper-plugins/csv-processor/README.md b/data-prepper-plugins/csv-processor/README.md index d6f5e81643..e8afcd0080 100644 --- a/data-prepper-plugins/csv-processor/README.md +++ b/data-prepper-plugins/csv-processor/README.md @@ -1,163 +1,7 @@ # CSV Processor This is a processor that takes in an Event and parses its CSV data into columns. -## Basic Usage -### User Specified Column Names -To get started, create the following `pipelines.yaml`. -```yaml -csv-pipeline: - source: - file: - path: "/full/path/to/ingest.csv" - record_type: "event" - processor: - - csv: - column_names: ["col1", "col2"] - sink: - - stdout: -``` -#### Example With User Specified Column Names Config: -If you wish to test the CSV Processor with the above config then you may find the following example useful. Create the following file `ingest.csv` and replace the `path` in the file source of your `pipeline.yaml` with the path of this file. -``` -1,2,3 -``` -When run, the processor will parse the message. Notice that since there are only two column names specified in the config, the third column name is autogenerated. -``` -{"message": "1,2,3", "col1": "1", "col2": "2", "column3": "3"} -``` -### Auto Detect Column Names -The following configuration auto detects the header of a CSV file ingested through S3 Source. See the [S3 Source Documentation](https://github.com/opensearch-project/data-prepper/tree/main/data-prepper-plugins/s3-source) for more information. -```yaml -csv-s3-pipeline: - source: - s3: - notification_type: "sqs" - codec: - newline: - skip_lines: 1 - header_destination: "header" - compression: none - sqs: - queue_url: "https://sqs..amazonaws.com//" - aws: - region: "" - processor: - - csv: - column_names_source_key: "header" - sink: - - stdout: -``` -#### Example With Auto Detect Column Names Config: -If you wish to test the CSV Processor with the above config then you may find the following example useful. Upload the following file `ingest.csv` to the S3 bucket that your SQS queue is attached to: -``` -Should,skip,this,line -a,b,c -1,2,3 -``` -When run, the processor will process the following event: -```json -{"header": "a,b,c", "message": "1,2,3"} -``` -And will parse it into the following. (Note that since `delete_header` is `true`, the header is deleted): -```json -{"message": "1,2,3", "a": "1", "b": "2", "c": "3"} -``` -## Configuration -* `source` — The field in the `Event` that will be parsed. - * Default: `message` - -* `quote_character` — The character used as a text qualifier for a single column of data. - * Default: Double quote (`"`) - -* `delimiter` — The character separating each column. - * Default: `,` - -* `delete_header` (Optional, boolean) — If specified, the header on the `Event` (`column_names_source_key`) will be deleted after the `Event` is parsed. If there’s no header on the `Event` then nothing will happen. - * Default: `true` - -* `column_names_source_key` — (Optional) The field in the Event that specifies the CSV column names, which will be autodetected. If there must be extra column names, they will be autogenerated according to their index. - * There is no default. - * Note: If `column_names` is also defined, the header in `column_names_source_key` will be used to generate the Event fields. - * Note: If too few columns are specified in this field, the remaining column names will be autogenerated. If too many column names are specified in this field, then the extra column names will be omitted. - -* `column_names` — (Optional) User-specified names for the CSV columns. - * Default: `[column1, column2, ..., columnN]` if there are `N` columns of data in the CSV record and `column_names_source_key` is not defined. - * Note: If `column_names_source_key` is defined, the header in `column_names_source_key` will be used to generate the Event fields. - * Note: If too few columns are specified in this field, the remaining column names will be autogenerated. If too many column names are specified in this field, then the extra column names will be omitted. - -## Metrics - -Apart from common metrics in [AbstractProcessor](https://github.com/opensearch-project/data-prepper/blob/main/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/processor/AbstractProcessor.java), the CSV Processor includes the following custom metric. - -**Counter** - -* `csvInvalidEvents`: The number of invalid Events. An invalid Event causes an Exception to be thrown when parsed. This is most commonly due to an unclosed quote. - -# CSV Sink/Output Codec - -This is an implementation of CSV Sink Codec that parses the Dataprepper Events into CSV rows and writes them into the underlying OutputStream. - -## Usages - -CSV Codec can be configured with sink plugins (e.g. S3 Sink) in the Pipeline file. - -## Configuration Options - -``` -pipeline: - ... - sink: - - s3: - aws: - region: us-east-1 - sts_role_arn: arn:aws:iam::123456789012:role/Data-Prepper - sts_header_overrides: - max_retries: 5 - bucket: bucket_name - object_key: - path_prefix: my-elb/%{yyyy}/%{MM}/%{dd}/ - threshold: - event_count: 2000 - maximum_size: 50mb - event_collect_timeout: 15s - codec: - csv: - delimiter: "," - header: - - Year - - Age - - Ethnic - - Sex - - Area - - count - exclude_keys: - - s3 - buffer_type: in_memory -``` - -### Note: - -1) If the user wants the tags to be a part of the resultant CSV Data and has given `tagsTargetKey` in the config file, the user also has to modify the header to accommodate the tags. Another header field has to be provided in the headers: - - ``` - header: - - Year - - Age - - Ethnic - - Sex - - Area - - - ``` - Please note that if this is not done, then the codec will throw an error: - `"CSV Row doesn't conform with the header."` - -## AWS Configuration - -### Codec Configuration: - -1) `exclude_keys`: Those keys of the events that the user wants to exclude while converting them to CSV Rows. -2) `delimiter`: The user can provide the delimiter of choice. -3) `header`: The user can provide the desired header for the resultant CSV file. +https://opensearch.org/docs/latest/data-prepper/pipelines/configuration/processors/csv ## Developer Guide This plugin is compatible with Java 8 and up. See diff --git a/data-prepper-plugins/csv-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/csv/CsvProcessor.java b/data-prepper-plugins/csv-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/csv/CsvProcessor.java index 55bd7301cb..a9f99e5862 100644 --- a/data-prepper-plugins/csv-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/csv/CsvProcessor.java +++ b/data-prepper-plugins/csv-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/csv/CsvProcessor.java @@ -5,10 +5,12 @@ package org.opensearch.dataprepper.plugins.processor.csv; +import org.opensearch.dataprepper.expression.ExpressionEvaluator; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.plugin.InvalidPluginConfigurationException; import org.opensearch.dataprepper.model.processor.AbstractProcessor; import org.opensearch.dataprepper.model.processor.Processor; import org.opensearch.dataprepper.model.record.Record; @@ -42,11 +44,23 @@ public class CsvProcessor extends AbstractProcessor, Record private final CsvProcessorConfig config; + private final ExpressionEvaluator expressionEvaluator; + @DataPrepperPluginConstructor - public CsvProcessor(final PluginMetrics pluginMetrics, final CsvProcessorConfig config) { + public CsvProcessor(final PluginMetrics pluginMetrics, + final CsvProcessorConfig config, + final ExpressionEvaluator expressionEvaluator) { super(pluginMetrics); this.csvInvalidEventsCounter = pluginMetrics.counter(CSV_INVALID_EVENTS); this.config = config; + this.expressionEvaluator = expressionEvaluator; + + if (config.getCsvWhen() != null + && !expressionEvaluator.isValidExpressionStatement(config.getCsvWhen())) { + throw new InvalidPluginConfigurationException( + String.format("csv_when value of %s is not a valid expression statement. " + + "See https://opensearch.org/docs/latest/data-prepper/pipelines/expression-syntax/ for valid expression syntax.", config.getCsvWhen())); + } } @Override @@ -58,16 +72,21 @@ public Collection> doExecute(final Collection> recor final Event event = record.getData(); - final String message = event.get(config.getSource(), String.class); + try { - if (Objects.isNull(message)) { - continue; - } + if (config.getCsvWhen() != null && !expressionEvaluator.evaluateConditional(config.getCsvWhen(), event)) { + continue; + } - final boolean userDidSpecifyHeaderEventKey = Objects.nonNull(config.getColumnNamesSourceKey()); - final boolean thisEventHasHeaderSource = userDidSpecifyHeaderEventKey && event.containsKey(config.getColumnNamesSourceKey()); + final String message = event.get(config.getSource(), String.class); + + if (Objects.isNull(message)) { + continue; + } + + final boolean userDidSpecifyHeaderEventKey = Objects.nonNull(config.getColumnNamesSourceKey()); + final boolean thisEventHasHeaderSource = userDidSpecifyHeaderEventKey && event.containsKey(config.getColumnNamesSourceKey()); - try { final MappingIterator> messageIterator = mapper.readerFor(List.class).with(schema).readValues(message); // otherwise the message is empty diff --git a/data-prepper-plugins/csv-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/csv/CsvProcessorConfig.java b/data-prepper-plugins/csv-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/csv/CsvProcessorConfig.java index 4e6500b675..ec5d685b7e 100644 --- a/data-prepper-plugins/csv-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/csv/CsvProcessorConfig.java +++ b/data-prepper-plugins/csv-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/csv/CsvProcessorConfig.java @@ -37,6 +37,9 @@ public class CsvProcessorConfig { @JsonProperty("column_names") private List columnNames; + @JsonProperty("csv_when") + private String csvWhen; + /** * The field of the Event that contains the CSV data to be processed. * @@ -93,6 +96,8 @@ public List getColumnNames() { return columnNames; } + public String getCsvWhen() { return csvWhen; } + @AssertTrue(message = "delimiter must be exactly one character.") boolean isValidDelimiter() { return delimiter.length() == 1; diff --git a/data-prepper-plugins/csv-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/csv/CsvProcessorIT.java b/data-prepper-plugins/csv-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/csv/CsvProcessorIT.java index 33e54c5952..417a997e12 100644 --- a/data-prepper-plugins/csv-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/csv/CsvProcessorIT.java +++ b/data-prepper-plugins/csv-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/csv/CsvProcessorIT.java @@ -5,6 +5,10 @@ package org.opensearch.dataprepper.plugins.processor.csv; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.expression.ExpressionEvaluator; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.record.Record; @@ -20,6 +24,7 @@ import static org.hamcrest.Matchers.equalTo; import static org.opensearch.dataprepper.test.helper.ReflectivelySetField.setField; +@ExtendWith(MockitoExtension.class) public class CsvProcessorIT { private static final String PLUGIN_NAME = "csv"; private static final String TEST_PIPELINE_NAME = "test_pipeline"; @@ -28,6 +33,9 @@ public class CsvProcessorIT { private CsvProcessor csvProcessor; private VpcFlowLogTypeGenerator vpcFlowLogTypeGenerator; + @Mock + private ExpressionEvaluator expressionEvaluator; + @BeforeEach void setup() { csvProcessorConfig = new CsvProcessorConfig(); @@ -39,7 +47,7 @@ void setup() { PluginMetrics pluginMetrics = PluginMetrics.fromNames(PLUGIN_NAME, TEST_PIPELINE_NAME); - csvProcessor = new CsvProcessor(pluginMetrics, csvProcessorConfig); + csvProcessor = new CsvProcessor(pluginMetrics, csvProcessorConfig, expressionEvaluator); vpcFlowLogTypeGenerator = new VpcFlowLogTypeGenerator(); } diff --git a/data-prepper-plugins/csv-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/csv/CsvProcessorTest.java b/data-prepper-plugins/csv-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/csv/CsvProcessorTest.java index be6da7e80f..5239679fab 100644 --- a/data-prepper-plugins/csv-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/csv/CsvProcessorTest.java +++ b/data-prepper-plugins/csv-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/csv/CsvProcessorTest.java @@ -5,16 +5,18 @@ package org.opensearch.dataprepper.plugins.processor.csv; -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.event.Event; -import org.opensearch.dataprepper.model.event.JacksonEvent; -import org.opensearch.dataprepper.model.record.Record; import io.micrometer.core.instrument.Counter; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.expression.ExpressionEvaluator; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.event.JacksonEvent; +import org.opensearch.dataprepper.model.plugin.InvalidPluginConfigurationException; +import org.opensearch.dataprepper.model.record.Record; import java.util.Arrays; import java.util.Collections; @@ -25,6 +27,7 @@ import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; @@ -39,6 +42,9 @@ class CsvProcessorTest { @Mock private Counter csvInvalidEventsCounter; + @Mock + private ExpressionEvaluator expressionEvaluator; + private CsvProcessor csvProcessor; @BeforeEach @@ -57,14 +63,22 @@ void setup() { } private CsvProcessor createObjectUnderTest() { - return new CsvProcessor(pluginMetrics, processorConfig); + return new CsvProcessor(pluginMetrics, processorConfig, expressionEvaluator); } @Test void do_nothing_when_source_is_null_value_or_does_not_exist_in_the_Event() { + final Record eventUnderTest = createMessageEvent(""); + + final String csvWhen = UUID.randomUUID().toString(); + when(processorConfig.getCsvWhen()).thenReturn(csvWhen); + when(expressionEvaluator.isValidExpressionStatement(csvWhen)).thenReturn(true); + when(expressionEvaluator.evaluateConditional(csvWhen, eventUnderTest.getData())).thenReturn(true); when(processorConfig.getSource()).thenReturn(UUID.randomUUID().toString()); + csvProcessor = createObjectUnderTest(); + final List> editedEvents = (List>) csvProcessor.doExecute(Collections.singletonList(eventUnderTest)); final Event parsedEvent = getSingleEvent(editedEvents); @@ -328,6 +342,35 @@ void test_when_invalidEvent_then_metricIncrementedAndNoException() { assertThat(parsedEvent.containsKey("column3"), equalTo(false)); } + @Test + void invalid_csv_when_throws_InvalidPluginConfigurationException() { + final String csvWhen = UUID.randomUUID().toString(); + + when(processorConfig.getCsvWhen()).thenReturn(csvWhen); + when(expressionEvaluator.isValidExpressionStatement(csvWhen)).thenReturn(false); + + assertThrows(InvalidPluginConfigurationException.class, this::createObjectUnderTest); + } + + @Test + void do_nothing_when_expression_evaluation_returns_false_for_event() { + final String csvWhen = UUID.randomUUID().toString(); + + when(processorConfig.getCsvWhen()).thenReturn(csvWhen); + when(expressionEvaluator.isValidExpressionStatement(csvWhen)).thenReturn(true); + + final Record eventUnderTest = createMessageEvent(""); + when(expressionEvaluator.evaluateConditional(csvWhen, eventUnderTest.getData())).thenReturn(false); + + csvProcessor = createObjectUnderTest(); + + + final List> editedEvents = (List>) csvProcessor.doExecute(Collections.singletonList(eventUnderTest)); + final Event parsedEvent = getSingleEvent(editedEvents); + + assertThat(parsedEvent, equalTo(eventUnderTest.getData())); + } + private Record createMessageEvent(final String message) { final Map eventData = new HashMap<>(); eventData.put("message",message); diff --git a/data-prepper-plugins/date-processor/README.md b/data-prepper-plugins/date-processor/README.md index b23962faab..3ba1647498 100644 --- a/data-prepper-plugins/date-processor/README.md +++ b/data-prepper-plugins/date-processor/README.md @@ -66,7 +66,7 @@ valid key and at least one pattern is required if match is configured. * `patterns`: List of possible patterns the timestamp value of key can have. The patterns are based on sequence of letters and symbols. The `patterns` support all the patterns listed in Java [DatetimeFormatter](https://docs.oracle.com/javase/8/docs/api/java/time/format/DateTimeFormatter.html). - and also supports `epoch_second`, `epoch_milli` and `epoch_nano` values which represents the timestamp as the number of seconds, milliseconds and nano seconds since epoch. Epoch values are always UTC time zone. + and also supports `epoch_second`, `epoch_milli`, `epoch_micro` and `epoch_nano` values which represents the timestamp as the number of seconds, milliseconds, microseconds and nano seconds since epoch. Epoch values are always UTC time zone. * Type: `List` The following example of date configuration will use `timestamp` key to match against given patterns and stores the timestamp in ISO 8601 diff --git a/data-prepper-plugins/date-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessor.java b/data-prepper-plugins/date-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessor.java index 6d43cfbff2..a494cf5334 100644 --- a/data-prepper-plugins/date-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessor.java +++ b/data-prepper-plugins/date-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessor.java @@ -5,6 +5,8 @@ package org.opensearch.dataprepper.plugins.processor.date; +import io.micrometer.core.instrument.Counter; +import org.apache.commons.lang3.tuple.Pair; import org.opensearch.dataprepper.expression.ExpressionEvaluator; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; @@ -13,23 +15,21 @@ import org.opensearch.dataprepper.model.processor.AbstractProcessor; import org.opensearch.dataprepper.model.processor.Processor; import org.opensearch.dataprepper.model.record.Record; -import io.micrometer.core.instrument.Counter; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.time.Instant; import java.time.LocalDate; -import java.time.ZonedDateTime; import java.time.ZoneId; +import java.time.ZonedDateTime; import java.time.format.DateTimeFormatter; import java.time.format.DateTimeFormatterBuilder; import java.time.temporal.ChronoField; import java.util.Collection; import java.util.List; -import java.util.Set; import java.util.Objects; +import java.util.Set; import java.util.stream.Collectors; -import org.apache.commons.lang3.tuple.Pair; @DataPrepperPlugin(name = "date", pluginType = Processor.class, pluginConfigurationType = DateProcessorConfig.class) public class DateProcessor extends AbstractProcessor, Record> { @@ -37,7 +37,7 @@ public class DateProcessor extends AbstractProcessor, Record> doExecute(Collection> records) { for(final Record record : records) { - if (Objects.nonNull(dateProcessorConfig.getDateWhen()) && !expressionEvaluator.evaluateConditional(dateProcessorConfig.getDateWhen(), record.getData())) { - continue; - } - - String zonedDateTime = null; - - if (Boolean.TRUE.equals(dateProcessorConfig.getFromTimeReceived())) { - zonedDateTime = getDateTimeFromTimeReceived(record); + try { + if (Objects.nonNull(dateProcessorConfig.getDateWhen()) && !expressionEvaluator.evaluateConditional(dateProcessorConfig.getDateWhen(), record.getData())) { + continue; + } - } else if (keyToParse != null && !keyToParse.isEmpty()) { - Pair result = getDateTimeFromMatch(record); - if (result != null) { - zonedDateTime = result.getLeft(); - Instant timeStamp = result.getRight(); - if (dateProcessorConfig.getToOriginationMetadata()) { - Event event = (Event)record.getData(); - event.getMetadata().setExternalOriginationTime(timeStamp); - event.getEventHandle().setExternalOriginationTime(timeStamp); + String zonedDateTime = null; + + if (Boolean.TRUE.equals(dateProcessorConfig.getFromTimeReceived())) { + zonedDateTime = getDateTimeFromTimeReceived(record); + + } else if (keyToParse != null && !keyToParse.isEmpty()) { + Pair result = getDateTimeFromMatch(record); + if (result != null) { + zonedDateTime = result.getLeft(); + Instant timeStamp = result.getRight(); + if (dateProcessorConfig.getToOriginationMetadata()) { + Event event = (Event)record.getData(); + event.getMetadata().setExternalOriginationTime(timeStamp); + event.getEventHandle().setExternalOriginationTime(timeStamp); + } } + populateDateProcessorMetrics(zonedDateTime); } - populateDateProcessorMetrics(zonedDateTime); - } - if (zonedDateTime != null) { - record.getData().put(dateProcessorConfig.getDestination(), zonedDateTime); + if (zonedDateTime != null) { + record.getData().put(dateProcessorConfig.getDestination(), zonedDateTime); + } + } catch (final Exception e) { + LOG.error("An exception occurred while attempting to process Event: ", e); } + } return records; } @@ -160,6 +165,9 @@ private Pair getEpochFormatOutput(Instant time) { return Pair.of(Long.toString(time.getEpochSecond()), time); } else if (outputFormat.equals("epoch_milli")) { return Pair.of(Long.toString(time.toEpochMilli()), time); + } else if (outputFormat.equals("epoch_micro")) { + long micro = (long)time.getEpochSecond() * 1000_000 + (long) time.getNano() / 1000; + return Pair.of(Long.toString(micro), time); } else { // epoch_nano. validation for valid epoch_ should be // done at init time long nano = (long)time.getEpochSecond() * 1000_000_000 + (long) time.getNano(); @@ -182,13 +190,20 @@ private Pair getFormattedDateTimeString(final String sourceTime } if (numberValue != null) { int timestampLength = sourceTimestamp.length(); - if (timestampLength > LENGTH_OF_EPOCH_IN_MILLIS) { + if (timestampLength > LENGTH_OF_EPOCH_MICROSECONDS) { if (epochFormatters.contains("epoch_nano")) { epochTime = Instant.ofEpochSecond(numberValue/1000_000_000, numberValue % 1000_000_000); } else { LOG.warn("Source time value is larger than epoch pattern configured. epoch_nano is expected but not present in the patterns list"); return null; } + } else if (timestampLength > LENGTH_OF_EPOCH_IN_MILLIS) { + if (epochFormatters.contains("epoch_micro")) { + epochTime = Instant.ofEpochSecond(numberValue/1000_000, (numberValue % 1000_000) * 1000); + } else { + LOG.warn("Source time value is larger than epoch pattern configured. epoch_micro is expected but not present in the patterns list"); + return null; + } } else if (timestampLength > LENGTH_OF_EPOCH_SECONDS) { if (epochFormatters.contains("epoch_milli")) { epochTime = Instant.ofEpochMilli(numberValue); diff --git a/data-prepper-plugins/date-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorConfig.java b/data-prepper-plugins/date-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorConfig.java index 8c9c37957d..a74b2e9d38 100644 --- a/data-prepper-plugins/date-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorConfig.java +++ b/data-prepper-plugins/date-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorConfig.java @@ -67,6 +67,7 @@ public boolean isValidPatterns() { public static boolean isValidPattern(final String pattern) { if (pattern.equals("epoch_second") || pattern.equals("epoch_milli") || + pattern.equals("epoch_micro") || pattern.equals("epoch_nano")) { return true; } diff --git a/data-prepper-plugins/date-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorConfigTest.java b/data-prepper-plugins/date-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorConfigTest.java index d6cc9a9149..b1dddfa013 100644 --- a/data-prepper-plugins/date-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorConfigTest.java +++ b/data-prepper-plugins/date-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorConfigTest.java @@ -80,6 +80,8 @@ void testValidAndInvalidOutputFormats() throws NoSuchFieldException, IllegalAcce assertThat(dateProcessorConfig.isValidOutputFormat(), equalTo(true)); setField(DateProcessorConfig.class, dateProcessorConfig, "outputFormat", "epoch_nano"); assertThat(dateProcessorConfig.isValidOutputFormat(), equalTo(true)); + setField(DateProcessorConfig.class, dateProcessorConfig, "outputFormat", "epoch_micro"); + assertThat(dateProcessorConfig.isValidOutputFormat(), equalTo(true)); setField(DateProcessorConfig.class, dateProcessorConfig, "outputFormat", "epoch_xyz"); assertThat(dateProcessorConfig.isValidOutputFormat(), equalTo(false)); setField(DateProcessorConfig.class, dateProcessorConfig, "outputFormat", "yyyy-MM-dd'T'HH:mm:ss.nnnnnnnnnXXX"); diff --git a/data-prepper-plugins/date-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorTests.java b/data-prepper-plugins/date-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorTests.java index 6027c1be77..c6688d08e3 100644 --- a/data-prepper-plugins/date-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorTests.java +++ b/data-prepper-plugins/date-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/date/DateProcessorTests.java @@ -5,11 +5,6 @@ package org.opensearch.dataprepper.plugins.processor.date; -import org.opensearch.dataprepper.expression.ExpressionEvaluator; -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.event.Event; -import org.opensearch.dataprepper.model.event.JacksonEvent; -import org.opensearch.dataprepper.model.record.Record; import io.micrometer.core.instrument.Counter; import org.apache.commons.lang3.LocaleUtils; import org.junit.jupiter.api.AfterEach; @@ -19,15 +14,20 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; -import org.junit.jupiter.params.provider.ValueSource; import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.expression.ExpressionEvaluator; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.event.JacksonEvent; +import org.opensearch.dataprepper.model.record.Record; import java.time.Instant; import java.time.LocalDate; -import java.time.LocalTime; import java.time.LocalDateTime; +import java.time.LocalTime; import java.time.ZoneId; import java.time.ZonedDateTime; import java.time.format.DateTimeFormatter; @@ -37,13 +37,15 @@ import java.util.List; import java.util.Locale; import java.util.Map; +import java.util.Random; import java.util.UUID; import java.util.stream.Stream; -import java.util.Random; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.notNullValue; import static org.junit.jupiter.params.provider.Arguments.arguments; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -281,12 +283,15 @@ private static Stream getInputOutputFormats() { Random random = new Random(); long millis = random.nextInt(1000); long nanos = random.nextInt(1000_000_000); + long micros = random.nextInt(1000_000); long epochMillis = epochSeconds * 1000L + millis; long epochNanos = epochSeconds * 1000_000_000L + nanos; + long epochMicros = epochSeconds * 1000_000L + micros; ZonedDateTime zdtSeconds = ZonedDateTime.ofInstant(Instant.ofEpochSecond(epochSeconds), java.time.ZoneId.systemDefault()); ZonedDateTime zdtMillis = ZonedDateTime.ofInstant(Instant.ofEpochMilli(epochMillis), java.time.ZoneId.systemDefault()); ZonedDateTime zdtNanos = ZonedDateTime.ofInstant(Instant.ofEpochSecond(epochSeconds, nanos), java.time.ZoneId.systemDefault()); + ZonedDateTime zdtMicros = ZonedDateTime.ofInstant(Instant.ofEpochSecond(epochSeconds, micros * 1000), java.time.ZoneId.systemDefault()); String testFormat = "yyyy-MMM-dd HH:mm:ss.SSS"; String testNanosFormat = "yyyy-MMM-dd HH:mm:ss.nnnnnnnnnXXX"; String defaultFormat = DateProcessorConfig.DEFAULT_OUTPUT_FORMAT; @@ -303,9 +308,14 @@ private static Stream getInputOutputFormats() { arguments("epoch_nano", epochNanos, "epoch_milli", Long.toString(epochNanos/1000_000)), arguments("epoch_nano", epochNanos, testNanosFormat, zdtNanos.format(DateTimeFormatter.ofPattern(testNanosFormat))), arguments("epoch_nano", epochNanos, defaultFormat, zdtNanos.format(DateTimeFormatter.ofPattern(defaultFormat))), + arguments("epoch_micro", epochMicros, "epoch_second", Long.toString(epochSeconds)), + arguments("epoch_micro", epochMicros, "epoch_milli", Long.toString(epochMicros/1000)), + arguments("epoch_micro", epochMicros, testFormat, zdtMicros.format(DateTimeFormatter.ofPattern(testFormat))), + arguments("epoch_micro", epochMicros, defaultFormat, zdtMicros.format(DateTimeFormatter.ofPattern(defaultFormat))), arguments(testNanosFormat, zdtNanos.format(DateTimeFormatter.ofPattern(testNanosFormat)), "epoch_second", Long.toString(epochSeconds)), arguments(testNanosFormat, zdtNanos.format(DateTimeFormatter.ofPattern(testNanosFormat)), "epoch_milli", Long.toString(epochNanos/1000_000)), arguments(testNanosFormat, zdtNanos.format(DateTimeFormatter.ofPattern(testNanosFormat)), "epoch_nano", Long.toString(epochNanos)), + arguments(testNanosFormat, zdtNanos.format(DateTimeFormatter.ofPattern(testNanosFormat)), "epoch_micro", Long.toString(epochNanos/1000)), arguments(testNanosFormat, zdtNanos.format(DateTimeFormatter.ofPattern(testNanosFormat)), defaultFormat, zdtNanos.format(DateTimeFormatter.ofPattern(defaultFormat))) ); } @@ -514,6 +524,24 @@ void match_without_year_test(String pattern) { verify(dateProcessingMatchSuccessCounter, times(1)).increment(); } + @Test + void date_processor_catches_exceptions_instead_of_throwing() { + when(mockDateProcessorConfig.getDateWhen()).thenReturn(UUID.randomUUID().toString()); + when(expressionEvaluator.evaluateConditional(any(String.class), any(Event.class))) + .thenThrow(RuntimeException.class); + + dateProcessor = createObjectUnderTest(); + + final Record record = buildRecordWithEvent(testData); + final List> processedRecords = (List>) dateProcessor.doExecute(Collections.singletonList(record)); + + assertThat(processedRecords, notNullValue()); + assertThat(processedRecords.size(), equalTo(1)); + assertThat(processedRecords.get(0), notNullValue()); + assertThat(processedRecords.get(0).getData(), notNullValue()); + assertThat(processedRecords.get(0).getData().toMap(), equalTo(record.getData().toMap())); + } + static Record buildRecordWithEvent(final Map data) { return new Record<>(JacksonEvent.builder() .withData(data) diff --git a/data-prepper-plugins/decompress-processor/build.gradle b/data-prepper-plugins/decompress-processor/build.gradle new file mode 100644 index 0000000000..9d67cffc3b --- /dev/null +++ b/data-prepper-plugins/decompress-processor/build.gradle @@ -0,0 +1,13 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +dependencies { + implementation 'commons-io:commons-io:2.15.1' + implementation project(':data-prepper-api') + implementation project(':data-prepper-plugins:common') + implementation 'com.fasterxml.jackson.core:jackson-databind' + implementation 'io.micrometer:micrometer-core' + testImplementation testLibs.mockito.inline +} \ No newline at end of file diff --git a/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressProcessor.java b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressProcessor.java new file mode 100644 index 0000000000..61e7b7e812 --- /dev/null +++ b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressProcessor.java @@ -0,0 +1,114 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.decompress; + +import com.google.common.base.Charsets; +import io.micrometer.core.instrument.Counter; +import org.apache.commons.io.IOUtils; +import org.opensearch.dataprepper.expression.ExpressionEvaluator; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; +import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.plugin.InvalidPluginConfigurationException; +import org.opensearch.dataprepper.model.processor.AbstractProcessor; +import org.opensearch.dataprepper.model.processor.Processor; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.plugins.processor.decompress.exceptions.DecodingException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStream; +import java.util.Collection; + +@DataPrepperPlugin(name = "decompress", pluginType = Processor.class, pluginConfigurationType = DecompressProcessorConfig.class) +public class DecompressProcessor extends AbstractProcessor, Record> { + + private static final Logger LOG = LoggerFactory.getLogger(DecompressProcessor.class); + static final String DECOMPRESSION_PROCESSING_ERRORS = "processingErrors"; + + private final DecompressProcessorConfig decompressProcessorConfig; + private final ExpressionEvaluator expressionEvaluator; + + private final Counter decompressionProcessingErrors; + + @DataPrepperPluginConstructor + public DecompressProcessor(final PluginMetrics pluginMetrics, + final DecompressProcessorConfig decompressProcessorConfig, + final ExpressionEvaluator expressionEvaluator) { + super(pluginMetrics); + this.decompressProcessorConfig = decompressProcessorConfig; + this.expressionEvaluator = expressionEvaluator; + this.decompressionProcessingErrors = pluginMetrics.counter(DECOMPRESSION_PROCESSING_ERRORS); + + if (decompressProcessorConfig.getDecompressWhen() != null + && !expressionEvaluator.isValidExpressionStatement(decompressProcessorConfig.getDecompressWhen())) { + throw new InvalidPluginConfigurationException( + String.format("decompress_when value of %s is not a valid expression statement. " + + "See https://opensearch.org/docs/latest/data-prepper/pipelines/expression-syntax/ for valid expression syntax.", decompressProcessorConfig.getDecompressWhen())); + } + } + + @Override + public Collection> doExecute(final Collection> records) { + for (final Record record : records) { + + try { + if (decompressProcessorConfig.getDecompressWhen() != null && !expressionEvaluator.evaluateConditional(decompressProcessorConfig.getDecompressWhen(), record.getData())) { + continue; + } + + for (final String key : decompressProcessorConfig.getKeys()) { + + final String compressedValue = record.getData().get(key, String.class); + + if (compressedValue == null) { + continue; + } + + final byte[] compressedValueAsBytes = decompressProcessorConfig.getEncodingType().getDecoderEngine().decode(compressedValue); + + try (final InputStream inputStream = decompressProcessorConfig.getDecompressionType().getDecompressionEngine().createInputStream(new ByteArrayInputStream(compressedValueAsBytes));){ + + final String decompressedString = IOUtils.toString(inputStream, Charsets.UTF_8); + record.getData().put(key, decompressedString); + } catch (final Exception e) { + LOG.error("Unable to decompress key {} using decompression type {}:", + key, decompressProcessorConfig.getDecompressionType(), e); + record.getData().getMetadata().addTags(decompressProcessorConfig.getTagsOnFailure()); + decompressionProcessingErrors.increment(); + } + } + } catch (final DecodingException e) { + LOG.error("Unable to decode key with base64: {}", e.getMessage()); + record.getData().getMetadata().addTags(decompressProcessorConfig.getTagsOnFailure()); + decompressionProcessingErrors.increment(); + } catch (final Exception e) { + LOG.error("An uncaught exception occurred while decompressing Events", e); + record.getData().getMetadata().addTags(decompressProcessorConfig.getTagsOnFailure()); + decompressionProcessingErrors.increment(); + } + } + + return records; + } + + @Override + public void prepareForShutdown() { + + } + + @Override + public boolean isReadyForShutdown() { + return true; + } + + @Override + public void shutdown() { + + } +} diff --git a/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressProcessorConfig.java b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressProcessorConfig.java new file mode 100644 index 0000000000..ce2d985277 --- /dev/null +++ b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressProcessorConfig.java @@ -0,0 +1,53 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.decompress; + +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.constraints.NotEmpty; +import jakarta.validation.constraints.NotNull; +import org.opensearch.dataprepper.plugins.processor.decompress.encoding.EncodingType; +import org.opensearch.dataprepper.plugins.processor.decompress.encoding.DecoderEngineFactory; + +import java.util.List; + +public class DecompressProcessorConfig { + + @JsonProperty("keys") + @NotEmpty + private List keys; + + @JsonProperty("type") + @NotNull + private DecompressionType decompressionType; + + @JsonProperty("decompress_when") + private String decompressWhen; + + @JsonProperty("tags_on_failure") + private List tagsOnFailure = List.of("_decompression_failure"); + + @JsonIgnore + private final EncodingType encodingType = EncodingType.BASE64; + + public List getKeys() { + return keys; + } + + public DecompressionEngineFactory getDecompressionType() { + return decompressionType; + } + + public DecoderEngineFactory getEncodingType() { return encodingType; } + + public String getDecompressWhen() { + return decompressWhen; + } + + public List getTagsOnFailure() { + return tagsOnFailure; + } +} diff --git a/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressionEngineFactory.java b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressionEngineFactory.java new file mode 100644 index 0000000000..581b3a1970 --- /dev/null +++ b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressionEngineFactory.java @@ -0,0 +1,12 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.decompress; + +import org.opensearch.dataprepper.model.codec.DecompressionEngine; + +public interface DecompressionEngineFactory { + public DecompressionEngine getDecompressionEngine(); +} diff --git a/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressionType.java b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressionType.java new file mode 100644 index 0000000000..88f64a52e1 --- /dev/null +++ b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressionType.java @@ -0,0 +1,44 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.decompress; + +import com.fasterxml.jackson.annotation.JsonCreator; +import org.opensearch.dataprepper.model.codec.DecompressionEngine; +import org.opensearch.dataprepper.plugins.codec.GZipDecompressionEngine; + +import java.util.Arrays; +import java.util.Map; +import java.util.stream.Collectors; + +public enum DecompressionType implements DecompressionEngineFactory { + GZIP("gzip"); + + private final String option; + + private static final Map OPTIONS_MAP = Arrays.stream(DecompressionType.values()) + .collect(Collectors.toMap( + value -> value.option, + value -> value + )); + + private static final Map DECOMPRESSION_ENGINE_MAP = Map.of( + "gzip", new GZipDecompressionEngine() + ); + + DecompressionType(final String option) { + this.option = option; + } + + @JsonCreator + static DecompressionType fromOptionValue(final String option) { + return OPTIONS_MAP.get(option); + } + + @Override + public DecompressionEngine getDecompressionEngine() { + return DECOMPRESSION_ENGINE_MAP.get(this.option); + } +} diff --git a/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/Base64DecoderEngine.java b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/Base64DecoderEngine.java new file mode 100644 index 0000000000..a6d59d84ed --- /dev/null +++ b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/Base64DecoderEngine.java @@ -0,0 +1,21 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.decompress.encoding; + +import org.opensearch.dataprepper.plugins.processor.decompress.exceptions.DecodingException; + +import java.util.Base64; + +public class Base64DecoderEngine implements DecoderEngine { + @Override + public byte[] decode(final String encodedValue) { + try { + return Base64.getDecoder().decode(encodedValue); + } catch (final Exception e) { + throw new DecodingException(String.format("There was an error decoding with the base64 encoding type: %s", e.getMessage())); + } + } +} diff --git a/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/DecoderEngine.java b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/DecoderEngine.java new file mode 100644 index 0000000000..ef443c273f --- /dev/null +++ b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/DecoderEngine.java @@ -0,0 +1,12 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.decompress.encoding; + +import org.opensearch.dataprepper.plugins.processor.decompress.exceptions.DecodingException; + +public interface DecoderEngine { + byte[] decode(final String encodedValue) throws DecodingException; +} diff --git a/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/DecoderEngineFactory.java b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/DecoderEngineFactory.java new file mode 100644 index 0000000000..89d8a6334c --- /dev/null +++ b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/DecoderEngineFactory.java @@ -0,0 +1,10 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.decompress.encoding; + +public interface DecoderEngineFactory { + DecoderEngine getDecoderEngine(); +} diff --git a/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/EncodingType.java b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/EncodingType.java new file mode 100644 index 0000000000..1e88412682 --- /dev/null +++ b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/EncodingType.java @@ -0,0 +1,42 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.decompress.encoding; + +import com.fasterxml.jackson.annotation.JsonCreator; + +import java.util.Arrays; +import java.util.Map; +import java.util.stream.Collectors; + +public enum EncodingType implements DecoderEngineFactory { + BASE64("base64"); + + private final String option; + + private static final Map OPTIONS_MAP = Arrays.stream(EncodingType.values()) + .collect(Collectors.toMap( + value -> value.option, + value -> value + )); + + private static final Map DECODER_ENGINE_MAP = Map.of( + "base64", new Base64DecoderEngine() + ); + + EncodingType(final String option) { + this.option = option; + } + + @JsonCreator + static EncodingType fromOptionValue(final String option) { + return OPTIONS_MAP.get(option); + } + + @Override + public DecoderEngine getDecoderEngine() { + return DECODER_ENGINE_MAP.get(this.option); + } +} diff --git a/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/exceptions/DecodingException.java b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/exceptions/DecodingException.java new file mode 100644 index 0000000000..f14e67d92c --- /dev/null +++ b/data-prepper-plugins/decompress-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/decompress/exceptions/DecodingException.java @@ -0,0 +1,12 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.decompress.exceptions; + +public class DecodingException extends RuntimeException { + public DecodingException(final String message) { + super(message); + } +} diff --git a/data-prepper-plugins/decompress-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressProcessorIT.java b/data-prepper-plugins/decompress-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressProcessorIT.java new file mode 100644 index 0000000000..f0b6a7da42 --- /dev/null +++ b/data-prepper-plugins/decompress-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressProcessorIT.java @@ -0,0 +1,72 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.decompress; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.expression.ExpressionEvaluator; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.plugins.processor.decompress.encoding.EncodingType; + +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.notNullValue; +import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.plugins.processor.decompress.DecompressProcessorTest.buildRecordWithEvent; + +@ExtendWith(MockitoExtension.class) +public class DecompressProcessorIT { + + private List keys; + + @Mock + private PluginMetrics pluginMetrics; + + @Mock + private ExpressionEvaluator expressionEvaluator; + + @Mock + private DecompressProcessorConfig decompressProcessorConfig; + + private DecompressProcessor createObjectUnderTest() { + return new DecompressProcessor(pluginMetrics, decompressProcessorConfig, expressionEvaluator); + } + + @BeforeEach + void setup() { + keys = List.of(UUID.randomUUID().toString()); + when(decompressProcessorConfig.getKeys()).thenReturn(keys); + } + + @ParameterizedTest + @CsvSource({"H4sIAAAAAAAAAPNIzcnJVyjPL8pJAQBSntaLCwAAAA==,Hello world", + "H4sIAAAAAAAAAwvJyCxWAKJEhYKcxMy8ktSKEoXikqLMvHQAkJ3GfRoAAAA=,This is a plaintext string"}) + void base64_encoded_gzip_is_decompressed_successfully(final String compressedValue, final String expectedDecompressedValue) { + when(decompressProcessorConfig.getEncodingType()).thenReturn(EncodingType.BASE64); + when(decompressProcessorConfig.getDecompressionType()).thenReturn(DecompressionType.GZIP); + + final DecompressProcessor objectUnderTest = createObjectUnderTest(); + final List> records = List.of(buildRecordWithEvent(Map.of(keys.get(0), compressedValue))); + + final List> result = (List>) objectUnderTest.doExecute(records); + + assertThat(result, notNullValue()); + assertThat(result.size(), equalTo(1)); + assertThat(result.get(0), notNullValue()); + assertThat(result.get(0).getData(), notNullValue()); + assertThat(result.get(0).getData().get(keys.get(0), String.class), equalTo(expectedDecompressedValue)); + } +} diff --git a/data-prepper-plugins/decompress-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressProcessorTest.java b/data-prepper-plugins/decompress-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressProcessorTest.java new file mode 100644 index 0000000000..8ad5cfc657 --- /dev/null +++ b/data-prepper-plugins/decompress-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressProcessorTest.java @@ -0,0 +1,223 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.decompress; + +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.Timer; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.expression.ExpressionEvaluator; +import org.opensearch.dataprepper.metrics.MetricNames; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.codec.DecompressionEngine; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.event.JacksonEvent; +import org.opensearch.dataprepper.model.plugin.InvalidPluginConfigurationException; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.plugins.processor.decompress.encoding.DecoderEngine; +import org.opensearch.dataprepper.plugins.processor.decompress.encoding.DecoderEngineFactory; +import org.opensearch.dataprepper.plugins.processor.decompress.exceptions.DecodingException; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.notNullValue; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.plugins.processor.decompress.DecompressProcessor.DECOMPRESSION_PROCESSING_ERRORS; + +@ExtendWith(MockitoExtension.class) +public class DecompressProcessorTest { + + private String key; + + @Mock + private DecompressionEngine decompressionEngine; + + @Mock + private DecoderEngine decoderEngine; + + @Mock + private PluginMetrics pluginMetrics; + + @Mock + private Counter decompressionProcessingErrors; + + @Mock + private ExpressionEvaluator expressionEvaluator; + + @Mock + private DecompressProcessorConfig decompressProcessorConfig; + + @Mock + private DecompressionEngineFactory decompressionType; + + @Mock + private DecoderEngineFactory encodingType; + + private DecompressProcessor createObjectUnderTest() { + return new DecompressProcessor(pluginMetrics, decompressProcessorConfig, expressionEvaluator); + } + + @BeforeEach + void setup() { + key = UUID.randomUUID().toString(); + + when(pluginMetrics.counter(MetricNames.RECORDS_IN)).thenReturn(mock(Counter.class)); + when(pluginMetrics.counter(MetricNames.RECORDS_OUT)).thenReturn(mock(Counter.class)); + when(pluginMetrics.timer(MetricNames.TIME_ELAPSED)).thenReturn(mock(Timer.class)); + when(pluginMetrics.counter(DECOMPRESSION_PROCESSING_ERRORS)).thenReturn(decompressionProcessingErrors); + } + + @Test + void decompression_returns_expected_output() throws IOException { + final String compressedValue = UUID.randomUUID().toString(); + final String expectedResult = UUID.randomUUID().toString(); + final byte[] decodedValue = expectedResult.getBytes(); + + when(decompressProcessorConfig.getKeys()).thenReturn(List.of(key)); + when(encodingType.getDecoderEngine()).thenReturn(decoderEngine); + when(decompressProcessorConfig.getEncodingType()).thenReturn(encodingType); + when(decompressProcessorConfig.getDecompressionType()).thenReturn(decompressionType); + when(decompressionType.getDecompressionEngine()).thenReturn(decompressionEngine); + when(decoderEngine.decode(compressedValue)).thenReturn(decodedValue); + when(decompressionEngine.createInputStream(any(InputStream.class))).thenReturn(new ByteArrayInputStream(decodedValue)); + + final List> records = List.of(buildRecordWithEvent(Map.of(key, compressedValue))); + + final DecompressProcessor objectUnderTest = createObjectUnderTest(); + final List> result = (List>) objectUnderTest.doExecute(records); + + assertThat(result, notNullValue()); + assertThat(result.size(), equalTo(1)); + assertThat(result.get(0), notNullValue()); + assertThat(result.get(0).getData(), notNullValue()); + assertThat(result.get(0).getData().get(key, String.class), equalTo(expectedResult)); + } + + @Test + void decompression_with_decoding_error_adds_tags_and_increments_error_metric() { + final String compressedValue = UUID.randomUUID().toString(); + final String tagForFailure = UUID.randomUUID().toString(); + + when(decompressProcessorConfig.getKeys()).thenReturn(List.of(key)); + when(encodingType.getDecoderEngine()).thenReturn(decoderEngine); + when(decompressProcessorConfig.getEncodingType()).thenReturn(encodingType); + when(decompressProcessorConfig.getTagsOnFailure()).thenReturn(List.of(tagForFailure)); + when(decoderEngine.decode(compressedValue)).thenThrow(DecodingException.class); + + final List> records = List.of(buildRecordWithEvent(Map.of(key, compressedValue))); + + final DecompressProcessor objectUnderTest = createObjectUnderTest(); + final List> result = (List>) objectUnderTest.doExecute(records); + + assertThat(result, notNullValue()); + assertThat(result.size(), equalTo(1)); + assertThat(result.get(0), notNullValue()); + assertThat(result.get(0).getData(), notNullValue()); + assertThat(result.get(0).getData().get(key, String.class), equalTo(compressedValue)); + assertThat(result.get(0).getData().getMetadata().getTags(), notNullValue()); + assertThat(result.get(0).getData().getMetadata().getTags().size(), equalTo(1)); + assertThat(result.get(0).getData().getMetadata().getTags().contains(tagForFailure), equalTo(true)); + + verifyNoInteractions(decompressionEngine); + verify(decompressionProcessingErrors).increment(); + } + + @Test + void exception_from_DecompressionEngine_adds_tags_and_increments_error_metric() throws IOException { + final String compressedValue = UUID.randomUUID().toString(); + final String expectedResult = UUID.randomUUID().toString(); + final byte[] decodedValue = expectedResult.getBytes(); + final String tagForFailure = UUID.randomUUID().toString(); + + when(decompressProcessorConfig.getKeys()).thenReturn(List.of(key)); + when(encodingType.getDecoderEngine()).thenReturn(decoderEngine); + when(decompressProcessorConfig.getEncodingType()).thenReturn(encodingType); + when(decompressProcessorConfig.getTagsOnFailure()).thenReturn(List.of(tagForFailure)); + when(decompressProcessorConfig.getDecompressionType()).thenReturn(decompressionType); + when(decompressionType.getDecompressionEngine()).thenReturn(decompressionEngine); + when(decoderEngine.decode(compressedValue)).thenReturn(decodedValue); + when(decompressionEngine.createInputStream(any(InputStream.class))).thenThrow(RuntimeException.class); + + final List> records = List.of(buildRecordWithEvent(Map.of(key, compressedValue))); + + final DecompressProcessor objectUnderTest = createObjectUnderTest(); + final List> result = (List>) objectUnderTest.doExecute(records); + + assertThat(result, notNullValue()); + assertThat(result.size(), equalTo(1)); + assertThat(result.get(0), notNullValue()); + assertThat(result.get(0).getData(), notNullValue()); + assertThat(result.get(0).getData().get(key, String.class), equalTo(compressedValue)); + assertThat(result.get(0).getData().getMetadata().getTags(), notNullValue()); + assertThat(result.get(0).getData().getMetadata().getTags().size(), equalTo(1)); + assertThat(result.get(0).getData().getMetadata().getTags().contains(tagForFailure), equalTo(true)); + + verify(decompressionProcessingErrors).increment(); + } + + @Test + void exception_from_expression_evaluator_adds_tags_and_increments_error_metric() { + final String decompressWhen = UUID.randomUUID().toString(); + final String compressedValue = UUID.randomUUID().toString(); + final String tagForFailure = UUID.randomUUID().toString(); + + when(decompressProcessorConfig.getTagsOnFailure()).thenReturn(List.of(tagForFailure)); + when(decompressProcessorConfig.getDecompressWhen()).thenReturn(decompressWhen); + when(expressionEvaluator.isValidExpressionStatement(decompressWhen)).thenReturn(true); + when(expressionEvaluator.evaluateConditional(eq(decompressWhen), any(Event.class))) + .thenThrow(RuntimeException.class); + + final List> records = List.of(buildRecordWithEvent(Map.of(key, compressedValue))); + + final DecompressProcessor objectUnderTest = createObjectUnderTest(); + final List> result = (List>) objectUnderTest.doExecute(records); + + assertThat(result, notNullValue()); + assertThat(result.size(), equalTo(1)); + assertThat(result.get(0), notNullValue()); + assertThat(result.get(0).getData(), notNullValue()); + assertThat(result.get(0).getData().get(key, String.class), equalTo(compressedValue)); + assertThat(result.get(0).getData().getMetadata().getTags(), notNullValue()); + assertThat(result.get(0).getData().getMetadata().getTags().size(), equalTo(1)); + assertThat(result.get(0).getData().getMetadata().getTags().contains(tagForFailure), equalTo(true)); + + verifyNoInteractions(decoderEngine, decompressionEngine); + verify(decompressionProcessingErrors).increment(); + } + + @Test + void invalid_expression_statement_throws_InvalidPluginConfigurationException() { + + final String decompressWhen = UUID.randomUUID().toString(); + when(decompressProcessorConfig.getDecompressWhen()).thenReturn(decompressWhen); + when(expressionEvaluator.isValidExpressionStatement(decompressWhen)).thenReturn(false); + + assertThrows(InvalidPluginConfigurationException.class, this::createObjectUnderTest); + } + + static Record buildRecordWithEvent(final Map data) { + return new Record<>(JacksonEvent.builder() + .withData(data) + .withEventType("event") + .build()); + } +} diff --git a/data-prepper-plugins/decompress-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressionTypeTest.java b/data-prepper-plugins/decompress-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressionTypeTest.java new file mode 100644 index 0000000000..287e0cdb1d --- /dev/null +++ b/data-prepper-plugins/decompress-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/decompress/DecompressionTypeTest.java @@ -0,0 +1,54 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.decompress; + +import org.junit.jupiter.api.extension.ExtensionContext; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.ArgumentsProvider; +import org.junit.jupiter.params.provider.ArgumentsSource; +import org.opensearch.dataprepper.model.codec.DecompressionEngine; +import org.opensearch.dataprepper.plugins.codec.GZipDecompressionEngine; + +import java.util.stream.Stream; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.instanceOf; +import static org.junit.jupiter.params.provider.Arguments.arguments; + +public class DecompressionTypeTest { + + @ParameterizedTest + @ArgumentsSource(EnumToStringNameArgumentsProvider.class) + void fromOptionValue_returns_expected_DecompressionType(final DecompressionType expectedEnumValue, final String enumName) { + assertThat(DecompressionType.fromOptionValue(enumName), equalTo(expectedEnumValue)); + } + + @ParameterizedTest + @ArgumentsSource(EnumToDecompressionEngineClassArgumentsProvider.class) + void getDecompressionEngine_returns_expected_DecompressionEngine(final DecompressionType enumValue, final Class decompressionEngineClass) { + assertThat(enumValue.getDecompressionEngine(), instanceOf(decompressionEngineClass)); + } + + private static class EnumToStringNameArgumentsProvider implements ArgumentsProvider { + @Override + public Stream provideArguments(final ExtensionContext context) { + return Stream.of( + arguments(DecompressionType.GZIP, "gzip") + ); + } + } + + private static class EnumToDecompressionEngineClassArgumentsProvider implements ArgumentsProvider { + @Override + public Stream provideArguments(final ExtensionContext context) { + return Stream.of( + arguments(DecompressionType.GZIP, GZipDecompressionEngine.class) + ); + } + } +} diff --git a/data-prepper-plugins/decompress-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/Base64DecoderEngineTest.java b/data-prepper-plugins/decompress-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/Base64DecoderEngineTest.java new file mode 100644 index 0000000000..989a4a067a --- /dev/null +++ b/data-prepper-plugins/decompress-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/Base64DecoderEngineTest.java @@ -0,0 +1,52 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.decompress.encoding; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.mockito.MockedStatic; +import org.opensearch.dataprepper.plugins.processor.decompress.exceptions.DecodingException; + +import java.util.Base64; +import java.util.UUID; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.when; + +public class Base64DecoderEngineTest { + + @ParameterizedTest + @CsvSource(value = {"Hello world,SGVsbG8gd29ybGQ=", "Test123,VGVzdDEyMw=="}) + void decode_correctly_decodes_base64(final String expectedDecodedValue, final String base64EncodedValue) { + final byte[] expectedDecodedBytes = expectedDecodedValue.getBytes(); + + final DecoderEngine objectUnderTest = new Base64DecoderEngine(); + + final byte[] decodedBytes = objectUnderTest.decode(base64EncodedValue); + + assertThat(decodedBytes, equalTo(expectedDecodedBytes)); + } + + @Test + void decode_throws_DecodingException_when_decoding_base64_throws_exception() { + final String encodedValue = UUID.randomUUID().toString(); + final Base64.Decoder decoder = mock(Base64.Decoder.class); + when(decoder.decode(encodedValue)).thenThrow(RuntimeException.class); + + try(final MockedStatic base64MockedStatic = mockStatic(Base64.class)) { + base64MockedStatic.when(Base64::getDecoder).thenReturn(decoder); + + final DecoderEngine objectUnderTest = new Base64DecoderEngine(); + + assertThrows(DecodingException.class, () -> objectUnderTest.decode(encodedValue)); + } + } +} diff --git a/data-prepper-plugins/decompress-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/EncodingTypeTest.java b/data-prepper-plugins/decompress-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/EncodingTypeTest.java new file mode 100644 index 0000000000..8a6075eb6c --- /dev/null +++ b/data-prepper-plugins/decompress-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/decompress/encoding/EncodingTypeTest.java @@ -0,0 +1,52 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.decompress.encoding; + +import org.junit.jupiter.api.extension.ExtensionContext; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.ArgumentsProvider; +import org.junit.jupiter.params.provider.ArgumentsSource; + +import java.util.stream.Stream; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.instanceOf; +import static org.junit.jupiter.params.provider.Arguments.arguments; + +public class EncodingTypeTest { + + @ParameterizedTest + @ArgumentsSource(EnumToStringNameArgumentsProvider.class) + void fromOptionValue_returns_expected_DecompressionType(final EncodingType expectedEnumValue, final String enumName) { + assertThat(EncodingType.fromOptionValue(enumName), equalTo(expectedEnumValue)); + } + + @ParameterizedTest + @ArgumentsSource(EnumToDecoderEngineClassArgumentsProvider.class) + void getDecompressionEngine_returns_expected_DecompressionEngine(final EncodingType enumValue, final Class decoderEngineClass) { + assertThat(enumValue.getDecoderEngine(), instanceOf(decoderEngineClass)); + } + + private static class EnumToStringNameArgumentsProvider implements ArgumentsProvider { + @Override + public Stream provideArguments(final ExtensionContext context) { + return Stream.of( + arguments(EncodingType.BASE64, "base64") + ); + } + } + + private static class EnumToDecoderEngineClassArgumentsProvider implements ArgumentsProvider { + @Override + public Stream provideArguments(final ExtensionContext context) { + return Stream.of( + arguments(EncodingType.BASE64, Base64DecoderEngine.class) + ); + } + } +} diff --git a/data-prepper-plugins/dissect-processor/README.md b/data-prepper-plugins/dissect-processor/README.md index 84bd7c286e..c75d7f3176 100644 --- a/data-prepper-plugins/dissect-processor/README.md +++ b/data-prepper-plugins/dissect-processor/README.md @@ -1,129 +1,5 @@ # Dissect Processor -The Dissect processor is useful when dealing with log files or messages that have a known pattern or structure. It extracts specific pieces of information from the text and map them to individual fields based on the defined Dissect patterns. +The dissect processor extracts values from an event and maps them to individual fields based on user-defined dissect patterns. - -## Basic Usage - -To get started with dissect processor using Data Prepper, create the following `pipeline.yaml`. -```yaml -dissect-pipeline: - source: - file: - path: "/full/path/to/dissect_logs_json.log" - record_type: "event" - format: "json" - processor: - - dissect: - map: - log: "%{Date} %{Time} %{Log_Type}: %{Message}" - sink: - - stdout: -``` - -Create the following file named `dissect_logs_json.log` and replace the `path` in the file source of your `pipeline.yaml` with the path of this file. - -``` -{"log": "07-25-2023 10:00:00 ERROR: Some error"} -``` - -The Dissect processor will retrieve the necessary fields from the `log` message, such as `Date`, `Time`, `Log_Type`, and `Message`, with the help of the pattern `%{Date} %{Time} %{Type}: %{Message}`, configured in the pipeline. - -When you run Data Prepper with this `pipeline.yaml` passed in, you should see the following standard output. - -``` -{ - "log" : "07-25-2023 10:00:00 ERROR: Some error", - "Date" : "07-25-2023" - "Time" : "10:00:00" - "Log_Type" : "ERROR" - "Message" : "Some error" -} -``` - -The fields `Date`, `Time`, `Log_Type`, and `Message` have been extracted from `log` value. - -## Configuration -* `map` (Required): `map` is required to specify the dissect patterns. It takes a `Map` with fields as keys and respective dissect patterns as values. - - -* `target_types` (Optional): A `Map` that specifies what the target type of specific field should be. Valid options are `integer`, `double`, `string`, and `boolean`. By default, all the values are `string`. Target types will be changed after the dissection process. - - -* `dissect_when` (Optional): A Data Prepper Expression string following the [Data Prepper Expression syntax](../../docs/expression_syntax.md). When configured, the processor will evaluate the expression before proceeding with the dissection process and perform the dissection if the expression evaluates to `true`. - -## Field Notations - -Symbols like `?, +, ->, /, &` can be used to perform logical extraction of data. - -* **Normal Field** : The field without a suffix or prefix. The field will be directly added to the output Event. - - Ex: `%{field_name}` - - -* **Skip Field** : ? can be used as a prefix to key to skip that field in the output JSON. - * Skip Field : `%{}` - * Named skip field : `%{?field_name}` - - - - -* **Append Field** : To append multiple values and put the final value in the field, we can use + before the field name in the dissect pattern - * **Usage**: - - Pattern : "%{+field_name}, %{+field_name}" - Text : "foo, bar" - - Output : {"field_name" : "foobar"} - - We can also define the order the concatenation with the help of suffix `/` . - - * **Usage**: - - Pattern : "%{+field_name/2}, %{+field_name/1}" - Text : "foo, bar" - - Output : {"field_name" : "barfoo"} - - If the order is not mentioned, the append operation will take place in the order of fields specified in the dissect pattern.

- -* **Indirect Field** : While defining a pattern, prefix the field with a `&` to assign the value found with this field to the value of another field found as the key. - * **Usage**: - - Pattern : "%{?field_name}, %{&field_name}" - Text: "foo, bar" - - Output : {“foo” : “bar”} - - Here we can see that `foo` which was captured from the skip field `%{?field_name}` is made the key to value captured form the field `%{&field_name}` - * **Usage**: - - Pattern : %{field_name}, %{&field_name} - Text: "foo, bar" - - Output : {“field_name”:“foo”, “foo”:“bar”} - - We can also indirectly assign the value to an appended field, along with `normal` field and `skip` field. - -### Padding - -* `->` operator can be used as a suffix to a field to indicate that white spaces after this field can be ignored. - * **Usage**: - - Pattern : %{field1→} %{field2} - Text : “firstname lastname” - - Output : {“field1” : “firstname”, “field2” : “lastname”} - -* This operator should be used as the right most suffix. - * **Usage**: - - Pattern : %{fieldname/1->} %{fieldname/2} - - If we use `->` before `/`, the `->` operator will also be considered part of the field name. - - -## Developer Guide -This plugin is compatible with Java 14. See -- [CONTRIBUTING](https://github.com/opensearch-project/data-prepper/blob/main/CONTRIBUTING.md) -- [monitoring](https://github.com/opensearch-project/data-prepper/blob/main/docs/monitoring.md) +See the [`dissect` processor documentation](https://opensearch.org/docs/latest/data-prepper/pipelines/configuration/processors/dissect/). diff --git a/data-prepper-plugins/dissect-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/dissect/DissectProcessor.java b/data-prepper-plugins/dissect-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/dissect/DissectProcessor.java index 5ff7f4ad56..0977d998a3 100644 --- a/data-prepper-plugins/dissect-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/dissect/DissectProcessor.java +++ b/data-prepper-plugins/dissect-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/dissect/DissectProcessor.java @@ -55,18 +55,17 @@ public DissectProcessor(PluginMetrics pluginMetrics, final DissectProcessorConfi public Collection> doExecute(Collection> records) { for (final Record record : records) { Event event = record.getData(); - String dissectWhen = dissectConfig.getDissectWhen(); - if (Objects.nonNull(dissectWhen) && !expressionEvaluator.evaluateConditional(dissectWhen, event)) { - continue; - } try{ - for(String field: dissectorMap.keySet()){ + String dissectWhen = dissectConfig.getDissectWhen(); + if (Objects.nonNull(dissectWhen) && !expressionEvaluator.evaluateConditional(dissectWhen, event)) { + continue; + } + for (String field: dissectorMap.keySet()){ if(event.containsKey(field)){ dissectField(event, field); } } - } - catch (Exception ex){ + } catch (Exception ex){ LOG.error(EVENT, "Error dissecting the event [{}] ", record.getData(), ex); } } diff --git a/data-prepper-plugins/dynamodb-source/README.md b/data-prepper-plugins/dynamodb-source/README.md index db620a6406..7b3369fc5e 100644 --- a/data-prepper-plugins/dynamodb-source/README.md +++ b/data-prepper-plugins/dynamodb-source/README.md @@ -1,85 +1,5 @@ # DynamoDB Source -This is a source plugin that supports retrieve data from DynamoDB tables. Basic use case of this source plugin is to -sync the data from DynamoDB tables to OpenSearch indexes. With this CDC support, customer can run the end to end data -sync pipeline and capture changed data in near real-time without writing any codes and without any downtime of business. -Such pipeline can run on multiple nodes in parallel to support data capture of large scale tables. +This source ingests data to Data Prepper from DynamoDB -This plugin can support below three different modes: - -1. Full load only: One time full data export and load -2. CDC Only: DynamoDB Stream -3. Full load + CDC: One time full export and load + DynamoDB Stream. - -## Usages - -To get started with this DynamoDB source, create the following source configuration: - -```yaml -source: - dynamodb: - tables: - - table_arn: "arn:aws:dynamodb:us-west-2:123456789012:table/my-table" - stream: - start_position: - export: - s3_bucket: "my-bucket" - s3_prefix: "export/" - aws: - region: "us-west-2" - sts_role_arn: "arn:aws:iam::123456789012:role/DataPrepperRole" - - coordinator: - dynamodb: - table_name: "coordinator-demo" - region: "us-west-2" - - -``` - -## Configurations - -### Shared Configurations: - -* coordinator (Required): Coordination store setting. This design create a custom coordinator based on existing - coordination store implementation. Only DynamoDB is tested so far. -* aws (Required): High level AWS Auth. Note Data Prepper will use the same AWS auth to access all tables, check - Security for more details. - * region - * sts_role_arn - -### Export Configurations: - -* s3_bucket (Required): The destination bucket to store the exported data files -* s3_prefix (Optional): Custom prefix. -* s3_sse_kms_key_id (Optional): A AWS KMS Customer Managed Key (CMK) to encrypt the export data files. The key id will - be the ARN of the Key, e.g. arn:aws:kms:us-west-2:123456789012:key/0a4bc22f-bb96-4ad4-80ca-63b12b3ec147 - -### Stream Configurations - -* start_position (Optional): start position of the stream, can be either TRIM_HORIZON or LATEST. If export is required, - this value will be ignored and set to LATEST by default. This is useful if customer don’t want to run initial export, - so they can - choose either from the beginning of the stream (up to 24 hours) or from latest (from the time point when pipeline is - started) - -## Metrics - -### Counter - -- `exportJobsSuccess`: measures total number of export jobs run with status completed. -- `exportJobsErrors`: measures total number of export jobs cannot be submitted or run with status failed. -- `exportFilesTotal`: measures total number of export files generated. -- `exportFilesSuccess`: measures total number of export files read (till the last line) successfully. -- `exportRecordsTotal`: measures total number of export records generated -- `exportRecordsSuccess`: measures total number of export records processed successfully . -- `exportRecordsErrors`: measures total number of export records processed failed -- `changeEventsSucceeded`: measures total number of changed events in total processed successfully -- `changeEventsFailed`: measures total number of changed events in total processed failed - -## Developer Guide - -This plugin is compatible with Java 17. See - -- [CONTRIBUTING](https://github.com/opensearch-project/data-prepper/blob/main/CONTRIBUTING.md) -- [monitoring](https://github.com/opensearch-project/data-prepper/blob/main/docs/monitoring.md) +See the [`dynamodb` source documentation](https://opensearch.org/docs/latest/data-prepper/pipelines/configuration/sources/dynamo-db/) diff --git a/data-prepper-plugins/dynamodb-source/build.gradle b/data-prepper-plugins/dynamodb-source/build.gradle index 30a84c2733..8c25a0482b 100644 --- a/data-prepper-plugins/dynamodb-source/build.gradle +++ b/data-prepper-plugins/dynamodb-source/build.gradle @@ -3,10 +3,6 @@ plugins { } -repositories { - mavenCentral() -} - dependencies { implementation project(path: ':data-prepper-api') diff --git a/data-prepper-plugins/flatten-processor/build.gradle b/data-prepper-plugins/flatten-processor/build.gradle new file mode 100644 index 0000000000..42a3a74cf5 --- /dev/null +++ b/data-prepper-plugins/flatten-processor/build.gradle @@ -0,0 +1,22 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +jacocoTestCoverageVerification { + dependsOn jacocoTestReport + violationRules { + rule { + limit { + minimum = 0.9 + } + } + } +} + +dependencies { + implementation project(':data-prepper-api') + implementation project(':data-prepper-plugins:common') + implementation 'com.fasterxml.jackson.core:jackson-databind' + implementation 'com.github.wnameless.json:json-flattener:0.16.6' +} \ No newline at end of file diff --git a/data-prepper-plugins/flatten-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/flatten/FlattenProcessor.java b/data-prepper-plugins/flatten-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/flatten/FlattenProcessor.java new file mode 100644 index 0000000000..a0786f76c1 --- /dev/null +++ b/data-prepper-plugins/flatten-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/flatten/FlattenProcessor.java @@ -0,0 +1,138 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.flatten; + +import com.github.wnameless.json.flattener.JsonFlattener; +import org.opensearch.dataprepper.expression.ExpressionEvaluator; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; +import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.processor.AbstractProcessor; +import org.opensearch.dataprepper.model.processor.Processor; +import org.opensearch.dataprepper.model.record.Record; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +@DataPrepperPlugin(name = "flatten", pluginType = Processor.class, pluginConfigurationType = FlattenProcessorConfig.class) +public class FlattenProcessor extends AbstractProcessor, Record> { + private static final Logger LOG = LoggerFactory.getLogger(FlattenProcessor.class); + + private static final String SEPARATOR = "/"; + private final FlattenProcessorConfig config; + private final ExpressionEvaluator expressionEvaluator; + + @DataPrepperPluginConstructor + public FlattenProcessor(final PluginMetrics pluginMetrics, final FlattenProcessorConfig config, final ExpressionEvaluator expressionEvaluator) { + super(pluginMetrics); + this.config = config; + this.expressionEvaluator = expressionEvaluator; + } + + @Override + public Collection> doExecute(final Collection> records) { + for (final Record record : records) { + final Event recordEvent = record.getData(); + + try { + if (config.getFlattenWhen() != null && !expressionEvaluator.evaluateConditional(config.getFlattenWhen(), recordEvent)) { + continue; + } + + final String sourceJson = recordEvent.getAsJsonString(config.getSource()); + + // adds ignoreReservedCharacters() so that dots in keys are ignored during flattening + // e.g., {"a.b": {"c": 1}} will be flattened as expected: {"a.b.c": 1}; otherwise, flattened result will be {"[\"a.b\"]c": 1} + Map flattenedJson = new JsonFlattener(sourceJson).ignoreReservedCharacters().flattenAsMap(); + + if (config.isRemoveProcessedFields()) { + final Map sourceMap = recordEvent.get(config.getSource(), Map.class); + for (final String keyInSource : sourceMap.keySet()) { + recordEvent.delete(getJsonPointer(config.getSource(), keyInSource)); + } + } + + if (config.isRemoveListIndices()) { + flattenedJson = removeListIndicesInKeys(flattenedJson); + } + + updateEvent(recordEvent, flattenedJson); + } catch (Exception e) { + LOG.error("Fail to perform flatten operation", e); + recordEvent.getMetadata().addTags(config.getTagsOnFailure()); + } + } + return records; + } + + @Override + public void prepareForShutdown() { + } + + @Override + public boolean isReadyForShutdown() { + return true; + } + + @Override + public void shutdown() { + } + + private String getJsonPointer(final String outerKey, final String innerKey) { + if (outerKey.isEmpty()) { + return SEPARATOR + innerKey; + } else { + return SEPARATOR + outerKey + SEPARATOR + innerKey; + } + } + + private Map removeListIndicesInKeys(final Map inputMap) { + final Map resultMap = new HashMap<>(); + + for (final Map.Entry entry : inputMap.entrySet()) { + final String keyWithoutIndices = removeListIndices(entry.getKey()); + addFieldsToMapWithMerge(keyWithoutIndices, entry.getValue(), resultMap); + } + return resultMap; + } + + private String removeListIndices(final String key) { + return key.replaceAll("\\[\\d+\\]", "[]"); + } + + private void addFieldsToMapWithMerge(String key, Object value, Map map) { + if (!map.containsKey(key)) { + map.put(key, value); + } else { + final Object currentValue = map.get(key); + if (currentValue instanceof List) { + ((List)currentValue).add(value); + } else { + List newValue = new ArrayList<>(); + newValue.add(currentValue); + newValue.add(value); + map.put(key, newValue); + } + } + } + + private void updateEvent(Event recordEvent, Map flattenedJson) { + if (config.getTarget().isEmpty()) { + // Target is root + for (final Map.Entry entry : flattenedJson.entrySet()) { + recordEvent.put(entry.getKey(), entry.getValue()); + } + } else { + recordEvent.put(config.getTarget(), flattenedJson); + } + } +} diff --git a/data-prepper-plugins/flatten-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/flatten/FlattenProcessorConfig.java b/data-prepper-plugins/flatten-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/flatten/FlattenProcessorConfig.java new file mode 100644 index 0000000000..7b2622b4f5 --- /dev/null +++ b/data-prepper-plugins/flatten-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/flatten/FlattenProcessorConfig.java @@ -0,0 +1,58 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.flatten; + +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.constraints.NotNull; + +import java.util.List; + +public class FlattenProcessorConfig { + + @NotNull + @JsonProperty("source") + private String source; + + @NotNull + @JsonProperty("target") + private String target; + + @JsonProperty("remove_processed_fields") + private boolean removeProcessedFields = false; + + @JsonProperty("remove_list_indices") + private boolean removeListIndices = false; + + @JsonProperty("flatten_when") + private String flattenWhen; + + @JsonProperty("tags_on_failure") + private List tagsOnFailure; + + public String getSource() { + return source; + } + + public String getTarget() { + return target; + } + + public boolean isRemoveProcessedFields() { + return removeProcessedFields; + } + + public boolean isRemoveListIndices() { + return removeListIndices; + } + + public String getFlattenWhen() { + return flattenWhen; + } + + public List getTagsOnFailure() { + return tagsOnFailure; + } +} diff --git a/data-prepper-plugins/flatten-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/flatten/FlattenProcessorConfigTest.java b/data-prepper-plugins/flatten-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/flatten/FlattenProcessorConfigTest.java new file mode 100644 index 0000000000..1b32c789e2 --- /dev/null +++ b/data-prepper-plugins/flatten-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/flatten/FlattenProcessorConfigTest.java @@ -0,0 +1,25 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.flatten; + +import org.junit.jupiter.api.Test; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +class FlattenProcessorConfigTest { + @Test + void testDefaultConfig() { + final FlattenProcessorConfig FlattenProcessorConfig = new FlattenProcessorConfig(); + + assertThat(FlattenProcessorConfig.getSource(), equalTo(null)); + assertThat(FlattenProcessorConfig.getTarget(), equalTo(null)); + assertThat(FlattenProcessorConfig.isRemoveListIndices(), equalTo(false)); + assertThat(FlattenProcessorConfig.isRemoveListIndices(), equalTo(false)); + assertThat(FlattenProcessorConfig.getFlattenWhen(), equalTo(null)); + assertThat(FlattenProcessorConfig.getTagsOnFailure(), equalTo(null)); + } +} diff --git a/data-prepper-plugins/flatten-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/flatten/FlattenProcessorTest.java b/data-prepper-plugins/flatten-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/flatten/FlattenProcessorTest.java new file mode 100644 index 0000000000..d8e3309c1b --- /dev/null +++ b/data-prepper-plugins/flatten-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/flatten/FlattenProcessorTest.java @@ -0,0 +1,291 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.flatten; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.expression.ExpressionEvaluator; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.event.JacksonEvent; +import org.opensearch.dataprepper.model.record.Record; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +class FlattenProcessorTest { + private static final String SOURCE_KEY = "source"; + private static final String TARGET_KEY = "target"; + + @Mock + private PluginMetrics pluginMetrics; + + @Mock + private FlattenProcessorConfig mockConfig; + + @Mock + private ExpressionEvaluator expressionEvaluator; + + @BeforeEach + void setUp() { + lenient().when(mockConfig.getSource()).thenReturn(""); + lenient().when(mockConfig.getTarget()).thenReturn(""); + lenient().when(mockConfig.isRemoveProcessedFields()).thenReturn(false); + lenient().when(mockConfig.isRemoveListIndices()).thenReturn(false); + lenient().when(mockConfig.getFlattenWhen()).thenReturn(null); + lenient().when(mockConfig.getTagsOnFailure()).thenReturn(new ArrayList<>()); + } + + @Test + void testFlattenEntireEventData() { + final FlattenProcessor processor = createObjectUnderTest(); + final Record testRecord = createTestRecord(createTestData()); + final List> resultRecord = (List>) processor.doExecute(Collections.singletonList(testRecord)); + + assertThat(resultRecord.size(), is(1)); + + final Event resultEvent = resultRecord.get(0).getData(); + Map resultData = resultEvent.get("", Map.class); + + assertFlattenedData(resultData); + + assertThat(resultData.containsKey("key2"), is(true)); + assertThat(resultData.containsKey("list1"), is(true)); + } + + @Test + void testFlattenEntireEventDataAndRemoveProcessedFields() { + when(mockConfig.isRemoveProcessedFields()).thenReturn(true); + + final FlattenProcessor processor = createObjectUnderTest(); + final Record testRecord = createTestRecord(createTestData()); + final List> resultRecord = (List>) processor.doExecute(Collections.singletonList(testRecord)); + + assertThat(resultRecord.size(), is(1)); + + final Event resultEvent = resultRecord.get(0).getData(); + Map resultData = resultEvent.get("", Map.class); + + assertFlattenedData(resultData); + + assertThat(resultData.containsKey("key2"), is(false)); + assertThat(resultData.containsKey("list1"), is(false)); + } + + @Test + void testFlattenEntireEventDataAndRemoveListIndices() { + when(mockConfig.isRemoveListIndices()).thenReturn(true); + + final FlattenProcessor processor = createObjectUnderTest(); + final Record testRecord = createTestRecord(createTestData()); + final List> resultRecord = (List>) processor.doExecute(Collections.singletonList(testRecord)); + + assertThat(resultRecord.size(), is(1)); + + final Event resultEvent = resultRecord.get(0).getData(); + Map resultData = resultEvent.get("", Map.class); + + assertThat(resultData.containsKey("key1"), is(true)); + assertThat(resultData.get("key1"), is("val1")); + + assertThat(resultData.containsKey("key1"), is(true)); + assertThat(resultData.get("key2.key3.key.4"), is("val2")); + + assertThat(resultData.containsKey("list1[].list2[].name"), is(true)); + assertThat(resultData.get("list1[].list2[].name"), is(List.of("name1", "name2"))); + + assertThat(resultData.containsKey("list1[].list2[].value"), is(true)); + assertThat(resultData.get("list1[].list2[].value"), is(List.of("value1", "value2"))); + } + + @Test + void testFlattenWithSpecificFieldsAsSourceAndTarget() { + when(mockConfig.getSource()).thenReturn(SOURCE_KEY); + when(mockConfig.getTarget()).thenReturn(TARGET_KEY); + + final FlattenProcessor processor = createObjectUnderTest(); + final Record testRecord = createTestRecord(Map.of(SOURCE_KEY, createTestData())); + final List> resultRecord = (List>) processor.doExecute(Collections.singletonList(testRecord)); + + assertThat(resultRecord.size(), is(1)); + + final Event resultEvent = resultRecord.get(0).getData(); + Map resultData = resultEvent.get(TARGET_KEY, Map.class); + + assertFlattenedData(resultData); + } + + @Test + void testFlattenWithSpecificFieldsAsSourceAndTargetAndRemoveProcessedFields() { + when(mockConfig.getSource()).thenReturn(SOURCE_KEY); + when(mockConfig.getTarget()).thenReturn(TARGET_KEY); + when(mockConfig.isRemoveProcessedFields()).thenReturn(true); + + final FlattenProcessor processor = createObjectUnderTest(); + final Record testRecord = createTestRecord(Map.of(SOURCE_KEY, createTestData())); + final List> resultRecord = (List>) processor.doExecute(Collections.singletonList(testRecord)); + + assertThat(resultRecord.size(), is(1)); + + final Event resultEvent = resultRecord.get(0).getData(); + Map resultData = resultEvent.get(TARGET_KEY, Map.class); + + assertFlattenedData(resultData); + + Map sourceData = resultEvent.get(SOURCE_KEY, Map.class); + assertThat(sourceData.containsKey("key1"), is(false)); + assertThat(sourceData.containsKey("key2"), is(false)); + assertThat(sourceData.containsKey("list1"), is(false)); + } + + @Test + void testFlattenWithSpecificFieldsAsSourceAndTargetAndRemoveListIndices() { + when(mockConfig.getSource()).thenReturn(SOURCE_KEY); + when(mockConfig.getTarget()).thenReturn(TARGET_KEY); + when(mockConfig.isRemoveListIndices()).thenReturn(true); + + final FlattenProcessor processor = createObjectUnderTest(); + final Record testRecord = createTestRecord(Map.of(SOURCE_KEY, createTestData())); + final List> resultRecord = (List>) processor.doExecute(Collections.singletonList(testRecord)); + + assertThat(resultRecord.size(), is(1)); + + final Event resultEvent = resultRecord.get(0).getData(); + Map resultData = resultEvent.get(TARGET_KEY, Map.class); + + assertThat(resultData.containsKey("key1"), is(true)); + assertThat(resultData.get("key1"), is("val1")); + + assertThat(resultData.containsKey("key1"), is(true)); + assertThat(resultData.get("key2.key3.key.4"), is("val2")); + + assertThat(resultData.containsKey("list1[].list2[].name"), is(true)); + assertThat(resultData.get("list1[].list2[].name"), is(List.of("name1", "name2"))); + + assertThat(resultData.containsKey("list1[].list2[].value"), is(true)); + assertThat(resultData.get("list1[].list2[].value"), is(List.of("value1", "value2"))); + } + + @Test + public void testEventNotProcessedWhenTheWhenConditionIsFalse() { + final String whenCondition = UUID.randomUUID().toString(); + when(mockConfig.getFlattenWhen()).thenReturn(whenCondition); + + final FlattenProcessor processor = createObjectUnderTest(); + final Record testRecord = createTestRecord(createTestData()); + when(expressionEvaluator.evaluateConditional(whenCondition, testRecord.getData())).thenReturn(false); + final List> resultRecord = (List>) processor.doExecute(Collections.singletonList(testRecord)); + + assertThat(resultRecord.size(), is(1)); + + final Event resultEvent = resultRecord.get(0).getData(); + Map resultData = resultEvent.get("", Map.class); + + assertThat(resultData.containsKey("key2.key3.key.4"), is(false)); + assertThat(resultData.containsKey("list1[0].list2[0].name"), is(false)); + assertThat(resultData.containsKey("list1[0].list2[0].value"), is(false)); + assertThat(resultData.containsKey("list1[0].list2[1].name"), is(false)); + assertThat(resultData.containsKey("list1[0].list2[1].value"), is(false)); + } + + @Test + void testFailureTagsAreAddedWhenException() { + // non-existing source key + when(mockConfig.getSource()).thenReturn("my-other-map"); + final List testTags = List.of("tag1", "tag2"); + when(mockConfig.getTagsOnFailure()).thenReturn(testTags); + + final FlattenProcessor processor = createObjectUnderTest(); + final Record testRecord = createTestRecord(createTestData()); + final List> resultRecord = (List>) processor.doExecute(Collections.singletonList(testRecord)); + + assertThat(resultRecord.size(), is(1)); + + final Event resultEvent = resultRecord.get(0).getData(); + assertThat(resultEvent.getMetadata().getTags(), is(new HashSet<>(testTags))); + } + + private FlattenProcessor createObjectUnderTest() { + return new FlattenProcessor(pluginMetrics, mockConfig, expressionEvaluator); + } + + private Map createTestData() { + // Json data: + // { + // "key1": "val1", + // "key2": { + // "key3": { + // "key.4": "val2" + // } + // }, + // "list1": [ + // { + // "list2": [ + // { + // "name": "name1", + // "value": "value1" + // }, + // { + // "name": "name2", + // "value": "value2" + // } + // ] + // } + // ] + //} + return Map.of( + "key1", "val1", + "key2", Map.of("key3", Map.of("key.4", "val2")), + "list1", List.of( + Map.of("list2", List.of( + Map.of("name", "name1", "value", "value1"), + Map.of("name", "name2", "value", "value2")) + ) + ) + ); + } + + private Record createTestRecord(Object data) { + final Event event = JacksonEvent.builder() + .withData(data) + .withEventType("event") + .build(); + return new Record<>(event); + } + + private void assertFlattenedData(Map resultData) { + assertThat(resultData.containsKey("key1"), is(true)); + assertThat(resultData.get("key1"), is("val1")); + + assertThat(resultData.containsKey("key2.key3.key.4"), is(true)); + assertThat(resultData.get("key2.key3.key.4"), is("val2")); + + assertThat(resultData.containsKey("list1[0].list2[0].name"), is(true)); + assertThat(resultData.get("list1[0].list2[0].name"), is("name1")); + + assertThat(resultData.containsKey("list1[0].list2[0].value"), is(true)); + assertThat(resultData.get("list1[0].list2[0].value"), is("value1")); + + assertThat(resultData.containsKey("list1[0].list2[1].name"), is(true)); + assertThat(resultData.get("list1[0].list2[1].name"), is("name2")); + + assertThat(resultData.containsKey("list1[0].list2[1].value"), is(true)); + assertThat(resultData.get("list1[0].list2[1].value"), is("value2")); + } +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/build.gradle b/data-prepper-plugins/geoip-processor/build.gradle index 916f78e2f1..6bd3788e93 100644 --- a/data-prepper-plugins/geoip-processor/build.gradle +++ b/data-prepper-plugins/geoip-processor/build.gradle @@ -1,15 +1,15 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + plugins{ - id 'de.undercouch.download' version '4.1.2' + id 'de.undercouch.download' version '5.5.0' } apply plugin: 'de.undercouch.download' import de.undercouch.gradle.tasks.download.Download -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - dependencies { implementation project(':data-prepper-api') implementation project(path: ':data-prepper-plugins:common') @@ -19,90 +19,65 @@ dependencies { implementation libs.commons.compress implementation 'org.mapdb:mapdb:3.0.8' implementation libs.commons.io - implementation 'software.amazon.awssdk:aws-sdk-java:2.20.67' + implementation 'software.amazon.awssdk:sts' implementation 'software.amazon.awssdk:s3-transfer-manager' - implementation 'software.amazon.awssdk.crt:aws-crt:0.21.17' + implementation 'software.amazon.awssdk.crt:aws-crt:0.29.9' implementation 'com.maxmind.geoip2:geoip2:4.0.1' implementation 'com.maxmind.db:maxmind-db:3.0.0' implementation 'org.hibernate.validator:hibernate-validator:8.0.1.Final' - implementation libs.commons.lang3 + implementation libs.armeria.core + testImplementation project(':data-prepper-core') testImplementation 'com.fasterxml.jackson.datatype:jackson-datatype-jsr310' testImplementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml' testImplementation project(':data-prepper-test-common') } -def geoIP2='GeoIP2' -def geoLite2= 'GeoLite2' -task downloadFile(type: Download) { - - def urls = [ - 'https://raw.githubusercontent.com/maxmind/MaxMind-DB/main/test-data/GeoIP2-City-Test.mmdb', - 'https://raw.githubusercontent.com/maxmind/MaxMind-DB/main/test-data/GeoIP2-Country-Test.mmdb', - 'https://raw.githubusercontent.com/maxmind/MaxMind-DB/main/test-data/GeoLite2-ASN-Test.mmdb' - ] - def mmdbFileExtension = '.mmdb' - def baseDirPath = 'src/test/resources/mmdb-file/geo-lite2/' - - urls.each { url -> - src(url) - dest(baseDirPath) - doLast { - def testFileName = url.substring(url.lastIndexOf('/') + 1) - def testMmdbSubString = testFileName.substring(testFileName.lastIndexOf('-')) - def fileName = testFileName.substring(0, testFileName.length() - testMmdbSubString.length()) +def downloadFiles = tasks.register('downloadFiles') - if(fileName.contains(geoIP2)) { - fileName = fileName.replace(geoIP2, geoLite2) - } - File sourceFile = file(baseDirPath+testFileName) - File destinationFile = file( baseDirPath+fileName+mmdbFileExtension) - sourceFile.renameTo(destinationFile) +def databaseNames = [ + 'GeoLite2-City-Test', + 'GeoLite2-Country-Test', + 'GeoLite2-ASN-Test' +] - } +databaseNames.forEach { databaseName -> { + def url = "https://raw.githubusercontent.com/maxmind/MaxMind-DB/main/test-data/${databaseName}.mmdb" + def gradleName = databaseName.replaceAll('-', '') + def downloadTask = tasks.register("download${gradleName}", Download) { + src(url) + dest "build/resources/test/mmdb-files/geo-lite2/${databaseName}.mmdb" + overwrite true } + downloadFiles.get().dependsOn downloadTask +}} +def enterpriseDatabaseNames = [ + 'GeoIP2-Enterprise-Test' +] -} -task downloadEnterpriseFile(type: Download) { - dependsOn downloadFile - def urls = [ - 'https://raw.githubusercontent.com/maxmind/MaxMind-DB/main/test-data/GeoIP2-Enterprise-Test.mmdb' - ] - def mmdbFileExtension = '.mmdb' - def baseDirPath = 'src/test/resources/mmdb-file/geo-enterprise/' - - urls.each { url -> - src(url) - def testFileName = url.substring(url.lastIndexOf('/') + 1) - def testMmdbSubString = testFileName.substring(testFileName.lastIndexOf('-')) - def fileName = testFileName.substring(0, testFileName.length() - testMmdbSubString.length()) - - dest(baseDirPath+testFileName) - doLast { - if(fileName.contains(geoIP2)) { - fileName = fileName.replace(geoIP2, geoLite2) - } - File sourceFile = file(baseDirPath+testFileName) - File destinationFile = file( baseDirPath+fileName+mmdbFileExtension) - sourceFile.renameTo(destinationFile) - } +enterpriseDatabaseNames.forEach { enterpriseDatabaseName -> { + def url = "https://raw.githubusercontent.com/maxmind/MaxMind-DB/main/test-data/${enterpriseDatabaseName}.mmdb" + def gradleName = enterpriseDatabaseName.replaceAll('-', '') + def downloadEnterpriseTask = tasks.register("download${gradleName}", Download) { + src(url) + dest "build/resources/test/mmdb-files/geo-ip2/${enterpriseDatabaseName}.mmdb" + overwrite true } + downloadFiles.get().dependsOn downloadEnterpriseTask +}} -} -/*task processTestResources(type: Copy) { - dependsOn downloadEnterpriseFile - from 'src/test/resources' // Source directory containing test resources - into 'build/resources/test' // Destination directory for processed test resources -}*/ -tasks.test.dependsOn 'processTestResources' -tasks.processTestResources.dependsOn 'downloadEnterpriseFile' test { useJUnitPlatform() + dependsOn(downloadFiles) +} + +checkstyleTest { + dependsOn(downloadFiles) } jacocoTestCoverageVerification { @@ -110,40 +85,10 @@ jacocoTestCoverageVerification { violationRules { rule { limit { - minimum = 0.1 // temporarily reduce coverage for the builds to pass + minimum = 0.84 } } } } -check.dependsOn jacocoTestCoverageVerification - -sourceSets { - integrationTest { - java { - compileClasspath += main.output + test.output - runtimeClasspath += main.output + test.output - srcDir file('src/integrationTest/java') - } - resources.srcDir file('src/integrationTest/resources') - } -} - -configurations { - integrationTestImplementation.extendsFrom testImplementation - integrationTestRuntime.extendsFrom testRuntime -} - -task integrationTest(type: Test) { - group = 'verification' - testClassesDirs = sourceSets.integrationTest.output.classesDirs - - useJUnitPlatform() - - classpath = sourceSets.integrationTest.runtimeClasspath - systemProperty 'tests.geoipProcessor.maxmindLicenseKey', System.getProperty('tests.geoipProcessor.maxmindLicenseKey') - - filter { - includeTestsMatching '*IT' - } -} \ No newline at end of file +check.dependsOn jacocoTestCoverageVerification \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/integrationTest/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorUrlServiceIT.java b/data-prepper-plugins/geoip-processor/src/integrationTest/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorUrlServiceIT.java index 70a80c0684..426686e099 100644 --- a/data-prepper-plugins/geoip-processor/src/integrationTest/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorUrlServiceIT.java +++ b/data-prepper-plugins/geoip-processor/src/integrationTest/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorUrlServiceIT.java @@ -81,7 +81,7 @@ public void setUp() throws JsonProcessingException { public GeoIPProcessorService createObjectUnderTest() { // TODO: pass in geoIpServiceConfig object - return new GeoIPProcessorService(null); + return new GeoIPProcessorService(null, null, null); } @Test @@ -93,7 +93,7 @@ void verify_enrichment_of_data_from_maxmind_url() throws UnknownHostException { if (IPValidationCheck.isPublicIpAddress(ipAddress)) { InetAddress inetAddress = InetAddress.getByName(ipAddress); //All attributes are considered by default with the null value - geoData = geoIPProcessorService.getGeoData(inetAddress, null); +// geoData = geoIPProcessorService.getGeoData(inetAddress, null); assertThat(geoData.get("country_iso_code"), equalTo("US")); assertThat(geoData.get("ip"), equalTo("8.8.8.8")); diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPDatabase.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPDatabase.java new file mode 100644 index 0000000000..ff22e138b0 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPDatabase.java @@ -0,0 +1,13 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor; + +public enum GeoIPDatabase { + CITY, + COUNTRY, + ASN, + ENTERPRISE; +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPField.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPField.java new file mode 100644 index 0000000000..5476b7ea88 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPField.java @@ -0,0 +1,72 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor; + +import java.util.Arrays; +import java.util.HashSet; +import java.util.Set; + +public enum GeoIPField { + CONTINENT_CODE("continent_code", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + CONTINENT_NAME("continent_name", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + COUNTRY_NAME("country_name", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + IS_COUNTRY_IN_EUROPEAN_UNION("is_country_in_european_union", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + COUNTRY_ISO_CODE("country_iso_code", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + COUNTRY_CONFIDENCE("country_confidence", GeoIPDatabase.ENTERPRISE), + REGISTERED_COUNTRY_NAME("registered_country_name", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + REGISTERED_COUNTRY_ISO_CODE("registered_country_iso_code", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + REPRESENTED_COUNTRY_NAME("represented_country_name", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + REPRESENTED_COUNTRY_ISO_CODE("represented_country_iso_code", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + REPRESENTED_COUNTRY_TYPE("represented_country_type", GeoIPDatabase.COUNTRY, GeoIPDatabase.ENTERPRISE), + CITY_NAME("city_name", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + CITY_CONFIDENCE("city_confidence", GeoIPDatabase.ENTERPRISE), + LOCATION("location", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + LATITUDE("latitude", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + LONGITUDE("longitude", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + LOCATION_ACCURACY_RADIUS("location_accuracy_radius", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + METRO_CODE("metro_code", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + TIME_ZONE("time_zone", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + POSTAL_CODE("postal_code", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + POSTAL_CODE_CONFIDENCE("postal_code_confidence", GeoIPDatabase.ENTERPRISE), + MOST_SPECIFIED_SUBDIVISION_NAME("most_specified_subdivision_name", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + MOST_SPECIFIED_SUBDIVISION_ISO_CODE("most_specified_subdivision_iso_code", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + MOST_SPECIFIED_SUBDIVISION_CONFIDENCE("most_specified_subdivision_confidence", GeoIPDatabase.ENTERPRISE), + LEAST_SPECIFIED_SUBDIVISION_NAME("least_specified_subdivision_name", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + LEAST_SPECIFIED_SUBDIVISION_ISO_CODE("least_specified_subdivision_iso_code", GeoIPDatabase.CITY, GeoIPDatabase.ENTERPRISE), + LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE("least_specified_subdivision_confidence", GeoIPDatabase.ENTERPRISE), + + ASN("asn", GeoIPDatabase.ASN), + ASN_ORGANIZATION("asn_organization", GeoIPDatabase.ASN), + NETWORK("network", GeoIPDatabase.ASN), + IP("ip", GeoIPDatabase.ASN); + + private final HashSet geoIPDatabases; + private final String fieldName; + + GeoIPField(final String fieldName, final GeoIPDatabase... geoIPDatabases) { + this.fieldName = fieldName; + this.geoIPDatabases = new HashSet<>(Arrays.asList(geoIPDatabases)); + } + + public static GeoIPField findByName(final String name) { + GeoIPField result = null; + for (GeoIPField geoIPField : values()) { + if (geoIPField.getFieldName().equalsIgnoreCase(name)) { + result = geoIPField; + break; + } + } + return result; + } + + public String getFieldName() { + return fieldName; + } + + public Set getGeoIPDatabases() { + return geoIPDatabases; + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessor.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessor.java index 3b3e42f9d0..7b09aaae73 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessor.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessor.java @@ -15,7 +15,9 @@ import org.opensearch.dataprepper.model.processor.Processor; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.processor.configuration.EntryConfig; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.EnrichFailedException; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader; +import org.opensearch.dataprepper.plugins.processor.exception.EngineFailureException; +import org.opensearch.dataprepper.plugins.processor.exception.EnrichFailedException; import org.opensearch.dataprepper.plugins.processor.extension.GeoIPProcessorService; import org.opensearch.dataprepper.plugins.processor.extension.GeoIpConfigSupplier; import org.opensearch.dataprepper.plugins.processor.utils.IPValidationCheck; @@ -23,11 +25,15 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.io.IOException; import java.net.InetAddress; +import java.net.UnknownHostException; +import java.util.ArrayList; import java.util.Collection; +import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; /** * Implementation class of geoIP-processor plugin. It is responsible for enrichment of @@ -37,16 +43,23 @@ public class GeoIPProcessor extends AbstractProcessor, Record> { private static final Logger LOG = LoggerFactory.getLogger(GeoIPProcessor.class); - //TODO: rename metrics - static final String GEO_IP_PROCESSING_MATCH = "geoIpProcessingMatch"; - static final String GEO_IP_PROCESSING_MISMATCH = "geoIpProcessingMismatch"; - private final Counter geoIpProcessingMatchCounter; - private final Counter geoIpProcessingMismatchCounter; + static final String GEO_IP_EVENTS_PROCESSED = "eventsProcessed"; + static final String GEO_IP_EVENTS_SUCCEEDED = "eventsSucceeded"; + static final String GEO_IP_EVENTS_FAILED = "eventsFailed"; + static final String GEO_IP_EVENTS_FAILED_ENGINE_EXCEPTION = "eventsFailedEngineException"; + static final String GEO_IP_EVENTS_FAILED_IP_NOT_FOUND = "eventsFailedIpNotFound"; + private final Counter geoIpEventsProcessed; + private final Counter geoIpEventsSucceeded; + private final Counter geoIpEventsFailed; + private final Counter geoIpEventsFailedEngineException; + private final Counter geoIpEventsFailedIPNotFound; private final GeoIPProcessorConfig geoIPProcessorConfig; - private final List tagsOnFailure; + private final List tagsOnEngineFailure; + private final List tagsOnIPNotFound; private final GeoIPProcessorService geoIPProcessorService; - private final String whenCondition; private final ExpressionEvaluator expressionEvaluator; + private final Map> entryFieldsMap; + final Map> entryDatabaseMap; /** * GeoIPProcessor constructor for initialization of required attributes @@ -60,13 +73,20 @@ public GeoIPProcessor(final PluginMetrics pluginMetrics, final GeoIpConfigSupplier geoIpConfigSupplier, final ExpressionEvaluator expressionEvaluator) { super(pluginMetrics); + this.geoIPProcessorService = geoIpConfigSupplier.getGeoIPProcessorService().orElseThrow(() -> + new IllegalStateException("geoip_service configuration is required when using geoip processor.")); this.geoIPProcessorConfig = geoIPProcessorConfig; - this.geoIPProcessorService = geoIpConfigSupplier.getGeoIPProcessorService(); - this.tagsOnFailure = geoIPProcessorConfig.getTagsOnFailure(); - this.whenCondition = geoIPProcessorConfig.getWhenCondition(); + this.tagsOnEngineFailure = geoIPProcessorConfig.getTagsOnEngineFailure(); + this.tagsOnIPNotFound = geoIPProcessorConfig.getTagsOnIPNotFound(); this.expressionEvaluator = expressionEvaluator; - this.geoIpProcessingMatchCounter = pluginMetrics.counter(GEO_IP_PROCESSING_MATCH); - this.geoIpProcessingMismatchCounter = pluginMetrics.counter(GEO_IP_PROCESSING_MISMATCH); + this.geoIpEventsProcessed = pluginMetrics.counter(GEO_IP_EVENTS_PROCESSED); + this.geoIpEventsSucceeded = pluginMetrics.counter(GEO_IP_EVENTS_SUCCEEDED); + this.geoIpEventsFailed = pluginMetrics.counter(GEO_IP_EVENTS_FAILED); + this.geoIpEventsFailedEngineException = pluginMetrics.counter(GEO_IP_EVENTS_FAILED_ENGINE_EXCEPTION); + this.geoIpEventsFailedIPNotFound = pluginMetrics.counter(GEO_IP_EVENTS_FAILED_IP_NOT_FOUND); + + this.entryFieldsMap = populateGeoIPFields(); + this.entryDatabaseMap = populateGeoIPDatabases(); } /** @@ -77,51 +97,169 @@ public GeoIPProcessor(final PluginMetrics pluginMetrics, @Override public Collection> doExecute(final Collection> records) { Map geoData; + try (final GeoIPDatabaseReader geoIPDatabaseReader = geoIPProcessorService.getGeoIPDatabaseReader()) { - for (final Record eventRecord : records) { - final Event event = eventRecord.getData(); - if (whenCondition != null && !expressionEvaluator.evaluateConditional(whenCondition, event)) { - continue; - } - for (EntryConfig entry : geoIPProcessorConfig.getEntries()) { - final String source = entry.getSource(); - final List attributes = entry.getFields(); - final String ipAddress = event.get(source, String.class); + for (final Record eventRecord : records) { + final Event event = eventRecord.getData(); + final String whenCondition = geoIPProcessorConfig.getWhenCondition(); + // continue if when condition is null or is false + // or if database reader is null or database reader is expired + if (checkConditionAndDatabaseReader(geoIPDatabaseReader, event, whenCondition)) continue; - //Lookup from DB - if (ipAddress != null && !ipAddress.isEmpty()) { + boolean eventSucceeded = true; + boolean ipNotFound = false; + boolean engineFailure = false; + + for (final EntryConfig entry : geoIPProcessorConfig.getEntries()) { + final String source = entry.getSource(); + final List fields = entryFieldsMap.get(entry); + final Set databases = entryDatabaseMap.get(entry); + String ipAddress = null; try { - if (IPValidationCheck.isPublicIpAddress(ipAddress)) { - geoData = geoIPProcessorService.getGeoData(InetAddress.getByName(ipAddress), attributes); - eventRecord.getData().put(entry.getTarget(), geoData); - geoIpProcessingMatchCounter.increment(); + ipAddress = event.get(source, String.class); + } catch (final Exception e) { + eventSucceeded = false; + ipNotFound = true; + LOG.error(DataPrepperMarkers.EVENT, "Failed to get IP address from [{}] in event: [{}]. Caused by:[{}]", + source, event, e.getMessage()); + } + + //Lookup from DB + if (ipAddress != null && !ipAddress.isEmpty()) { + try { + if (IPValidationCheck.isPublicIpAddress(ipAddress)) { + geoData = geoIPDatabaseReader.getGeoData(InetAddress.getByName(ipAddress), fields, databases); + if (geoData.isEmpty()) { + ipNotFound = true; + eventSucceeded = false; + } else { + eventRecord.getData().put(entry.getTarget(), geoData); + } + } else { + // no enrichment if IP is not public + ipNotFound = true; + eventSucceeded = false; + } + } catch (final UnknownHostException e) { + ipNotFound = true; + eventSucceeded = false; + LOG.error(DataPrepperMarkers.EVENT, "Failed to validate IP address: [{}] in event: [{}]. Caused by:[{}]", + ipAddress, event, e.getMessage()); + LOG.error("Failed to validate IP address: [{}]. Caused by:[{}]", ipAddress, e.getMessage()); + } catch (final EnrichFailedException e) { + ipNotFound = true; + eventSucceeded = false; + LOG.error(DataPrepperMarkers.EVENT, "IP address not found in database for IP: [{}] in event: [{}]. Caused by:[{}]", + ipAddress, event, e.getMessage()); + LOG.error("IP address not found in database for IP: [{}]. Caused by:[{}]", ipAddress, e.getMessage()); + } catch (final EngineFailureException e) { + engineFailure = true; + eventSucceeded = false; + LOG.error(DataPrepperMarkers.EVENT, "Failed to get Geo data for event: [{}] for the IP address [{}]. Caused by:{}", + event, ipAddress, e.getMessage()); + LOG.error("Failed to get Geo data for the IP address [{}]. Caused by:{}", ipAddress, e.getMessage()); } - } catch (final IOException | EnrichFailedException ex) { - geoIpProcessingMismatchCounter.increment(); - event.getMetadata().addTags(tagsOnFailure); - LOG.error(DataPrepperMarkers.EVENT, "Failed to get Geo data for event: [{}] for the IP address [{}]", event, ipAddress, ex); + } else { + //No Enrichment if IP is null or empty + eventSucceeded = false; + ipNotFound = true; } - } else { - //No Enrichment. - event.getMetadata().addTags(tagsOnFailure); } + + updateTagsAndMetrics(event, eventSucceeded, ipNotFound, engineFailure); } + } catch (final Exception e) { + LOG.error("Encountered exception in geoip processor.", e); } return records; } + private void updateTagsAndMetrics(final Event event, final boolean eventSucceeded, final boolean ipNotFound, final boolean engineFailure) { + if (ipNotFound) { + event.getMetadata().addTags(tagsOnIPNotFound); + geoIpEventsFailedIPNotFound.increment(); + } + if (engineFailure) { + event.getMetadata().addTags(tagsOnEngineFailure); + geoIpEventsFailedEngineException.increment(); + } + if (eventSucceeded) { + geoIpEventsSucceeded.increment(); + } else { + geoIpEventsFailed.increment(); + } + } + + private boolean checkConditionAndDatabaseReader(final GeoIPDatabaseReader geoIPDatabaseReader, final Event event, final String whenCondition) { + if (whenCondition != null && !expressionEvaluator.evaluateConditional(whenCondition, event)) { + return true; + } + geoIpEventsProcessed.increment(); + + // if database reader is not created or if all database readers are expired + if (geoIPDatabaseReader == null || geoIPDatabaseReader.isExpired()) { + event.getMetadata().addTags(tagsOnEngineFailure); + geoIpEventsFailed.increment(); + return true; + } + return false; + } + + private Map> populateGeoIPFields() { + final Map> entryConfigFieldsMap = new HashMap<>(); + for (final EntryConfig entry: geoIPProcessorConfig.getEntries()) { + final List includeFields = entry.getIncludeFields(); + final List excludeFields = entry.getExcludeFields(); + List geoIPFields = new ArrayList<>(); + if (includeFields != null && !includeFields.isEmpty()) { + for (final String field : includeFields) { + final GeoIPField geoIPField = GeoIPField.findByName(field); + if (geoIPField != null) { + geoIPFields.add(geoIPField); + } + } + } else if (excludeFields != null) { + final List excludeGeoIPFields = new ArrayList<>(); + for (final String field : excludeFields) { + final GeoIPField geoIPField = GeoIPField.findByName(field); + if (geoIPField != null) { + excludeGeoIPFields.add(geoIPField); + } + } + geoIPFields = new ArrayList<>(List.of(GeoIPField.values())); + geoIPFields.removeAll(excludeGeoIPFields); + } + entryConfigFieldsMap.put(entry, geoIPFields); + } + return entryConfigFieldsMap; + } + + private Map> populateGeoIPDatabases() { + final Map> entryConfigGeoIPDatabaseMap = new HashMap<>(); + for (final EntryConfig entry : geoIPProcessorConfig.getEntries()) { + final List geoIPFields = entryFieldsMap.get(entry); + final Set geoIPDatabasesToUse = new HashSet<>(); + for (final GeoIPField geoIPField : geoIPFields) { + final Set geoIPDatabases = geoIPField.getGeoIPDatabases(); + geoIPDatabasesToUse.addAll(geoIPDatabases); + } + entryConfigGeoIPDatabaseMap.put(entry, geoIPDatabasesToUse); + } + return entryConfigGeoIPDatabaseMap; + } + @Override public void prepareForShutdown() { } @Override public boolean isReadyForShutdown() { + geoIPProcessorService.shutdown(); return true; } @Override public void shutdown() { - //TODO: delete mmdb files - LOG.info("GeoIP plugin Shutdown"); + } } \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorConfig.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorConfig.java index ca0b00e8e7..cc2445b1f0 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorConfig.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorConfig.java @@ -18,19 +18,21 @@ */ public class GeoIPProcessorConfig { - @JsonProperty("entries") + @Valid @NotNull @Size(min = 1) - @Valid + @JsonProperty("entries") private List entries; - @JsonProperty("tags_on_failure") - private List tagsOnFailure; + @JsonProperty("tags_on_engine_failure") + private List tagsOnEngineFailure; + + @JsonProperty("tags_on_ip_not_found") + private List tagsOnIPNotFound; @JsonProperty("geoip_when") private String whenCondition; - /** * Get List of entries * @return List of EntryConfig @@ -40,11 +42,19 @@ public List getEntries() { } /** - * Get the List of failure tags - * @return List of failure tags + * Get the List of engine failure tags + * @return List of engine failure tags + */ + public List getTagsOnEngineFailure() { + return tagsOnEngineFailure; + } + + /** + * Get the List of invalid IP / IP not found in database tags + * @return List of invalid IP / IP not found in database tags */ - public List getTagsOnFailure() { - return tagsOnFailure; + public List getTagsOnIPNotFound() { + return tagsOnIPNotFound; } /** diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/configuration/EntryConfig.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/configuration/EntryConfig.java index c1e95373a3..d9c22d6143 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/configuration/EntryConfig.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/configuration/EntryConfig.java @@ -6,12 +6,13 @@ package org.opensearch.dataprepper.plugins.processor.configuration; import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.constraints.AssertTrue; import jakarta.validation.constraints.NotEmpty; import java.util.List; public class EntryConfig { - static final String DEFAULT_TARGET = "geolocation"; + static final String DEFAULT_TARGET = "geo"; @JsonProperty("source") @NotEmpty private String source; @@ -19,8 +20,11 @@ public class EntryConfig { @JsonProperty("target") private String target = DEFAULT_TARGET; - @JsonProperty("fields") - private List fields; + @JsonProperty("include_fields") + private List includeFields; + + @JsonProperty("exclude_fields") + private List excludeFields; public String getSource() { return source; @@ -30,7 +34,19 @@ public String getTarget() { return target; } - public List getFields() { - return fields; + public List getIncludeFields() { + return includeFields; + } + + public List getExcludeFields() { + return excludeFields; + } + + @AssertTrue(message = "include_fields and exclude_fields are mutually exclusive. include_fields or exclude_fields is required.") + boolean areFieldsValid() { + if (includeFields == null && excludeFields == null) { + return false; + } + return includeFields == null || excludeFields == null; } } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/AutoCountingDatabaseReader.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/AutoCountingDatabaseReader.java new file mode 100644 index 0000000000..b6cac46508 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/AutoCountingDatabaseReader.java @@ -0,0 +1,54 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.databaseenrich; + +import org.opensearch.dataprepper.plugins.processor.GeoIPDatabase; +import org.opensearch.dataprepper.plugins.processor.GeoIPField; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.net.InetAddress; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; + +public class AutoCountingDatabaseReader implements GeoIPDatabaseReader { + private static final Logger LOG = LoggerFactory.getLogger(AutoCountingDatabaseReader.class); + private final GeoIPDatabaseReader delegateDatabaseReader; + private final AtomicInteger closeCount; + + public AutoCountingDatabaseReader(final GeoIPDatabaseReader geoIPDatabaseReader) { + this.delegateDatabaseReader = geoIPDatabaseReader; + this.closeCount = new AtomicInteger(1); + } + + @Override + public Map getGeoData(final InetAddress inetAddress, + final List fields, + final Set geoIPDatabases) { + return delegateDatabaseReader.getGeoData(inetAddress, fields, geoIPDatabases); + } + + @Override + public boolean isExpired() { + return delegateDatabaseReader.isExpired(); + } + + @Override + public void retain() { + closeCount.incrementAndGet(); + } + + @Override + public void close() throws Exception { + final int count = closeCount.decrementAndGet(); + if (count == 0) { + LOG.debug("Closing old geoip database readers"); + delegateDatabaseReader.close(); + } + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIP2DatabaseReader.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIP2DatabaseReader.java new file mode 100644 index 0000000000..e428695246 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIP2DatabaseReader.java @@ -0,0 +1,174 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.databaseenrich; + +import com.maxmind.geoip2.DatabaseReader; +import com.maxmind.geoip2.exception.GeoIp2Exception; +import com.maxmind.geoip2.model.AsnResponse; +import com.maxmind.geoip2.model.EnterpriseResponse; +import com.maxmind.geoip2.record.City; +import com.maxmind.geoip2.record.Continent; +import com.maxmind.geoip2.record.Country; +import com.maxmind.geoip2.record.Location; +import com.maxmind.geoip2.record.Postal; +import com.maxmind.geoip2.record.RepresentedCountry; +import com.maxmind.geoip2.record.Subdivision; +import org.opensearch.dataprepper.plugins.processor.GeoIPDatabase; +import org.opensearch.dataprepper.plugins.processor.GeoIPField; +import org.opensearch.dataprepper.plugins.processor.exception.DatabaseReaderInitializationException; +import org.opensearch.dataprepper.plugins.processor.exception.EnrichFailedException; +import org.opensearch.dataprepper.plugins.processor.exception.NoValidDatabaseFoundException; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPFileManager; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DatabaseReaderBuilder; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.io.IOException; +import java.net.InetAddress; +import java.nio.file.Path; +import java.time.Instant; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig.GEOIP2_ENTERPRISE; + +public class GeoIP2DatabaseReader implements GeoIPDatabaseReader, AutoCloseable { + private static final Logger LOG = LoggerFactory.getLogger(GeoIP2DatabaseReader.class); + private static final String MAXMIND_GEOIP2_DATABASE_TYPE = "geoip2"; + private final DatabaseReaderBuilder databaseReaderBuilder; + private final String databasePath; + private final int cacheSize; + private final AtomicInteger closeCount; + private final GeoIPFileManager geoIPFileManager; + private final AtomicBoolean isEnterpriseDatabaseExpired; + private DatabaseReader enterpriseDatabaseReader; + private Instant enterpriseDatabaseBuildDate; + + public GeoIP2DatabaseReader(final DatabaseReaderBuilder databaseReaderBuilder, + final GeoIPFileManager geoIPFileManager, + final String databasePath, final int cacheSize) { + this.databaseReaderBuilder = databaseReaderBuilder; + this.geoIPFileManager = geoIPFileManager; + this.databasePath = databasePath; + this.cacheSize = cacheSize; + closeCount = new AtomicInteger(1); + this.isEnterpriseDatabaseExpired = new AtomicBoolean(false); + buildDatabaseReaders(); + } + + private void buildDatabaseReaders() { + try { + final Optional enterpriseDatabaseName = getDatabaseName(GEOIP2_ENTERPRISE, databasePath, MAXMIND_GEOIP2_DATABASE_TYPE); + + if (enterpriseDatabaseName.isPresent()) { + enterpriseDatabaseReader = databaseReaderBuilder.buildReader(Path.of( + databasePath + File.separator + enterpriseDatabaseName.get()), cacheSize); + enterpriseDatabaseBuildDate = enterpriseDatabaseReader.getMetadata().getBuildDate().toInstant(); + } + } catch (final IOException ex) { + throw new DatabaseReaderInitializationException("Exception while creating GeoIP2 DatabaseReaders due to: " + ex.getMessage()); + } + + if (enterpriseDatabaseReader == null) { + throw new NoValidDatabaseFoundException("Unable to initialize GeoIP2 database, make sure it is valid."); + } + } + @Override + public Map getGeoData(final InetAddress inetAddress, final List fields, final Set geoIPDatabases) { + Map geoData = new HashMap<>(); + + try { + if (enterpriseDatabaseReader != null && !isEnterpriseDatabaseExpired.get() && geoIPDatabases.contains(GeoIPDatabase.ENTERPRISE)) { + final Optional optionalEnterpriseResponse = enterpriseDatabaseReader.tryEnterprise(inetAddress); + optionalEnterpriseResponse.ifPresent(response -> processEnterpriseResponse(response, geoData, fields)); + } + + if (enterpriseDatabaseReader != null && !isEnterpriseDatabaseExpired.get() && geoIPDatabases.contains(GeoIPDatabase.ASN)) { + final Optional asnResponse = enterpriseDatabaseReader.tryAsn(inetAddress); + asnResponse.ifPresent(response -> processAsnResponse(response, geoData, fields)); + } + + } catch (final GeoIp2Exception e) { + throw new EnrichFailedException("Address not found in database."); + } catch (final IOException e) { + throw new EnrichFailedException("Failed to close database readers gracefully. It can be due to expired databases"); + } + return geoData; + } + + private void processEnterpriseResponse(final EnterpriseResponse enterpriseResponse, final Map geoData, final List fields) { + final Continent continent = enterpriseResponse.getContinent(); + final Country country = enterpriseResponse.getCountry(); + final Country registeredCountry = enterpriseResponse.getRegisteredCountry(); + final RepresentedCountry representedCountry = enterpriseResponse.getRepresentedCountry(); + + final City city = enterpriseResponse.getCity(); + final Location location = enterpriseResponse.getLocation(); + final Postal postal = enterpriseResponse.getPostal(); + final Subdivision mostSpecificSubdivision = enterpriseResponse.getMostSpecificSubdivision(); + final Subdivision leastSpecificSubdivision = enterpriseResponse.getLeastSpecificSubdivision(); + + extractContinentFields(continent, geoData, fields); + extractCountryFields(country, geoData, fields, true); + extractRegisteredCountryFields(registeredCountry, geoData, fields); + extractRepresentedCountryFields(representedCountry, geoData, fields); + extractCityFields(city, geoData, fields, true); + extractLocationFields(location, geoData, fields); + extractPostalFields(postal, geoData, fields, true); + extractMostSpecifiedSubdivisionFields(mostSpecificSubdivision, geoData, fields, true); + extractLeastSpecifiedSubdivisionFields(leastSpecificSubdivision, geoData, fields, true); + } + + private void processAsnResponse(final AsnResponse asnResponse, final Map geoData, final List fields) { + extractAsnFields(asnResponse, geoData, fields); + } + + @Override + public boolean isExpired() { + final Instant instant = Instant.now(); + if (enterpriseDatabaseReader == null) { + return true; + } + if (isEnterpriseDatabaseExpired.get()) { + return true; + } + if (enterpriseDatabaseBuildDate.plus(MAX_EXPIRY_DURATION).isBefore(instant)) { + isEnterpriseDatabaseExpired.set(true); + closeReader(); + } + return isEnterpriseDatabaseExpired.get(); + } + + @Override + public void retain() { + + } + + @Override + public void close() { + closeReader(); + } + + private void closeReader() { + try { + if (enterpriseDatabaseReader != null) { + enterpriseDatabaseReader.close(); + } + } catch (final IOException e) { + LOG.debug("Failed to close Maxmind database readers due to: {}. Force closing readers.", e.getMessage()); + } + + // delete database directory + final File file = new File(databasePath); + geoIPFileManager.deleteDirectory(file); + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIPDatabaseReader.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIPDatabaseReader.java new file mode 100644 index 0000000000..c4e7f022fe --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIPDatabaseReader.java @@ -0,0 +1,354 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.databaseenrich; + +import com.maxmind.geoip2.model.AsnResponse; +import com.maxmind.geoip2.record.City; +import com.maxmind.geoip2.record.Continent; +import com.maxmind.geoip2.record.Country; +import com.maxmind.geoip2.record.Location; +import com.maxmind.geoip2.record.Postal; +import com.maxmind.geoip2.record.RepresentedCountry; +import com.maxmind.geoip2.record.Subdivision; +import org.opensearch.dataprepper.plugins.processor.GeoIPDatabase; +import org.opensearch.dataprepper.plugins.processor.GeoIPField; + +import java.io.File; +import java.net.InetAddress; +import java.time.Duration; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +/** + * Interface for storing and maintaining MaxMind database readers + */ + +public interface GeoIPDatabaseReader extends AutoCloseable { + String MAXMIND_DATABASE_EXTENSION = ".mmdb"; + Duration MAX_EXPIRY_DURATION = Duration.ofDays(30); + String LAT = "lat"; + String LON = "lon"; + + /** + * Gets the geo data from the {@link com.maxmind.geoip2.DatabaseReader} + * + * @param inetAddress InetAddress + * @return Map of geo field and value pairs from IP address + * + * @since 2.7 + */ + Map getGeoData(InetAddress inetAddress, List fields, Set geoIPDatabases); + + /** + * Gets if the database is expired from metadata or last updated timestamp + * + * @return boolean indicating if database is expired + * + * @since 2.7 + */ + boolean isExpired(); + + /** + * Retains the reader which prevents from closing if it's being used + * + * @since 2.7 + */ + void retain(); + + /** + * Enrich attributes + * @param geoData geoData + * @param fieldName fieldName + * @param fieldValue fieldValue + */ + default void enrichData(final Map geoData, final String fieldName, final Object fieldValue) { + if (!geoData.containsKey(fieldName) && fieldValue != null) { + geoData.put(fieldName, fieldValue); + } + } + + default Optional getDatabaseName(final String database, final String databasePath, final String databaseType) { + final File file = new File(databasePath); + if (file.exists() && file.isDirectory()) { + final String[] list = file.list(); + for (final String fileName: list) { + final String lowerCaseFileName = fileName.toLowerCase(); + if (lowerCaseFileName.contains(database.toLowerCase()) + && fileName.endsWith(MAXMIND_DATABASE_EXTENSION) + && lowerCaseFileName.contains(databaseType)) { + return Optional.of(fileName); + } + } + } + return Optional.empty(); + } + + default void extractContinentFields(final Continent continent, + final Map geoData, + final List fields) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + switch (field) { + case CONTINENT_CODE: + enrichData(geoData, GeoIPField.CONTINENT_CODE.getFieldName(), continent.getCode()); + break; + case CONTINENT_NAME: + enrichData(geoData, GeoIPField.CONTINENT_NAME.getFieldName(), continent.getName()); + break; + } + } + } else { + // add all fields + enrichData(geoData, GeoIPField.CONTINENT_CODE.getFieldName(), continent.getCode()); + enrichData(geoData, GeoIPField.CONTINENT_NAME.getFieldName(), continent.getName()); + } + } + + default void extractCountryFields(final Country country, + final Map geoData, + final List fields, + final boolean isEnterpriseDatabase) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + switch (field) { + case COUNTRY_NAME: + enrichData(geoData, GeoIPField.COUNTRY_NAME.getFieldName(), country.getName()); + break; + case IS_COUNTRY_IN_EUROPEAN_UNION: + enrichData(geoData, GeoIPField.IS_COUNTRY_IN_EUROPEAN_UNION.getFieldName(), country.isInEuropeanUnion()); + break; + case COUNTRY_ISO_CODE: + enrichData(geoData, GeoIPField.COUNTRY_ISO_CODE.getFieldName(), country.getIsoCode()); + break; + case COUNTRY_CONFIDENCE: + if (isEnterpriseDatabase) + enrichData(geoData, GeoIPField.COUNTRY_CONFIDENCE.getFieldName(), country.getConfidence()); + break; + } + } + } else { + // add all fields + enrichData(geoData, GeoIPField.COUNTRY_NAME.getFieldName(), country.getName()); + enrichData(geoData, GeoIPField.IS_COUNTRY_IN_EUROPEAN_UNION.getFieldName(), country.isInEuropeanUnion()); + enrichData(geoData, GeoIPField.COUNTRY_ISO_CODE.getFieldName(), country.getIsoCode()); + if (isEnterpriseDatabase) + enrichData(geoData,GeoIPField. COUNTRY_CONFIDENCE.getFieldName(), country.getConfidence()); + } + } + + default void extractRegisteredCountryFields(final Country registeredCountry, + final Map geoData, + final List fields) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + switch (field) { + case REGISTERED_COUNTRY_NAME: + enrichData(geoData, GeoIPField.REGISTERED_COUNTRY_NAME.getFieldName(), registeredCountry.getName()); + break; + case REGISTERED_COUNTRY_ISO_CODE: + enrichData(geoData, GeoIPField.REGISTERED_COUNTRY_ISO_CODE.getFieldName(), registeredCountry.getIsoCode()); + break; + } + } + } else { + // add all fields + enrichData(geoData, GeoIPField.REGISTERED_COUNTRY_NAME.getFieldName(), registeredCountry.getName()); + enrichData(geoData, GeoIPField.REGISTERED_COUNTRY_ISO_CODE.getFieldName(), registeredCountry.getIsoCode()); + } + } + + default void extractRepresentedCountryFields(final RepresentedCountry representedCountry, + final Map geoData, + final List fields) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + switch (field) { + case REPRESENTED_COUNTRY_NAME: + enrichData(geoData, GeoIPField.REPRESENTED_COUNTRY_NAME.getFieldName(), representedCountry.getName()); + break; + case REPRESENTED_COUNTRY_ISO_CODE: + enrichData(geoData, GeoIPField.REPRESENTED_COUNTRY_ISO_CODE.getFieldName(), representedCountry.getIsoCode()); + break; + case REPRESENTED_COUNTRY_TYPE: + enrichData(geoData, GeoIPField.REPRESENTED_COUNTRY_TYPE.getFieldName(), representedCountry.getType()); + break; + } + } + } else { + // add all fields + enrichData(geoData, GeoIPField.REPRESENTED_COUNTRY_NAME.getFieldName(), representedCountry.getName()); + enrichData(geoData, GeoIPField.REPRESENTED_COUNTRY_ISO_CODE.getFieldName(), representedCountry.getIsoCode()); + enrichData(geoData, GeoIPField.REPRESENTED_COUNTRY_TYPE.getFieldName(), representedCountry.getType()); + } + } + + default void extractCityFields(final City city, + final Map geoData, + final List fields, + final boolean isEnterpriseDatabase) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + if (field.equals(GeoIPField.CITY_NAME)) { + enrichData(geoData, GeoIPField.CITY_NAME.getFieldName(), city.getName()); + } else if (isEnterpriseDatabase && field.equals(GeoIPField.CITY_CONFIDENCE)) { + enrichData(geoData, GeoIPField.CITY_CONFIDENCE.getFieldName(), city.getConfidence()); + } + } + } else{ + enrichData(geoData, GeoIPField.CITY_NAME.getFieldName(), city.getName()); + if (isEnterpriseDatabase) + enrichData(geoData, GeoIPField.CITY_CONFIDENCE.getFieldName(), city.getConfidence()); + } + } + + default void extractLocationFields(final Location location, + final Map geoData, + final List fields) { + final Map locationObject = new HashMap<>(); + locationObject.put(LAT, location.getLatitude()); + locationObject.put(LON, location.getLongitude()); + + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + switch (field) { + case LOCATION: + enrichData(geoData, GeoIPField.LOCATION.getFieldName(), locationObject); + break; + case LATITUDE: + enrichData(geoData, GeoIPField.LATITUDE.getFieldName(), location.getLatitude()); + break; + case LONGITUDE: + enrichData(geoData, GeoIPField.LONGITUDE.getFieldName(), location.getLongitude()); + break; + case METRO_CODE: + enrichData(geoData, GeoIPField.METRO_CODE.getFieldName(), location.getMetroCode()); + break; + case TIME_ZONE: + enrichData(geoData, GeoIPField.TIME_ZONE.getFieldName(), location.getTimeZone()); + break; + case LOCATION_ACCURACY_RADIUS: + enrichData(geoData, GeoIPField.LOCATION_ACCURACY_RADIUS.getFieldName(), location.getAccuracyRadius()); + break; + } + } + } else{ + // add all fields - latitude & longitude will be part of location key + enrichData(geoData, GeoIPField.LOCATION.getFieldName(), locationObject); + enrichData(geoData, GeoIPField.METRO_CODE.getFieldName(), location.getMetroCode()); + enrichData(geoData, GeoIPField.TIME_ZONE.getFieldName(), location.getTimeZone()); + enrichData(geoData, GeoIPField.LOCATION_ACCURACY_RADIUS.getFieldName(), location.getAccuracyRadius()); + } + } + + default void extractPostalFields(final Postal postal, + final Map geoData, + final List fields, + final boolean isEnterpriseDatabase) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + if (field.equals(GeoIPField.POSTAL_CODE)) { + enrichData(geoData, GeoIPField.POSTAL_CODE.getFieldName(), postal.getCode()); + } else if (isEnterpriseDatabase && field.equals(GeoIPField.POSTAL_CODE_CONFIDENCE)) { + enrichData(geoData, GeoIPField.POSTAL_CODE_CONFIDENCE.getFieldName(), postal.getConfidence()); + } + } + } else{ + enrichData(geoData, GeoIPField.POSTAL_CODE.getFieldName(), postal.getCode()); + if (isEnterpriseDatabase) + enrichData(geoData, GeoIPField.POSTAL_CODE_CONFIDENCE.getFieldName(), postal.getConfidence()); + } + } + + default void extractMostSpecifiedSubdivisionFields(final Subdivision subdivision, + final Map geoData, + final List fields, + final boolean isEnterpriseDatabase) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + switch (field) { + case MOST_SPECIFIED_SUBDIVISION_NAME: + enrichData(geoData, GeoIPField.MOST_SPECIFIED_SUBDIVISION_NAME.getFieldName(), subdivision.getName()); + break; + case MOST_SPECIFIED_SUBDIVISION_ISO_CODE: + enrichData(geoData, GeoIPField.MOST_SPECIFIED_SUBDIVISION_ISO_CODE.getFieldName(), subdivision.getIsoCode()); + break; + case MOST_SPECIFIED_SUBDIVISION_CONFIDENCE: + if (isEnterpriseDatabase) + enrichData(geoData, GeoIPField.MOST_SPECIFIED_SUBDIVISION_CONFIDENCE.getFieldName(), subdivision.getConfidence()); + break; + } + } + } else { + // add all fields + enrichData(geoData, GeoIPField.MOST_SPECIFIED_SUBDIVISION_NAME.getFieldName(), subdivision.getName()); + enrichData(geoData, GeoIPField.MOST_SPECIFIED_SUBDIVISION_ISO_CODE.getFieldName(), subdivision.getIsoCode()); + if (isEnterpriseDatabase) + enrichData(geoData, GeoIPField.MOST_SPECIFIED_SUBDIVISION_CONFIDENCE.getFieldName(), subdivision.getConfidence()); + } + } + + default void extractLeastSpecifiedSubdivisionFields(final Subdivision subdivision, + final Map geoData, + final List fields, + final boolean isEnterpriseDatabase) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + switch (field) { + case LEAST_SPECIFIED_SUBDIVISION_NAME: + enrichData(geoData, GeoIPField.LEAST_SPECIFIED_SUBDIVISION_NAME.getFieldName(), subdivision.getName()); + break; + case LEAST_SPECIFIED_SUBDIVISION_ISO_CODE: + enrichData(geoData, GeoIPField.LEAST_SPECIFIED_SUBDIVISION_ISO_CODE.getFieldName(), subdivision.getIsoCode()); + break; + case LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE: + if (isEnterpriseDatabase) + enrichData(geoData, GeoIPField.LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE.getFieldName(), subdivision.getConfidence()); + break; + } + } + } else { + // add all fields + enrichData(geoData, GeoIPField.LEAST_SPECIFIED_SUBDIVISION_NAME.getFieldName(), subdivision.getName()); + enrichData(geoData, GeoIPField.LEAST_SPECIFIED_SUBDIVISION_ISO_CODE.getFieldName(), subdivision.getIsoCode()); + if (isEnterpriseDatabase) + enrichData(geoData, GeoIPField.LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE.getFieldName(), subdivision.getConfidence()); + } + } + + default void extractAsnFields(final AsnResponse asnResponse, + final Map geoData, + final List fields) { + if (!fields.isEmpty()) { + for (final GeoIPField field : fields) { + switch (field) { + case ASN: + enrichData(geoData, GeoIPField.ASN.getFieldName(), asnResponse.getAutonomousSystemNumber()); + break; + case ASN_ORGANIZATION: + enrichData(geoData, GeoIPField.ASN_ORGANIZATION.getFieldName(), asnResponse.getAutonomousSystemOrganization()); + break; + case NETWORK: + enrichData(geoData, GeoIPField.NETWORK.getFieldName(), asnResponse.getNetwork().toString()); + break; + case IP: + enrichData(geoData, GeoIPField.IP.getFieldName(), asnResponse.getIpAddress()); + break; + } + } + } else { + // add all fields + enrichData(geoData, GeoIPField.ASN.getFieldName(), asnResponse.getAutonomousSystemNumber()); + enrichData(geoData, GeoIPField.ASN_ORGANIZATION.getFieldName(), asnResponse.getAutonomousSystemOrganization()); + enrichData(geoData, GeoIPField.NETWORK.getFieldName(), asnResponse.getNetwork().toString()); + enrichData(geoData, GeoIPField.IP.getFieldName(), asnResponse.getIpAddress()); + } + } + + +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoLite2DatabaseReader.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoLite2DatabaseReader.java new file mode 100644 index 0000000000..72304505ad --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoLite2DatabaseReader.java @@ -0,0 +1,254 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.databaseenrich; + +import com.maxmind.geoip2.DatabaseReader; +import com.maxmind.geoip2.exception.GeoIp2Exception; +import com.maxmind.geoip2.model.AsnResponse; +import com.maxmind.geoip2.model.CityResponse; +import com.maxmind.geoip2.model.CountryResponse; +import com.maxmind.geoip2.record.City; +import com.maxmind.geoip2.record.Continent; +import com.maxmind.geoip2.record.Country; +import com.maxmind.geoip2.record.Location; +import com.maxmind.geoip2.record.Postal; +import com.maxmind.geoip2.record.RepresentedCountry; +import com.maxmind.geoip2.record.Subdivision; +import org.opensearch.dataprepper.plugins.processor.GeoIPDatabase; +import org.opensearch.dataprepper.plugins.processor.GeoIPField; +import org.opensearch.dataprepper.plugins.processor.exception.DatabaseReaderInitializationException; +import org.opensearch.dataprepper.plugins.processor.exception.EngineFailureException; +import org.opensearch.dataprepper.plugins.processor.exception.EnrichFailedException; +import org.opensearch.dataprepper.plugins.processor.exception.NoValidDatabaseFoundException; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPFileManager; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DatabaseReaderBuilder; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.io.IOException; +import java.net.InetAddress; +import java.nio.file.Path; +import java.time.Instant; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig.GEOLITE2_ASN; +import static org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig.GEOLITE2_CITY; +import static org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig.GEOLITE2_COUNTRY; + +public class GeoLite2DatabaseReader implements GeoIPDatabaseReader, AutoCloseable { + private static final Logger LOG = LoggerFactory.getLogger(GeoLite2DatabaseReader.class); + static final String MAXMIND_GEOLITE2_DATABASE_TYPE = "geolite2"; + private final DatabaseReaderBuilder databaseReaderBuilder; + private final String databasePath; + private final int cacheSize; + private final AtomicBoolean isCountryDatabaseExpired; + private final AtomicBoolean isCityDatabaseExpired; + private final AtomicBoolean isAsnDatabaseExpired; + private final GeoIPFileManager geoIPFileManager; + private DatabaseReader cityDatabaseReader; + private DatabaseReader countryDatabaseReader; + private DatabaseReader asnDatabaseReader; + private Instant cityDatabaseBuildDate; + private Instant countryDatabaseBuildDate; + private Instant asnDatabaseBuildDate; + + public GeoLite2DatabaseReader(final DatabaseReaderBuilder databaseReaderBuilder, + final GeoIPFileManager geoIPFileManager, + final String databasePath, final int cacheSize) { + this.databaseReaderBuilder = databaseReaderBuilder; + this.geoIPFileManager = geoIPFileManager; + this.databasePath = databasePath; + this.cacheSize = cacheSize; + this.isCountryDatabaseExpired = new AtomicBoolean(false); + this.isCityDatabaseExpired = new AtomicBoolean(false); + this.isAsnDatabaseExpired = new AtomicBoolean(false); + buildDatabaseReaders(); + } + + private void buildDatabaseReaders() { + try { + final Optional cityDatabaseName = getDatabaseName(GEOLITE2_CITY, databasePath, MAXMIND_GEOLITE2_DATABASE_TYPE); + final Optional countryDatabaseName = getDatabaseName(GEOLITE2_COUNTRY, databasePath, MAXMIND_GEOLITE2_DATABASE_TYPE); + final Optional asnDatabaseName = getDatabaseName(GEOLITE2_ASN, databasePath, MAXMIND_GEOLITE2_DATABASE_TYPE); + + if (cityDatabaseName.isPresent()) { + cityDatabaseReader = databaseReaderBuilder.buildReader(Path.of(databasePath + File.separator + cityDatabaseName.get()), cacheSize); + cityDatabaseBuildDate = cityDatabaseReader.getMetadata().getBuildDate().toInstant(); + } + if (countryDatabaseName.isPresent()) { + countryDatabaseReader = databaseReaderBuilder.buildReader(Path.of(databasePath + File.separator + countryDatabaseName.get()), cacheSize); + countryDatabaseBuildDate = countryDatabaseReader.getMetadata().getBuildDate().toInstant(); + } + if (asnDatabaseName.isPresent()) { + asnDatabaseReader = databaseReaderBuilder.buildReader(Path.of(databasePath + File.separator + asnDatabaseName.get()), cacheSize); + asnDatabaseBuildDate = asnDatabaseReader.getMetadata().getBuildDate().toInstant(); + } + + } catch (final IOException ex) { + throw new DatabaseReaderInitializationException("Exception while creating GeoLite2 DatabaseReaders due to: " + ex.getMessage()); + } + + if (cityDatabaseReader == null && countryDatabaseReader == null && asnDatabaseReader == null) { + throw new NoValidDatabaseFoundException("Unable to initialize any GeoLite2 database, make sure they are valid."); + } + } + + @Override + public Map getGeoData(final InetAddress inetAddress, final List fields, final Set geoIPDatabases) { + final Map geoData = new HashMap<>(); + + try { + if (countryDatabaseReader != null && !isCountryDatabaseExpired.get() && geoIPDatabases.contains(GeoIPDatabase.COUNTRY)) { + final Optional countryResponse = countryDatabaseReader.tryCountry(inetAddress); + countryResponse.ifPresent(response -> processCountryResponse(response, geoData, fields)); + } + + if (cityDatabaseReader != null && !isCityDatabaseExpired.get() && geoIPDatabases.contains(GeoIPDatabase.CITY)) { + final Optional cityResponse = cityDatabaseReader.tryCity(inetAddress); + cityResponse.ifPresent(response -> processCityResponse(response, geoData, fields, geoIPDatabases)); + } + + if (asnDatabaseReader != null && !isAsnDatabaseExpired.get() && geoIPDatabases.contains(GeoIPDatabase.ASN)) { + final Optional asnResponse = asnDatabaseReader.tryAsn(inetAddress); + asnResponse.ifPresent(response -> processAsnResponse(response, geoData, fields)); + } + + } catch (final GeoIp2Exception e) { + throw new EnrichFailedException("Address not found in database."); + } catch (final IOException e) { + throw new EngineFailureException("Failed to close database readers gracefully. It can be due to expired databases."); + } + return geoData; + } + + private void processCountryResponse(final CountryResponse countryResponse, final Map geoData, final List fields) { + final Continent continent = countryResponse.getContinent(); + final Country country = countryResponse.getCountry(); + final Country registeredCountry = countryResponse.getRegisteredCountry(); + final RepresentedCountry representedCountry = countryResponse.getRepresentedCountry(); + + + extractContinentFields(continent, geoData, fields); + extractCountryFields(country, geoData, fields, false); + extractRegisteredCountryFields(registeredCountry, geoData, fields); + extractRepresentedCountryFields(representedCountry, geoData, fields); + } + + private void processCityResponse(final CityResponse cityResponse, + final Map geoData, + final List fields, + final Set geoIPDatabases) { + // Continent and Country fields are added from City database only if they are not extracted from Country database + if (!geoIPDatabases.contains(GeoIPDatabase.COUNTRY)) { + final Continent continent = cityResponse.getContinent(); + final Country country = cityResponse.getCountry(); + final Country registeredCountry = cityResponse.getRegisteredCountry(); + final RepresentedCountry representedCountry = cityResponse.getRepresentedCountry(); + + extractContinentFields(continent, geoData, fields); + extractCountryFields(country, geoData, fields, false); + extractRegisteredCountryFields(registeredCountry, geoData, fields); + extractRepresentedCountryFields(representedCountry, geoData, fields); + } + + final City city = cityResponse.getCity(); + final Location location = cityResponse.getLocation(); + final Postal postal = cityResponse.getPostal(); + final Subdivision mostSpecificSubdivision = cityResponse.getMostSpecificSubdivision(); + final Subdivision leastSpecificSubdivision = cityResponse.getLeastSpecificSubdivision(); + + extractCityFields(city, geoData, fields, false); + extractLocationFields(location, geoData, fields); + extractPostalFields(postal, geoData, fields, false); + extractMostSpecifiedSubdivisionFields(mostSpecificSubdivision, geoData, fields, false); + extractLeastSpecifiedSubdivisionFields(leastSpecificSubdivision, geoData, fields, false); + } + + private void processAsnResponse(final AsnResponse asnResponse, final Map geoData, final List fields) { + extractAsnFields(asnResponse, geoData, fields); + } + + @Override + public void retain() { + + } + + @Override + public void close() { + closeReaders(); + } + + @Override + public boolean isExpired() { + final Instant instant = Instant.now(); + return isDatabaseExpired(instant, countryDatabaseReader, isCountryDatabaseExpired, countryDatabaseBuildDate, GEOLITE2_COUNTRY) && + isDatabaseExpired(instant, cityDatabaseReader, isCityDatabaseExpired, cityDatabaseBuildDate, GEOLITE2_CITY) && + isDatabaseExpired(instant, asnDatabaseReader, isAsnDatabaseExpired, asnDatabaseBuildDate, GEOLITE2_ASN); + } + + private boolean isDatabaseExpired(final Instant instant, + final DatabaseReader databaseReader, + final AtomicBoolean isDatabaseExpired, + final Instant databaseBuildDate, + final String databaseName) { + if (databaseReader == null) { + // no need to delete - no action needed + return true; + } + if (isDatabaseExpired.get()) { + // Another thread already updated status to expired - no action needed + return true; + } + if (databaseBuildDate.plus(MAX_EXPIRY_DURATION).isBefore(instant)) { + isDatabaseExpired.set(true); + closeReader(databaseReader, databaseName); + } + return isDatabaseExpired.get(); + } + + private void closeReaders() { + try { + if (cityDatabaseReader != null) { + cityDatabaseReader.close(); + } + if (countryDatabaseReader != null) { + countryDatabaseReader.close(); + } + if (asnDatabaseReader != null) { + asnDatabaseReader.close(); + } + } catch (final IOException e) { + LOG.debug("Failed to close Maxmind database readers due to: {}. Force closing readers.", e.getMessage()); + } + + // delete database directory + final File file = new File(databasePath); + geoIPFileManager.deleteDirectory(file); + } + + private void closeReader(final DatabaseReader databaseReader, final String databaseName) { + try { + if (databaseReader != null) { + databaseReader.close(); + } + } catch (final IOException e) { + LOG.debug("Failed to close Maxmind database readers due to: {}. Force closing readers.", e.getMessage()); + } + + // delete database file + final Optional fileName = getDatabaseName(databaseName, databasePath, MAXMIND_GEOLITE2_DATABASE_TYPE); + fileName.ifPresent(response -> { + File file = new File(databasePath + File.separator + response); + geoIPFileManager.deleteFile(file); + }); + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoData.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoData.java deleted file mode 100644 index 3730be7990..0000000000 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoData.java +++ /dev/null @@ -1,65 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.databaseenrich; - -import java.net.InetAddress; -import java.util.HashMap; -import java.util.List; -import java.util.Map; - -/** - * Interface for Downloading through S3 or URl or from local path - */ -public interface GetGeoData { - - public final String GeoLite2CityDB = "GeoLite2-City.mmdb"; - public final String GeoLite2CountryDB = "GeoLite2-Country.mmdb"; - public final String GeoLite2AsnDB = "GeoLite2-ASN.mmdb"; - public final String GeoIP2EnterpriseDB = "GeoIP2-Enterprise.mmdb"; - - void switchDatabaseReader(); - void closeReader(); - Map getGeoData(InetAddress inetAddress, List attributes, String tempDestDir); - - /** - * Enrich attributes - * @param geoData geoData - * @param attributeName attributeName - * @param attributeValue attributeValue - */ - default void enrichData(final Map geoData, final String attributeName, final String attributeValue) { - if (attributeValue != null) { - geoData.put(attributeName, attributeValue); - } - } - - /** - * Enrich region iso code - * @param geoData geoData - * @param countryIso countryIso - * @param subdivisionIso subdivisionIso - */ - default void enrichRegionIsoCode(final Map geoData, final String countryIso, final String subdivisionIso) { - if (countryIso != null && subdivisionIso != null) { - enrichData(geoData, "region_iso_code", countryIso + "-" + subdivisionIso); - } - } - - /** - * Enrich Location Data - * @param geoData geoData - * @param latitude latitude - * @param longitude longitude - */ - default void enrichLocationData(final Map geoData, final Double latitude, final Double longitude) { - if (latitude != null && longitude != null) { - Map locationObject = new HashMap<>(); - locationObject.put("lat", latitude); - locationObject.put("lon", longitude); - geoData.put("location", locationObject); - } - } -} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoIP2Data.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoIP2Data.java deleted file mode 100644 index 70d888f587..0000000000 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoIP2Data.java +++ /dev/null @@ -1,191 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.databaseenrich; - -import com.maxmind.geoip2.DatabaseReader; -import com.maxmind.geoip2.exception.GeoIp2Exception; -import com.maxmind.geoip2.model.EnterpriseResponse; -import com.maxmind.geoip2.record.City; -import com.maxmind.geoip2.record.Continent; -import com.maxmind.geoip2.record.Country; -import com.maxmind.geoip2.record.Subdivision; -import com.maxmind.geoip2.record.Location; -import com.maxmind.geoip2.record.Postal; -import org.opensearch.dataprepper.plugins.processor.extension.GeoIPProcessorService; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DatabaseReaderCreate; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSource; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.io.File; -import java.io.IOException; -import java.net.InetAddress; -import java.nio.file.Path; -import java.util.HashMap; -import java.util.List; -import java.util.Map; - -/** - * Implementation class for enrichment of enterprise data - */ -public class GetGeoIP2Data implements GetGeoData { - - private static final Logger LOG = LoggerFactory.getLogger(GetGeoIP2Data.class); - public static final String COUNTRY_NAME = "country_name"; - public static final String CONTINENT_NAME = "continent_name"; - public static final String REGION_NAME = "region_name"; - public static final String CITY_NAME = "city_name"; - public static final String COUNTRY_ISO_CODE = "country_iso_code"; - public static final String IP = "ip"; - public static final String REGION_ISO_CODE = "region_iso_code"; - public static final String TIMEZONE = "timezone"; - public static final String LOCATION = "location"; - public static final String POSTAL = "postal"; - private DatabaseReader readerEnterprise; - private Country country; - private Continent continent; - private City city; - private Location location; - private Subdivision subdivision; - private String dbPath; - private int cacheSize; - private Postal postal; - private String tempDestDir; - - /** - * GetGeoLite2Data constructor for initialisation of attributes - * @param dbPath dbPath - * @param cacheSize cacheSize - */ - public GetGeoIP2Data(final String dbPath, final int cacheSize) { - this.dbPath = dbPath; - this.cacheSize = cacheSize; - initDatabaseReader(); - } - - /** - * Initialise all the DatabaseReader - */ - public void initDatabaseReader() { - try { - readerEnterprise = DatabaseReaderCreate.createLoader(Path.of(dbPath + File.separator + tempDestDir + File.separator + GeoIP2EnterpriseDB), cacheSize); - } catch (final IOException ex) { - LOG.error("Exception while creating GeoIP2 DatabaseReader: {0}", ex); - } - } - - /** - * Switch all the DatabaseReader - */ - @Override - public void switchDatabaseReader() { - LOG.info("Switching GeoIP2 DatabaseReader"); - closeReader(); - System.gc(); - File file = new File(dbPath); - DBSource.deleteDirectory(file); - initDatabaseReader(); - } - - /** - * Enrich the GeoData - * @param inetAddress inetAddress - * @param attributes attributes - * @return enriched data Map - */ - public Map getGeoData(InetAddress inetAddress, List attributes, String tempDestDir) { - Map geoData = new HashMap<>(); - if (GeoIPProcessorService.downloadReady) { - this.tempDestDir = tempDestDir; - GeoIPProcessorService.downloadReady = false; - switchDatabaseReader(); - } - try { - EnterpriseResponse enterpriseResponse = readerEnterprise.enterprise(inetAddress); - country = enterpriseResponse.getCountry(); - subdivision = enterpriseResponse.getMostSpecificSubdivision(); - city = enterpriseResponse.getCity(); - location = enterpriseResponse.getLocation(); - continent = enterpriseResponse.getContinent(); - postal = enterpriseResponse.getPostal(); - } catch (IOException | GeoIp2Exception ex) { - LOG.info("Look up Exception : {0}", ex); - } - - try { - if ((attributes != null) && (!attributes.isEmpty())) { - for (String attribute : attributes) { - switch (attribute) { - case IP: - enrichData(geoData, IP, inetAddress.getHostAddress()); - break; - case COUNTRY_ISO_CODE: - enrichData(geoData, COUNTRY_ISO_CODE, country.getIsoCode()); - break; - case COUNTRY_NAME: - enrichData(geoData, COUNTRY_NAME, country.getName()); - break; - case CONTINENT_NAME: - enrichData(geoData, CONTINENT_NAME, continent.getName()); - break; - case REGION_ISO_CODE: - // ISO 3166-2 code for country subdivisions. - // See iso.org/iso-3166-country-codes.html - enrichRegionIsoCode(geoData, country.getIsoCode(), subdivision.getIsoCode()); - break; - case REGION_NAME: - enrichData(geoData, REGION_NAME, subdivision.getName()); - break; - case CITY_NAME: - enrichData(geoData, CITY_NAME, city.getName()); - break; - case TIMEZONE: - enrichData(geoData, TIMEZONE, location.getTimeZone()); - break; - case LOCATION: - enrichLocationData(geoData, location.getLatitude(), location.getLongitude()); - break; - case POSTAL: - enrichData(geoData, "postalCode", postal.getCode()); - break; - } - } - } else { - - enrichData(geoData, IP, inetAddress.getHostAddress()); - enrichData(geoData, COUNTRY_ISO_CODE, country.getIsoCode()); - enrichData(geoData, COUNTRY_NAME, country.getName()); - enrichData(geoData, CONTINENT_NAME, continent.getName()); - - enrichRegionIsoCode(geoData, country.getIsoCode(), subdivision.getIsoCode()); - - enrichData(geoData, REGION_NAME, subdivision.getName()); - enrichData(geoData, CITY_NAME, city.getName()); - enrichData(geoData, "postalCode", postal.getCode()); - - enrichData(geoData, TIMEZONE, location.getTimeZone()); - enrichLocationData(geoData, location.getLatitude(), location.getLongitude()); - } - } catch (Exception ex) { - throw new EnrichFailedException("Enrichment failed exception" + ex); - } - return geoData; - } - - - /** - * Close the DatabaseReader - */ - @Override - public void closeReader() { - try { - if (readerEnterprise != null) - readerEnterprise.close(); - } catch (IOException ex) { - LOG.info("Close Enterprise DatabaseReader Exception : {0}", ex); - } - } -} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoLite2Data.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoLite2Data.java deleted file mode 100644 index 0d719541e0..0000000000 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoLite2Data.java +++ /dev/null @@ -1,262 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.databaseenrich; - -import com.maxmind.db.Network; -import com.maxmind.geoip2.DatabaseReader; -import com.maxmind.geoip2.exception.GeoIp2Exception; -import com.maxmind.geoip2.model.AsnResponse; -import com.maxmind.geoip2.model.CityResponse; -import com.maxmind.geoip2.model.CountryResponse; -import com.maxmind.geoip2.record.City; -import com.maxmind.geoip2.record.Continent; -import com.maxmind.geoip2.record.Country; -import com.maxmind.geoip2.record.Subdivision; -import com.maxmind.geoip2.record.Location; -import org.opensearch.dataprepper.plugins.processor.extension.GeoIPProcessorService; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSource; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DatabaseReaderCreate; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.io.File; -import java.io.IOException; -import java.net.InetAddress; -import java.nio.file.Path; -import java.util.HashMap; -import java.util.List; -import java.util.Map; - -/** - * Implementation class for enrichment of geoip lite2 data - */ -public class GetGeoLite2Data implements GetGeoData { - - private static final Logger LOG = LoggerFactory.getLogger(GetGeoLite2Data.class); - public static final String COUNTRY_NAME = "country_name"; - public static final String CONTINENT_NAME = "continent_name"; - public static final String REGION_NAME = "region_name"; - public static final String CITY_NAME = "city_name"; - public static final String ORGANIZATION_NAME = "organization_name"; - public static final String NETWORK = "network"; - public static final String COUNTRY_ISO_CODE = "country_iso_code"; - public static final String IP = "ip"; - public static final String REGION_ISO_CODE = "region_iso_code"; - public static final String TIMEZONE = "timezone"; - public static final String LOCATION = "location"; - public static final String ASN = "asn"; - private DatabaseReader readerCity; - private DatabaseReader readerCountry; - private DatabaseReader readerAsn; - private Country country; - private Continent continent; - private City city; - private Location location; - private Subdivision subdivision; - private Long asn; - private String organizationName; - private Network network; - private String dbPath; - private int cacheSize; - private CityResponse responseCity; - private CountryResponse responseCountry; - private AsnResponse responseAsn; - private String tempDestDir; - - - /** - * GetGeoLite2Data constructor for initialisation of attributes - * @param dbPath dbPath - * @param cacheSize cacheSize - */ - public GetGeoLite2Data(final String dbPath, final int cacheSize) { - this.dbPath = dbPath; - this.cacheSize = cacheSize; - initDatabaseReader(); - } - - /** - * Initialise all the DatabaseReader - */ - private void initDatabaseReader() { - try { - readerCity = DatabaseReaderCreate.createLoader(Path.of(dbPath + File.separator + GeoLite2CityDB), cacheSize); - readerCountry = DatabaseReaderCreate.createLoader(Path.of(dbPath + File.separator + GeoLite2CountryDB), cacheSize); - readerAsn = DatabaseReaderCreate.createLoader(Path.of(dbPath + File.separator + GeoLite2AsnDB), cacheSize); - } catch (final IOException ex) { - LOG.error("Exception while creating GeoLite2 DatabaseReader: {0}", ex); - } - } - - /** - * Switch all the DatabaseReader - */ - @Override - public void switchDatabaseReader() { - LOG.info("Switching GeoLite2 DatabaseReaders"); - closeReaderCity(); - closeReaderCountry(); - closeReaderAsn(); - System.gc(); - final File file = new File(dbPath); - DBSource.deleteDirectory(file); - dbPath = tempDestDir; - initDatabaseReader(); - } - - /** - * Enrich the GeoData - * @param inetAddress inetAddress - * @param attributes attributes - * @return enriched data Map - */ - @Override - public Map getGeoData(final InetAddress inetAddress, final List attributes, final String tempDestDir) { - Map geoData = new HashMap<>(); - if (GeoIPProcessorService.downloadReady) { - this.tempDestDir = tempDestDir; - GeoIPProcessorService.downloadReady = false; - switchDatabaseReader(); - } - try { - responseCountry = readerCountry.country(inetAddress); - country = responseCountry.getCountry(); - continent = responseCountry.getContinent(); - - responseCity = readerCity.city(inetAddress); - city = responseCity.getCity(); - location = responseCity.getLocation(); - subdivision = responseCity.getMostSpecificSubdivision(); - - responseAsn = readerAsn.asn(inetAddress); - asn = responseAsn.getAutonomousSystemNumber(); - organizationName = responseAsn.getAutonomousSystemOrganization(); - network = responseAsn.getNetwork(); - } catch (IOException | GeoIp2Exception ex) { - LOG.info("Look up Exception : {0}", ex); - } - - try { - if ((attributes != null) && (!attributes.isEmpty())) { - for (String attribute : attributes) { - switch (attribute) { - case IP: - enrichData(geoData, IP, inetAddress.getHostAddress()); - break; - case COUNTRY_ISO_CODE: - enrichData(geoData, COUNTRY_ISO_CODE, country.getIsoCode()); - break; - case COUNTRY_NAME: - enrichData(geoData, COUNTRY_NAME, country.getName()); - break; - case CONTINENT_NAME: - enrichData(geoData, CONTINENT_NAME, continent.getName()); - break; - case REGION_ISO_CODE: - // ISO 3166-2 code for country subdivisions. - // See iso.org/iso-3166-country-codes.html - enrichRegionIsoCode(geoData, country.getIsoCode(), subdivision.getIsoCode()); - break; - case REGION_NAME: - enrichData(geoData, REGION_NAME, subdivision.getName()); - break; - case CITY_NAME: - enrichData(geoData, CITY_NAME, city.getName()); - break; - case TIMEZONE: - enrichData(geoData, TIMEZONE, location.getTimeZone()); - break; - case LOCATION: - enrichLocationData(geoData, location.getLatitude(), location.getLongitude()); - break; - case ASN: - if (asn != null) { - geoData.put(ASN, asn); - } - break; - case ORGANIZATION_NAME: - enrichData(geoData, ORGANIZATION_NAME, organizationName); - break; - case NETWORK: - enrichData(geoData, NETWORK,network!=null? network.toString():null); - break; - } - } - } else { - - enrichData(geoData, IP, inetAddress.getHostAddress()); - enrichData(geoData, COUNTRY_ISO_CODE, country.getIsoCode()); - enrichData(geoData, COUNTRY_NAME, country.getName()); - enrichData(geoData, CONTINENT_NAME, continent.getName()); - enrichRegionIsoCode(geoData, country.getIsoCode(), subdivision.getIsoCode()); - - enrichData(geoData, REGION_NAME, subdivision.getName()); - enrichData(geoData, CITY_NAME, city.getName()); - enrichData(geoData, TIMEZONE, location.getTimeZone()); - enrichLocationData(geoData, location.getLatitude(), location.getLongitude()); - - if (asn != null) { - geoData.put(ASN, asn); - } - - enrichData(geoData, ORGANIZATION_NAME, organizationName); - enrichData(geoData, NETWORK,network!=null? network.toString():null); - } - } catch (Exception ex) { - throw new EnrichFailedException("Enrichment failed exception" + ex); - } - return geoData; - } - - /** - * Close the all DatabaseReader - */ - @Override - public void closeReader() { - closeReaderCity(); - closeReaderCountry(); - closeReaderAsn(); - } - - /** - * Close the City DatabaseReader - */ - private void closeReaderCity() { - try { - if (readerCity != null) - readerCity.close(); - readerCity = null; - } catch (IOException ex) { - LOG.info("Close City DatabaseReader Exception : {0}", ex); - } - } - - /** - * Close the Country DatabaseReader - */ - private void closeReaderCountry() { - try { - if (readerCountry != null) - readerCountry.close(); - readerCountry = null; - } catch (IOException ex) { - LOG.info("Close Country DatabaseReader Exception : {0}", ex); - } - } - - /** - * Close the ASN DatabaseReader - */ - private void closeReaderAsn() { - try { - if (readerAsn != null) - readerAsn.close(); - readerAsn = null; - } catch (IOException ex) { - LOG.info("Close Asn DatabaseReader Exception : {0}", ex); - } - } -} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/DatabaseReaderInitializationException.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/DatabaseReaderInitializationException.java new file mode 100644 index 0000000000..f51f249220 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/DatabaseReaderInitializationException.java @@ -0,0 +1,12 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +public class DatabaseReaderInitializationException extends EngineFailureException { + public DatabaseReaderInitializationException(final String exceptionMsg) { + super(exceptionMsg); + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/DownloadFailedException.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/DownloadFailedException.java similarity index 50% rename from data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/DownloadFailedException.java rename to data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/DownloadFailedException.java index b667c3273e..05cf8d9613 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/DownloadFailedException.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/DownloadFailedException.java @@ -3,13 +3,13 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.dataprepper.plugins.processor.databaseenrich; +package org.opensearch.dataprepper.plugins.processor.exception; /** * Implementation class for DownloadFailedException Custom exception */ -public class DownloadFailedException extends RuntimeException { - public DownloadFailedException(String exceptionMsg) { +public class DownloadFailedException extends EngineFailureException { + public DownloadFailedException(final String exceptionMsg) { super(exceptionMsg); } } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/EngineFailureException.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/EngineFailureException.java new file mode 100644 index 0000000000..f5958b521a --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/EngineFailureException.java @@ -0,0 +1,12 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +public class EngineFailureException extends RuntimeException { + public EngineFailureException(final String exceptionMsg) { + super(exceptionMsg); + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/EnrichFailedException.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/EnrichFailedException.java similarity index 67% rename from data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/EnrichFailedException.java rename to data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/EnrichFailedException.java index 9182d83979..0f3dd2bfba 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/EnrichFailedException.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/EnrichFailedException.java @@ -3,13 +3,13 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.dataprepper.plugins.processor.databaseenrich; +package org.opensearch.dataprepper.plugins.processor.exception; /** * Implementation class for EnrichFailedException Custom exception */ public class EnrichFailedException extends RuntimeException { - public EnrichFailedException(String exceptionMsg) { + public EnrichFailedException(final String exceptionMsg) { super(exceptionMsg); } } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/InvalidIPAddressException.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/InvalidIPAddressException.java new file mode 100644 index 0000000000..134941b5ad --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/InvalidIPAddressException.java @@ -0,0 +1,12 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +public class InvalidIPAddressException extends EnrichFailedException { + public InvalidIPAddressException(final String exceptionMsg) { + super(exceptionMsg); + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/NoValidDatabaseFoundException.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/NoValidDatabaseFoundException.java new file mode 100644 index 0000000000..218d8c0901 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/exception/NoValidDatabaseFoundException.java @@ -0,0 +1,12 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +public class NoValidDatabaseFoundException extends EngineFailureException { + public NoValidDatabaseFoundException(final String exceptionMsg) { + super(exceptionMsg); + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/DefaultGeoIpConfigSupplier.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/DefaultGeoIpConfigSupplier.java index 08e50381bf..6c07304171 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/DefaultGeoIpConfigSupplier.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/DefaultGeoIpConfigSupplier.java @@ -5,15 +5,30 @@ package org.opensearch.dataprepper.plugins.processor.extension; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPDatabaseManager; + +import java.util.Optional; +import java.util.concurrent.locks.ReentrantReadWriteLock; + public class DefaultGeoIpConfigSupplier implements GeoIpConfigSupplier { private final GeoIpServiceConfig geoIpServiceConfig; + private final ReentrantReadWriteLock.ReadLock readLock; + private final GeoIPDatabaseManager geoIPDatabaseManager; - public DefaultGeoIpConfigSupplier(final GeoIpServiceConfig geoIpServiceConfig) { + public DefaultGeoIpConfigSupplier(final GeoIpServiceConfig geoIpServiceConfig, + final GeoIPDatabaseManager geoIPDatabaseManager, + final ReentrantReadWriteLock.ReadLock readLock + ) { this.geoIpServiceConfig = geoIpServiceConfig; + this.geoIPDatabaseManager = geoIPDatabaseManager; + this.readLock = readLock; } @Override - public GeoIPProcessorService getGeoIPProcessorService() { - return new GeoIPProcessorService(geoIpServiceConfig); + public Optional getGeoIPProcessorService() { + if (geoIpServiceConfig != null) + return Optional.of(new GeoIPProcessorService(geoIpServiceConfig, geoIPDatabaseManager, readLock)); + else + return Optional.empty(); } } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIPProcessorService.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIPProcessorService.java index 7df7af684a..bfa4150868 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIPProcessorService.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIPProcessorService.java @@ -5,138 +5,86 @@ package org.opensearch.dataprepper.plugins.processor.extension; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSourceOptions; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoDataFactory; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.S3DBService; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSource; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.HttpDBDownloadService; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.LocalDBDownloadService; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.DownloadFailedException; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.GetGeoData; -import org.opensearch.dataprepper.plugins.processor.utils.DbSourceIdentification; -import org.opensearch.dataprepper.plugins.processor.utils.LicenseTypeCheck; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPDatabaseManager; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.io.File; -import java.net.InetAddress; -import java.time.Duration; -import java.util.List; -import java.util.Map; -import java.util.Objects; +import java.time.Instant; +import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.ReentrantReadWriteLock; /** * Implementation class of geoIP-processor plugin service class. * It is responsible for calling of mmdb files download */ public class GeoIPProcessorService { - private static final Logger LOG = LoggerFactory.getLogger(GeoIPProcessorService.class); - public static final String DATABASE_1 = "first_database_path"; - public static final String DATABASE_2 = "second_database_path"; - private static final String TEMP_PATH_FOLDER = "geoip"; - private final GeoDataFactory geoDataFactory; - private GetGeoData geoData; - private List databasePaths; - private final String tempPath; - private final ScheduledExecutorService scheduledExecutorService; - private final DBSourceOptions dbSourceOptions; private final MaxMindConfig maxMindConfig; - private final LicenseTypeCheck licenseTypeCheck; - public static volatile boolean downloadReady; - private boolean toggle; - private String flipDatabase; - private boolean isDuringInitialization; + private final GeoIPDatabaseManager geoIPDatabaseManager; + private final ReentrantReadWriteLock.ReadLock readLock; + private ExecutorService executorService = null; /** * GeoIPProcessorService constructor for initialization of required attributes * * @param geoIpServiceConfig geoIpServiceConfig */ - public GeoIPProcessorService(final GeoIpServiceConfig geoIpServiceConfig) { - this.toggle = false; + public GeoIPProcessorService(final GeoIpServiceConfig geoIpServiceConfig, + final GeoIPDatabaseManager geoIPDatabaseManager, + final ReentrantReadWriteLock.ReadLock readLock + ) { this.maxMindConfig = geoIpServiceConfig.getMaxMindConfig(); - this.databasePaths = maxMindConfig.getDatabasePaths(); - this.isDuringInitialization = true; - flipDatabase = DATABASE_1; - - licenseTypeCheck = new LicenseTypeCheck(); - geoDataFactory = new GeoDataFactory(maxMindConfig, licenseTypeCheck); - this.tempPath = System.getProperty("java.io.tmpdir") + File.separator + TEMP_PATH_FOLDER; + this.geoIPDatabaseManager = geoIPDatabaseManager; + this.readLock = readLock; - dbSourceOptions = DbSourceIdentification.getDatabasePathType(databasePaths); - final Duration checkInterval = Objects.requireNonNull(maxMindConfig.getDatabaseRefreshInterval()); - scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(); - scheduledExecutorService - .scheduleAtFixedRate(this::downloadThroughURLandS3, 0L, checkInterval.toSeconds(), TimeUnit.SECONDS); + try { + geoIPDatabaseManager.initiateDatabaseDownload(); + } catch (final Exception e) { + LOG.error("Failed to initialize geoip processor due to: {}. Will update with backoff.", e.getMessage()); + } + } - synchronized (this) { - try { - while (!downloadReady) { - wait(); - } - } catch (final InterruptedException ex) { - LOG.info("Thread interrupted while waiting for download to complete: {0}", ex); - Thread.currentThread().interrupt(); - } - if (downloadReady) { - geoData = geoDataFactory.create(flipDatabase); + public GeoIPDatabaseReader getGeoIPDatabaseReader() { + readLock.lock(); + try { + final GeoIPDatabaseReader geoIPDatabaseReader = geoIPDatabaseManager.getGeoIPDatabaseReader(); + if (geoIPDatabaseReader != null) { + geoIPDatabaseReader.retain(); } + checkAndUpdateDatabases(); + return geoIPDatabaseReader; + } catch (final Exception e) { + LOG.error("Failed to update databases: {}", e.getMessage()); + return null; + } + finally { + readLock.unlock(); } - downloadReady = false; } - /** - * Calling download method based on the database path type - */ - public synchronized void downloadThroughURLandS3() { - DBSource dbSource; - toggle = !toggle; - if (!toggle) { - flipDatabase = DATABASE_1; - } else { - flipDatabase = DATABASE_2; + private synchronized void checkAndUpdateDatabases() { + if (geoIPDatabaseManager.getNextUpdateAt().isBefore(Instant.now())) { + LOG.info("Trying to update geoip Database readers"); + geoIPDatabaseManager.setNextUpdateAt(Instant.now().plus(maxMindConfig.getDatabaseRefreshInterval())); + executorService = Executors.newSingleThreadExecutor(); + executorService.execute(geoIPDatabaseManager::updateDatabaseReader); + executorService.shutdown(); } + } - try { - switch (dbSourceOptions) { - case URL: - dbSource = new HttpDBDownloadService(flipDatabase); - dbSource.initiateDownload(databasePaths); - downloadReady = true; - break; - case S3: - dbSource = new S3DBService(maxMindConfig.getAwsAuthenticationOptionsConfig(), flipDatabase); - dbSource.initiateDownload(databasePaths); - downloadReady = true; - break; - case PATH: - dbSource = new LocalDBDownloadService(flipDatabase); - dbSource.initiateDownload(databasePaths); - downloadReady = true; - break; - } - } catch (final Exception ex) { - if (isDuringInitialization) { - throw new DownloadFailedException("Download failed due to: " + ex); - } else { - LOG.error("Download failed due to: {0}. Using previously loaded database files.", ex); + public void shutdown() { + if (executorService != null) { + try { + if (!executorService.awaitTermination(30, TimeUnit.SECONDS)) { + executorService.shutdownNow(); + } + } catch (final InterruptedException e) { + executorService.shutdownNow(); } } - isDuringInitialization = false; - notifyAll(); - } - - /** - * Method to call enrichment of data based on license type - * @param inetAddress inetAddress - * @param attributes attributes - * @return Enriched Map - */ - public Map getGeoData(final InetAddress inetAddress, final List attributes) { - return geoData.getGeoData(inetAddress, attributes, tempPath + File.separator + flipDatabase); + geoIPDatabaseManager.deleteDatabasesOnShutdown(); } } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigExtension.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigExtension.java index 200541339d..6f3b9d9b0d 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigExtension.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigExtension.java @@ -9,18 +9,35 @@ import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; import org.opensearch.dataprepper.model.plugin.ExtensionPlugin; import org.opensearch.dataprepper.model.plugin.ExtensionPoints; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPFileManager; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DatabaseReaderBuilder; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPDatabaseManager; +import org.opensearch.dataprepper.plugins.processor.utils.LicenseTypeCheck; -@DataPrepperExtensionPlugin(modelType = GeoIpServiceConfig.class, rootKeyJsonPath = "/geoip_service") +import java.util.concurrent.locks.ReentrantReadWriteLock; + +@DataPrepperExtensionPlugin(modelType = GeoIpServiceConfig.class, rootKeyJsonPath = "/geoip_service", allowInPipelineConfigurations = true) public class GeoIpConfigExtension implements ExtensionPlugin { private final DefaultGeoIpConfigSupplier defaultGeoIpConfigSupplier; @DataPrepperPluginConstructor public GeoIpConfigExtension(final GeoIpServiceConfig geoIpServiceConfig) { - this.defaultGeoIpConfigSupplier = new DefaultGeoIpConfigSupplier(geoIpServiceConfig != null ? geoIpServiceConfig : new GeoIpServiceConfig()); + final ReentrantReadWriteLock reentrantReadWriteLock = new ReentrantReadWriteLock(true); + GeoIPDatabaseManager geoIPDatabaseManager = null; + if (geoIpServiceConfig != null) { + geoIPDatabaseManager = new GeoIPDatabaseManager( + geoIpServiceConfig.getMaxMindConfig(), + new LicenseTypeCheck(), + new DatabaseReaderBuilder(), + new GeoIPFileManager(), + reentrantReadWriteLock.writeLock() + ); + } + this.defaultGeoIpConfigSupplier = new DefaultGeoIpConfigSupplier(geoIpServiceConfig, geoIPDatabaseManager, reentrantReadWriteLock.readLock()); } @Override public void apply(final ExtensionPoints extensionPoints) { - extensionPoints.addExtensionProvider(new GeoIpConfigProvider(this.defaultGeoIpConfigSupplier)); + extensionPoints.addExtensionProvider(new GeoIpConfigProvider(this.defaultGeoIpConfigSupplier)); } } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigSupplier.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigSupplier.java index a754c9a745..8e00ef6f3e 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigSupplier.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigSupplier.java @@ -5,6 +5,8 @@ package org.opensearch.dataprepper.plugins.processor.extension; +import java.util.Optional; + /** * Interface for supplying {@link GeoIPProcessorService} to {@link GeoIpConfigExtension} * @@ -16,5 +18,5 @@ public interface GeoIpConfigSupplier { * * @since 2.7 */ - GeoIPProcessorService getGeoIPProcessorService(); + Optional getGeoIPProcessorService(); } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpServiceConfig.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpServiceConfig.java index c311f773ca..3dc974e698 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpServiceConfig.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpServiceConfig.java @@ -11,12 +11,8 @@ public class GeoIpServiceConfig { private static final MaxMindConfig DEFAULT_MAXMIND_CONFIG = new MaxMindConfig(); - public GeoIpServiceConfig() { - // This default constructor is used if geoip_service is not configured - } - - @JsonProperty("maxmind") @Valid + @JsonProperty("maxmind") private MaxMindConfig maxMindConfig = DEFAULT_MAXMIND_CONFIG; /** diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindConfig.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindConfig.java index 3fea586e67..608d962d8e 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindConfig.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindConfig.java @@ -11,44 +11,54 @@ import jakarta.validation.constraints.Min; import org.hibernate.validator.constraints.time.DurationMax; import org.hibernate.validator.constraints.time.DurationMin; +import org.opensearch.dataprepper.plugins.processor.utils.DatabaseSourceIdentification; +import java.io.File; +import java.net.URI; +import java.net.URISyntaxException; import java.time.Duration; import java.util.ArrayList; import java.util.List; public class MaxMindConfig { + private static final boolean DEFAULT_INSECURE = false; private static final String S3_PREFIX = "s3://"; - - //TODO: Add validations to database paths - //TODO: Make default path to be a public CDN endpoint - private static final List DEFAULT_DATABASE_PATHS = new ArrayList<>(); private static final Duration DEFAULT_DATABASE_REFRESH_INTERVAL = Duration.ofDays(7); - private static final int DEFAULT_CACHE_SIZE = 4096; + private static final int DEFAULT_CACHE_COUNT = 4096; + static final String DEFAULT_DATABASE_DESTINATION = System.getProperty("data-prepper.dir") + File.separator + "data"; - @JsonProperty("database_paths") - private List databasePaths = DEFAULT_DATABASE_PATHS; + @Valid + @JsonProperty("databases") + private MaxMindDatabaseConfig maxMindDatabaseConfig = new MaxMindDatabaseConfig(); @JsonProperty("database_refresh_interval") @DurationMin(days = 1) @DurationMax(days = 30) private Duration databaseRefreshInterval = DEFAULT_DATABASE_REFRESH_INTERVAL; - @JsonProperty("cache_size") + @JsonProperty("cache_count") @Min(1) //TODO: Add a Max limit on cache size - private int cacheSize = DEFAULT_CACHE_SIZE; + private int cacheSize = DEFAULT_CACHE_COUNT; - //TODO: Add a destination path to store database files - @JsonProperty("aws") @Valid + @JsonProperty("aws") private AwsAuthenticationOptionsConfig awsAuthenticationOptionsConfig; + @JsonProperty("insecure") + private boolean insecure = DEFAULT_INSECURE; + + @JsonProperty("database_destination") + private String databaseDestination = DEFAULT_DATABASE_DESTINATION; + public MaxMindConfig() { // This default constructor is used if maxmind is not configured } @AssertTrue(message = "aws should be configured if any path in database_paths is S3 bucket path.") - boolean isAwsAuthenticationOptionsRequired() { + public boolean isAwsAuthenticationOptionsValid() { + final List databasePaths = new ArrayList<>(maxMindDatabaseConfig.getDatabasePaths().values()); + for (final String databasePath : databasePaths) { if (databasePath.startsWith(S3_PREFIX)) { return awsAuthenticationOptionsConfig != null; @@ -57,14 +67,28 @@ boolean isAwsAuthenticationOptionsRequired() { return true; } + @AssertTrue(message = "database_paths should be https endpoint if using URL and if insecure is set to false") + public boolean isHttpsEndpointOrInsecure() throws URISyntaxException { + if (insecure) { + return true; + } + final List databasePaths = new ArrayList<>(maxMindDatabaseConfig.getDatabasePaths().values()); + for (final String databasePath : databasePaths) { + if (DatabaseSourceIdentification.isURL(databasePath)) { + return new URI(databasePath).getScheme().equals("https"); + } + } + return true; + } + /** - * Gets the MaxMind database paths + * Gets Map of database name and database path * - * @return The MaxMind database paths + * @return Map * @since 2.7 */ - public List getDatabasePaths() { - return databasePaths; + public MaxMindDatabaseConfig getMaxMindDatabaseConfig() { + return maxMindDatabaseConfig; } /** @@ -96,4 +120,14 @@ public int getCacheSize() { public AwsAuthenticationOptionsConfig getAwsAuthenticationOptionsConfig() { return awsAuthenticationOptionsConfig; } + + /** + * Gets the destination folder to store database files + * + * @return The destination folder + * @since 2.7 + */ + public String getDatabaseDestination() { + return databaseDestination + File.separator + "geoip"; + } } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindDatabaseConfig.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindDatabaseConfig.java new file mode 100644 index 0000000000..7580c4000e --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindDatabaseConfig.java @@ -0,0 +1,76 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension; + +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.constraints.AssertTrue; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSourceOptions; +import org.opensearch.dataprepper.plugins.processor.utils.DatabaseSourceIdentification; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +public class MaxMindDatabaseConfig { + static final String DEFAULT_CITY_ENDPOINT = "https://geoip.maps.opensearch.org/v1/mmdb/geolite2-city/manifest.json"; + static final String DEFAULT_COUNTRY_ENDPOINT = "https://geoip.maps.opensearch.org/v1/mmdb/geolite2-country/manifest.json"; + static final String DEFAULT_ASN_ENDPOINT = "https://geoip.maps.opensearch.org/v1/mmdb/geolite2-asn/manifest.json"; + public static final String GEOLITE2_COUNTRY = "geolite2-country"; + public static final String GEOLITE2_CITY = "geolite2-city"; + public static final String GEOLITE2_ASN = "geolite2-asn"; + public static final String GEOIP2_ENTERPRISE = "geoip2-enterprise"; + private Map databases = null; + @JsonProperty("city") + private String cityDatabase; + + @JsonProperty("country") + private String countryDatabase; + + @JsonProperty("asn") + private String asnDatabase; + + @JsonProperty("enterprise") + private String enterpriseDatabase; + + @AssertTrue(message = "MaxMind GeoLite2 databases cannot be used along with enterprise database.") + public boolean isDatabasesValid() { + return enterpriseDatabase == null || (cityDatabase == null && countryDatabase == null && asnDatabase == null); + } + + @AssertTrue(message = "database_paths should be S3 URI or HTTP endpoint or local directory") + public boolean isPathsValid() { + final List databasePaths = new ArrayList<>(getDatabasePaths().values()); + + final DBSourceOptions dbSourceOptions = DatabaseSourceIdentification.getDatabasePathType(databasePaths); + return dbSourceOptions != null; + } + + public Map getDatabasePaths() { + if (databases == null) { + databases = new HashMap<>(); + if (countryDatabase == null && cityDatabase == null && asnDatabase == null && enterpriseDatabase == null) { + databases.put(GEOLITE2_COUNTRY, DEFAULT_COUNTRY_ENDPOINT); + databases.put(GEOLITE2_CITY, DEFAULT_CITY_ENDPOINT); + databases.put(GEOLITE2_ASN, DEFAULT_ASN_ENDPOINT); + } else { + if (countryDatabase != null) { + databases.put(GEOLITE2_COUNTRY, countryDatabase); + } + if (cityDatabase != null) { + databases.put(GEOLITE2_CITY, cityDatabase); + } + if (asnDatabase != null) { + databases.put(GEOLITE2_ASN, asnDatabase); + } + if (enterpriseDatabase != null) { + databases.put(GEOIP2_ENTERPRISE, enterpriseDatabase); + } + } + } + return databases; + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSource.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSource.java index 3ead66744c..88abe0b0f9 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSource.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSource.java @@ -5,66 +5,21 @@ package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - import javax.net.ssl.TrustManager; import javax.net.ssl.SSLContext; import javax.net.ssl.X509TrustManager; import javax.net.ssl.HttpsURLConnection; import javax.net.ssl.SSLSession; import javax.net.ssl.HostnameVerifier; -import java.io.File; -import java.io.UncheckedIOException; import java.security.KeyManagementException; import java.security.NoSuchAlgorithmException; import java.security.SecureRandom; import java.security.cert.CertificateException; import java.security.cert.X509Certificate; -import java.util.List; public interface DBSource { - - public static final Logger LOG = LoggerFactory.getLogger(DBSource.class); - public String tempFolderPath = System.getProperty("java.io.tmpdir")+ File.separator +"GeoIP"; - public String tarFolderPath = tempFolderPath + "/tar"; - public String downloadTarFilepath = tarFolderPath + "/out.tar.gz"; - public void initiateDownload(List config) throws Exception; - - /** - * create Folder If Not Exist - * @param outputFilePath Output File Path - * @return File - */ - static File createFolderIfNotExist(String outputFilePath) { - final File destFile = new File(outputFilePath); - try { - if (!destFile.exists()) { - destFile.mkdirs(); - } - } - catch (UncheckedIOException ex) { - LOG.info("Create Folder If NotExist Exception {0}", ex); - } - return destFile; - } - - /** - * Delete Directory - * @param file file - */ - static void deleteDirectory(File file) { - - if (file.exists()) { - for (final File subFile : file.listFiles()) { - if (subFile.isDirectory()) { - deleteDirectory(subFile); - } - subFile.delete(); - } - file.delete(); - } - } + String MAXMIND_DATABASE_EXTENSION = ".mmdb"; + void initiateDownload() throws Exception; /** * initiateSSL diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSourceOptions.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSourceOptions.java index 485274dd9d..60eacff72a 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSourceOptions.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSourceOptions.java @@ -17,7 +17,8 @@ public enum DBSourceOptions { PATH("path"), URL("url"), - S3("s3"); + S3("s3"), + HTTP_MANIFEST("http_manifest"); private final String option; diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderCreate.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderBuilder.java similarity index 80% rename from data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderCreate.java rename to data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderBuilder.java index 47546c6c9f..841b4c7fbc 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderCreate.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderBuilder.java @@ -15,7 +15,9 @@ /** * Implementation class for DatabaseReader Creation */ -public class DatabaseReaderCreate { +public class DatabaseReaderBuilder { + public DatabaseReaderBuilder() { + } /** * Creates DatabaseReader instance based on in memory or cache type @@ -23,7 +25,7 @@ public class DatabaseReaderCreate { * @param cacheSize cacheSize * @return DatabaseReader */ - public static DatabaseReader createLoader(final Path databasePath, final int cacheSize) throws IOException { + public DatabaseReader buildReader(final Path databasePath, final int cacheSize) throws IOException { return new DatabaseReader.Builder(databasePath.toFile()) .fileMode(Reader.FileMode.MEMORY_MAPPED) .withCache(new CHMCache(cacheSize)) diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoDataFactory.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoDataFactory.java deleted file mode 100644 index e05453d37f..0000000000 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoDataFactory.java +++ /dev/null @@ -1,39 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; - -import org.opensearch.dataprepper.plugins.processor.databaseenrich.GetGeoData; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.GetGeoIP2Data; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.GetGeoLite2Data; -import org.opensearch.dataprepper.plugins.processor.extension.MaxMindConfig; -import org.opensearch.dataprepper.plugins.processor.utils.LicenseTypeCheck; - -import java.io.File; - -import static org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSource.tempFolderPath; - -public class GeoDataFactory { - private final MaxMindConfig maxMindConfig; - private final LicenseTypeCheck licenseTypeCheck; - - public GeoDataFactory(final MaxMindConfig maxMindConfig, final LicenseTypeCheck licenseTypeCheck) { - this.maxMindConfig = maxMindConfig; - this.licenseTypeCheck = licenseTypeCheck; - } - - /** - * Creates GetGeoData class based on LicenseTypeOptions - */ - public GetGeoData create(final String databasePath) { - final String finalPath = tempFolderPath + File.separator + databasePath; - final LicenseTypeOptions licenseType = licenseTypeCheck.isGeoLite2OrEnterpriseLicense(finalPath); - if (licenseType.equals(LicenseTypeOptions.FREE)) { - return new GetGeoLite2Data(finalPath, maxMindConfig.getCacheSize()); - } else { - return new GetGeoIP2Data(finalPath, maxMindConfig.getCacheSize()); - } - } -} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPDatabaseManager.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPDatabaseManager.java new file mode 100644 index 0000000000..307b4b9276 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPDatabaseManager.java @@ -0,0 +1,213 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; + +import com.linecorp.armeria.client.retry.Backoff; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.AutoCountingDatabaseReader; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIP2DatabaseReader; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoLite2DatabaseReader; +import org.opensearch.dataprepper.plugins.processor.exception.DownloadFailedException; +import org.opensearch.dataprepper.plugins.processor.exception.NoValidDatabaseFoundException; +import org.opensearch.dataprepper.plugins.processor.extension.MaxMindConfig; +import org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig; +import org.opensearch.dataprepper.plugins.processor.utils.DatabaseSourceIdentification; +import org.opensearch.dataprepper.plugins.processor.utils.LicenseTypeCheck; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.ReentrantReadWriteLock; +import java.util.concurrent.locks.ReentrantReadWriteLock.WriteLock; + +public class GeoIPDatabaseManager { + private static final Logger LOG = LoggerFactory.getLogger(GeoIPDatabaseManager.class); + public static final String FIRST_DATABASE_DIR = "first_database"; + public static final String SECOND_DATABASE_DIR = "second_database"; + private static final long INITIAL_DELAY = Duration.ofMinutes(1).toMillis(); + private static final long MAXIMUM_DELAY = Duration.ofHours(1).toMillis(); + private static final double JITTER_RATE = 0.15; + private final MaxMindConfig maxMindConfig; + private final LicenseTypeCheck licenseTypeCheck; + private final DatabaseReaderBuilder databaseReaderBuilder; + private final MaxMindDatabaseConfig maxMindDatabaseConfig; + private final WriteLock writeLock; + private final int cacheSize; + private final GeoIPFileManager geoIPFileManager; + private final AtomicInteger failedAttemptCount; + private final Backoff backoff; + private final DBSourceOptions dbSourceOptions; + private String currentDatabaseDir; + private GeoIPDatabaseReader geoIPDatabaseReader; + private boolean databaseDirToggle; + private Instant nextUpdateAt; + + public GeoIPDatabaseManager(final MaxMindConfig maxMindConfig, + final LicenseTypeCheck licenseTypeCheck, + final DatabaseReaderBuilder databaseReaderBuilder, + final GeoIPFileManager geoIPFileManager, + final ReentrantReadWriteLock.WriteLock writeLock + ) { + this.maxMindConfig = maxMindConfig; + this.licenseTypeCheck = licenseTypeCheck; + this.databaseReaderBuilder = databaseReaderBuilder; + this.geoIPFileManager = geoIPFileManager; + this.maxMindDatabaseConfig = maxMindConfig.getMaxMindDatabaseConfig(); + this.writeLock = writeLock; + this.cacheSize = maxMindConfig.getCacheSize(); + this.failedAttemptCount = new AtomicInteger(0); + this.backoff = Backoff.exponential(INITIAL_DELAY, MAXIMUM_DELAY) + .withJitter(JITTER_RATE) + .withMaxAttempts(Integer.MAX_VALUE); + this.dbSourceOptions = DatabaseSourceIdentification.getDatabasePathType(new ArrayList<>( + maxMindDatabaseConfig.getDatabasePaths().values())); + + } + + public void initiateDatabaseDownload() { + try { + downloadDatabases(); + geoIPDatabaseReader = createReader(); + nextUpdateAt = Instant.now().plus(maxMindConfig.getDatabaseRefreshInterval()); + failedAttemptCount.set(0); + } catch (final Exception e) { + final Duration delay = Duration.ofMillis(applyBackoff()); + nextUpdateAt = Instant.now().plus(delay); + throw new DownloadFailedException(e.getMessage()); + } + } + + public void updateDatabaseReader() { + try { + downloadDatabases(); + switchDatabase(); + LOG.info("Updated geoip database readers"); + failedAttemptCount.set(0); + } catch (final Exception e) { + LOG.error("Failed to download database and create database readers, will try to use old databases if they exist. {}", e.getMessage()); + final Duration delay = Duration.ofMillis(applyBackoff()); + nextUpdateAt = Instant.now().plus(delay); + final File file = new File(currentDatabaseDir); + geoIPFileManager.deleteDirectory(file); + switchDirectory(); + } + } + + private void switchDatabase() { + writeLock.lock(); + try { + final GeoIPDatabaseReader newGeoipDatabaseReader = createReader(); + final GeoIPDatabaseReader oldGeoipDatabaseReader = geoIPDatabaseReader; + geoIPDatabaseReader = newGeoipDatabaseReader; + if (oldGeoipDatabaseReader != null) { + oldGeoipDatabaseReader.close(); + } + } catch (Exception e) { + LOG.error("Failed to close geoip database readers due to: {}", e.getMessage()); + } finally { + writeLock.unlock(); + } + } + + private void downloadDatabases() throws Exception { + DBSource dbSource; + switchDirectory(); + + final String destinationPath = maxMindConfig.getDatabaseDestination() + File.separator + currentDatabaseDir; + geoIPFileManager.createDirectoryIfNotExist(destinationPath); + switch (dbSourceOptions) { + case HTTP_MANIFEST: + dbSource = new ManifestDownloadService(destinationPath, maxMindDatabaseConfig); + dbSource.initiateDownload(); + break; + case URL: + dbSource = new HttpDBDownloadService(destinationPath, geoIPFileManager, maxMindDatabaseConfig); + dbSource.initiateDownload(); + break; + case S3: + dbSource = new S3DBService(maxMindConfig.getAwsAuthenticationOptionsConfig(), destinationPath, maxMindDatabaseConfig); + dbSource.initiateDownload(); + break; + case PATH: + dbSource = new LocalDBDownloadService(destinationPath, maxMindDatabaseConfig); + dbSource.initiateDownload(); + break; + } + } + + private GeoIPDatabaseReader createReader() { + final String finalPath = maxMindConfig.getDatabaseDestination() + File.separator + currentDatabaseDir; + final LicenseTypeOptions licenseType = licenseTypeCheck.isGeoLite2OrEnterpriseLicense(finalPath); + if (licenseType == null) { + throw new NoValidDatabaseFoundException("At least one valid database is required."); + } + GeoIPDatabaseReader newGeoIPDatabaseReader; + if (licenseType.equals(LicenseTypeOptions.FREE)) { + newGeoIPDatabaseReader = new AutoCountingDatabaseReader( + new GeoLite2DatabaseReader(databaseReaderBuilder, geoIPFileManager, finalPath, cacheSize)); + } else if (licenseType.equals(LicenseTypeOptions.ENTERPRISE)) { + newGeoIPDatabaseReader = new AutoCountingDatabaseReader( + new GeoIP2DatabaseReader(databaseReaderBuilder, geoIPFileManager, finalPath, cacheSize)); + } else { + throw new NoValidDatabaseFoundException("No valid database found to initialize database readers."); + } + return newGeoIPDatabaseReader; + } + + private void switchDirectory() { + databaseDirToggle = !databaseDirToggle; + if (databaseDirToggle) { + currentDatabaseDir = FIRST_DATABASE_DIR; + } else { + currentDatabaseDir = SECOND_DATABASE_DIR; + } + } + + private long applyBackoff() { + final long delayMillis = backoff.nextDelayMillis(failedAttemptCount.incrementAndGet()); + if (delayMillis < 0) { + // retries exhausted + LOG.info("Retries exhausted to download database. Will retry based on refresh interval"); + } + final Duration delayDuration = Duration.ofMillis(delayMillis); + LOG.info("Failed to download databases, will retry after {} seconds", delayDuration.getSeconds()); + return delayMillis; + } + + public GeoIPDatabaseReader getGeoIPDatabaseReader() { + return geoIPDatabaseReader; + } + + public Instant getNextUpdateAt() { + return nextUpdateAt; + } + + public void setNextUpdateAt(final Instant nextUpdateAt) { + this.nextUpdateAt = nextUpdateAt; + } + + public void deleteDatabasesOnShutdown() { + geoIPFileManager.deleteDirectory(new File(maxMindConfig.getDatabaseDestination() + File.separator + FIRST_DATABASE_DIR)); + geoIPFileManager.deleteDirectory(new File(maxMindConfig.getDatabaseDestination() + File.separator + SECOND_DATABASE_DIR)); + } + + public void deleteDirectory(final File file) { + + if (file.exists()) { + for (final File subFile : file.listFiles()) { + if (subFile.isDirectory()) { + deleteDirectory(subFile); + } + subFile.delete(); + } + file.delete(); + } + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPFileManager.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPFileManager.java new file mode 100644 index 0000000000..66dda2db70 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPFileManager.java @@ -0,0 +1,35 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; + +import java.io.File; + +public class GeoIPFileManager { + public void deleteDirectory(final File file) { + + if (file.exists()) { + for (final File subFile : file.listFiles()) { + if (subFile.isDirectory()) { + deleteDirectory(subFile); + } + subFile.delete(); + } + file.delete(); + } + } + + public void deleteFile(final File file) { + file.delete(); + } + + public void createDirectoryIfNotExist(final String outputFilePath) { + final File destFile = new File(outputFilePath); + if (!destFile.exists()) { + destFile.mkdirs(); + } + } + +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/HttpDBDownloadService.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/HttpDBDownloadService.java index d407f9bdb7..52d0c7d1a7 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/HttpDBDownloadService.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/HttpDBDownloadService.java @@ -8,8 +8,8 @@ import org.apache.commons.compress.archivers.tar.TarArchiveEntry; import org.apache.commons.compress.archivers.tar.TarArchiveInputStream; import org.apache.commons.compress.utils.IOUtils; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; +import org.opensearch.dataprepper.plugins.processor.exception.DownloadFailedException; +import org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig; import java.io.BufferedInputStream; import java.io.File; @@ -17,41 +17,48 @@ import java.io.IOException; import java.io.FileInputStream; import java.net.URL; -import java.util.List; +import java.util.Set; import java.util.zip.GZIPInputStream; /** * Implementation class for Download through Url */ public class HttpDBDownloadService implements DBSource { - - private static final Logger LOG = LoggerFactory.getLogger(HttpDBDownloadService.class); - private final String prefixDir; + private final String destinationDirectory; private static final int DEFAULT_BYTE_SIZE = 1024; + private final GeoIPFileManager geoIPFileManager; + private final MaxMindDatabaseConfig maxMindDatabaseConfig; /** * HttpDBDownloadService constructor for initialisation of attributes - * @param prefixDir prefixDir + * @param destinationDirectory destinationDirectory */ - public HttpDBDownloadService(String prefixDir) { - this.prefixDir = prefixDir; + public HttpDBDownloadService(final String destinationDirectory, + final GeoIPFileManager geoIPFileManager, + final MaxMindDatabaseConfig maxMindDatabaseConfig) { + this.destinationDirectory = destinationDirectory; + this.geoIPFileManager = geoIPFileManager; + this.maxMindDatabaseConfig = maxMindDatabaseConfig; } /** * Initialisation of Download through Url - * @param urlList urlList */ - public void initiateDownload(List urlList) { - final File tmpDir = DBSource.createFolderIfNotExist(tempFolderPath + File.separator + prefixDir); - for(String url : urlList) { - DBSource.createFolderIfNotExist(tarFolderPath); + public void initiateDownload() { + final String tarDir = destinationDirectory + File.separator + "tar"; + final String downloadTarFilepath = tarDir + File.separator + "out.tar.gz"; + final Set databasePaths = maxMindDatabaseConfig.getDatabasePaths().keySet(); + for (final String key: databasePaths) { + geoIPFileManager.createDirectoryIfNotExist(tarDir); try { initiateSSL(); - buildRequestAndDownloadFile(url); - decompressAndUntarFile(tarFolderPath, downloadTarFilepath, tmpDir); - deleteTarFolder(tarFolderPath); + buildRequestAndDownloadFile(maxMindDatabaseConfig.getDatabasePaths().get(key), downloadTarFilepath); + final File tarFile = decompressAndgetTarFile(tarDir, downloadTarFilepath); + unTarFile(tarFile, new File(destinationDirectory), key); + deleteTarFolder(tarDir); } catch (Exception ex) { - LOG.info("InitiateDownload Exception {0} " , ex); + throw new DownloadFailedException("Failed to download from " + maxMindDatabaseConfig.getDatabasePaths().get(key) + + " due to: " + ex.getMessage()); } } } @@ -60,19 +67,18 @@ public void initiateDownload(List urlList) { * Decompress and untar the file * @param tarFolderPath tarFolderPath * @param downloadTarFilepath downloadTarFilepath - * @param tmpDir tmpDir + * + * @return File Tar file */ - private void decompressAndUntarFile(String tarFolderPath, String downloadTarFilepath, File tmpDir) { + private File decompressAndgetTarFile(final String tarFolderPath, final String downloadTarFilepath) { try { final File inputFile = new File(downloadTarFilepath); final String outputFile = getFileName(inputFile, tarFolderPath); File tarFile = new File(outputFile); // Decompress file - tarFile = deCompressGZipFile(inputFile, tarFile); - // Untar file - unTarFile(tarFile, tmpDir); + return deCompressGZipFile(inputFile, tarFile); } catch (IOException ex) { - LOG.info("Decompress and untar the file Exception {0} " , ex); + throw new DownloadFailedException("Failed to decompress GZip file." + ex.getMessage()); } } @@ -80,17 +86,17 @@ private void decompressAndUntarFile(String tarFolderPath, String downloadTarFile * Build Request And DownloadFile * @param url url */ - public void buildRequestAndDownloadFile(String... url) { - downloadDBFileFromMaxmind(url[0], downloadTarFilepath); + public void buildRequestAndDownloadFile(final String url, final String downloadTarFilepath) { + downloadDBFileFromMaxmind(url, downloadTarFilepath); } /** * Delete Tar Folder * @param tarFolder Tar Folder */ - private static void deleteTarFolder(String tarFolder) { + private void deleteTarFolder(String tarFolder) { final File file = new File(tarFolder); - DBSource.deleteDirectory(file); + geoIPFileManager.deleteDirectory(file); if (file.exists()) { file.delete(); } @@ -110,7 +116,7 @@ private static void downloadDBFileFromMaxmind(String maxmindDownloadUrl, String fileOutputStream.write(dataBuffer, 0, bytesRead); } } catch (IOException ex) { - LOG.info("download DB File FromMaxmind Exception {0} " , ex); + throw new DownloadFailedException("Failed to download from " + maxmindDownloadUrl + " due to: " + ex.getMessage()); } } @@ -151,19 +157,21 @@ private static String getFileName(File inputFile, String outputFolder) { /** * unTarFile * @param tarFile tar File - * @param destFile dest File + * @param destDir dest directory + * @param fileName File name + * * @throws IOException ioexception */ - private static void unTarFile(File tarFile, File destFile) throws IOException { + private static void unTarFile(final File tarFile, final File destDir, final String fileName) throws IOException { final FileInputStream fileInputStream = new FileInputStream(tarFile); final TarArchiveInputStream tarArchiveInputStream = new TarArchiveInputStream(fileInputStream); TarArchiveEntry tarEntry = null; while ((tarEntry = tarArchiveInputStream.getNextTarEntry()) != null) { - if(tarEntry.getName().endsWith(".mmdb")) { - String fileName = destFile + File.separator + tarEntry.getName().split("/")[1]; - final File outputFile = new File(fileName); + if(tarEntry.getName().endsWith(MAXMIND_DATABASE_EXTENSION)) { + final File outputFile = new File(destDir + File.separator + fileName + MAXMIND_DATABASE_EXTENSION); + if (tarEntry.isDirectory()) { if (!outputFile.exists()) { outputFile.mkdirs(); diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/LocalDBDownloadService.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/LocalDBDownloadService.java index d4a85218fa..bda5dd59cd 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/LocalDBDownloadService.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/LocalDBDownloadService.java @@ -5,36 +5,38 @@ package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; -import org.apache.commons.io.FileUtils; +import com.google.common.io.Files; +import org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig; import java.io.File; -import java.util.List; +import java.util.Set; /** * Implementation class for Download through local path */ public class LocalDBDownloadService implements DBSource { - private final String prefixDir; + private final String destinationDirectory; + private final MaxMindDatabaseConfig maxMindDatabaseConfig; /** * LocalDBDownloadService constructor for initialisation of attributes - * @param prefixDir prefixDir + * @param destinationDirectory destinationDirectory */ - public LocalDBDownloadService(final String prefixDir) { - this.prefixDir = prefixDir; + public LocalDBDownloadService(final String destinationDirectory, final MaxMindDatabaseConfig maxMindDatabaseConfig) { + this.destinationDirectory = destinationDirectory; + this.maxMindDatabaseConfig = maxMindDatabaseConfig; } /** * Initialisation of Download from local file path - * @param config config */ @Override - public void initiateDownload(List config) throws Exception { - String destPath = tempFolderPath + File.separator + prefixDir; - DBSource.createFolderIfNotExist(destPath); - File srcDatabaseConfigPath = new File(config.get(0)); - File destDatabaseConfigPath = new File(destPath); - FileUtils.copyDirectory(srcDatabaseConfigPath, destDatabaseConfigPath); + public void initiateDownload() throws Exception { + final Set strings = maxMindDatabaseConfig.getDatabasePaths().keySet(); + for (final String key: strings) { + Files.copy(new File(maxMindDatabaseConfig.getDatabasePaths().get(key)), + new File(destinationDirectory + File.separator + key + MAXMIND_DATABASE_EXTENSION)); + } } } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/Manifest.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/Manifest.java new file mode 100644 index 0000000000..f29bd3c2fb --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/Manifest.java @@ -0,0 +1,47 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; + +import com.fasterxml.jackson.annotation.JsonProperty; + +public class Manifest { + @JsonProperty("url") + private String url; + @JsonProperty("db_name") + private String dbName; + @JsonProperty("sha256_hash") + private String sha256Hash; + @JsonProperty("valid_for_in_days") + private int validForInDays; + @JsonProperty("updated_at_in_epoch_milli") + private long updatedAt; + @JsonProperty("provider") + private String provider; + + public String getUrl() { + return url; + } + + public String getDbName() { + return dbName; + } + + public String getSha256Hash() { + return sha256Hash; + } + + public int getValidForInDays() { + return validForInDays; + } + + public Long getUpdatedAt() { + return updatedAt; + } + + public String getProvider() { + return provider; + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/ManifestDownloadService.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/ManifestDownloadService.java new file mode 100644 index 0000000000..809b68a713 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/ManifestDownloadService.java @@ -0,0 +1,127 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.opensearch.dataprepper.plugins.processor.exception.DownloadFailedException; +import org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig; + +import java.io.BufferedInputStream; +import java.io.File; +import java.io.FileInputStream; +import java.io.FileOutputStream; +import java.io.IOException; +import java.net.HttpURLConnection; +import java.net.URL; +import java.util.Set; +import java.util.zip.ZipEntry; +import java.util.zip.ZipInputStream; + +public class ManifestDownloadService implements DBSource { + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + private static final int DEFAULT_BYTE_SIZE = 1024; + private static final String ZIP_FILE_EXTENSION = ".zip"; + private final String directoryName; + private final MaxMindDatabaseConfig maxMindDatabaseConfig; + + public ManifestDownloadService(final String directoryName, final MaxMindDatabaseConfig maxMindDatabaseConfig) { + this.directoryName = directoryName; + this.maxMindDatabaseConfig = maxMindDatabaseConfig; + } + + @Override + public void initiateDownload() { + final Set databasePaths = maxMindDatabaseConfig.getDatabasePaths().keySet(); + for (final String key: databasePaths) { + final Manifest manifest = deserializeManifestFile(maxMindDatabaseConfig.getDatabasePaths().get(key)); + + final String manifestFilePath = manifest.getDbName(); + final String zipFileName = manifestFilePath.substring(0, manifestFilePath.lastIndexOf(".")).concat(ZIP_FILE_EXTENSION); + final String zipFilePath = directoryName + File.separator + zipFileName; + + downloadZipFile(manifest.getUrl(), zipFilePath); + unzipDownloadedFile(zipFilePath, directoryName, key + MAXMIND_DATABASE_EXTENSION); + } + } + + private Manifest deserializeManifestFile(final String CDNEndpoint) { + HttpURLConnection httpURLConnection = null; + try { + final URL url = new URL(CDNEndpoint); + httpURLConnection = (HttpURLConnection) url.openConnection(); + httpURLConnection.addRequestProperty("User-Agent", "Custom-User-Agent"); + + final Manifest manifest = OBJECT_MAPPER.readValue(httpURLConnection.getInputStream(), Manifest.class); + httpURLConnection.disconnect(); + + return manifest; + } catch (final IOException ex) { + if (httpURLConnection != null) { + httpURLConnection.disconnect(); + } + throw new DownloadFailedException("Exception occurred while reading manifest.json file due to: " + ex.getMessage()); + } + } + + private void downloadZipFile(final String databaseUrl, final String destinationPath) { + HttpURLConnection httpURLConnection; + try { + final URL url = new URL(databaseUrl); + httpURLConnection = (HttpURLConnection) url.openConnection(); + // CDN endpoint returns 403 without User Agent. + httpURLConnection.addRequestProperty("User-Agent", "Data Prepper"); + } catch (IOException ex) { + throw new DownloadFailedException("Exception occurred while opening connection due to: " + ex.getMessage()); + } + + try (final BufferedInputStream in = new BufferedInputStream(httpURLConnection.getInputStream()); + final FileOutputStream fileOutputStream = new FileOutputStream(destinationPath)) { + final byte[] dataBuffer = new byte[DEFAULT_BYTE_SIZE]; + int bytesRead; + while ((bytesRead = in.read(dataBuffer, 0, DEFAULT_BYTE_SIZE)) != -1) { + fileOutputStream.write(dataBuffer, 0, bytesRead); + } + httpURLConnection.disconnect(); + } catch (final IOException ex) { + httpURLConnection.disconnect(); + throw new DownloadFailedException("Exception occurred while downloading MaxMind database due to: " + ex.getMessage()); + } + } + + private void unzipDownloadedFile(final String zipFilePath, final String outputFilePath, final String fileName) { + final File inputFile = new File(zipFilePath); + final File outputDir = new File(outputFilePath); + if (!outputDir.exists()) { + outputDir.mkdirs(); + } + + final byte[] buffer = new byte[DEFAULT_BYTE_SIZE]; + + try (final FileInputStream fileInputStream = new FileInputStream(inputFile); + final ZipInputStream zipInputStream = new ZipInputStream(fileInputStream)) { + + ZipEntry zipEntry = zipInputStream.getNextEntry(); + while (zipEntry != null && zipEntry.getName().endsWith(MAXMIND_DATABASE_EXTENSION)) { + final File newFile = new File(outputDir + File.separator + fileName); + + final FileOutputStream fileOutputStream = new FileOutputStream(newFile); + int len; + while ((len = zipInputStream.read(buffer)) > 0) { + fileOutputStream.write(buffer, 0, len); + } + fileOutputStream.close(); + zipInputStream.closeEntry(); + zipEntry = zipInputStream.getNextEntry(); + } + + // deleting zip file after unzipping + inputFile.delete(); + } catch (final IOException e) { + inputFile.delete(); + throw new DownloadFailedException("Exception occurred while unzipping the database file due to: " + e.getMessage()); + } + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/S3DBService.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/S3DBService.java index 11793a8189..c2c10b80bd 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/S3DBService.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/S3DBService.java @@ -5,105 +5,80 @@ package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.DownloadFailedException; +import org.opensearch.dataprepper.plugins.processor.exception.DownloadFailedException; import org.opensearch.dataprepper.plugins.processor.extension.AwsAuthenticationOptionsConfig; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; +import org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig; import software.amazon.awssdk.services.s3.S3AsyncClient; import software.amazon.awssdk.transfer.s3.S3TransferManager; -import software.amazon.awssdk.transfer.s3.model.DirectoryDownload; -import software.amazon.awssdk.transfer.s3.model.DownloadDirectoryRequest; +import software.amazon.awssdk.transfer.s3.model.DownloadFileRequest; +import software.amazon.awssdk.transfer.s3.model.FileDownload; import java.io.File; import java.net.URI; import java.net.URISyntaxException; -import java.nio.file.Paths; -import java.util.List; +import java.util.Set; /** * Implementation class for Download through S3 */ public class S3DBService implements DBSource { - - private static final Logger LOG = LoggerFactory.getLogger(S3DBService.class); - private String bucketName; - private String bucketPath; private final AwsAuthenticationOptionsConfig awsAuthenticationOptionsConfig; - private final String prefixDir; + private final String destinationDirectory; + private final MaxMindDatabaseConfig maxMindDatabaseConfig; /** * S3DBService constructor for initialisation of attributes * - * @param prefixDir prefixDir + * @param destinationDirectory destinationDirectory */ public S3DBService(final AwsAuthenticationOptionsConfig awsAuthenticationOptionsConfig, - final String prefixDir) { + final String destinationDirectory, + final MaxMindDatabaseConfig maxMindDatabaseConfig) { this.awsAuthenticationOptionsConfig = awsAuthenticationOptionsConfig; - this.prefixDir = prefixDir; + this.destinationDirectory = destinationDirectory; + this.maxMindDatabaseConfig = maxMindDatabaseConfig; } /** * Initialisation of Download through Url - * @param s3URLs s3URLs */ - public void initiateDownload(List s3URLs) { - for (String s3Url : s3URLs) { + public void initiateDownload() { + final Set databasePaths = maxMindDatabaseConfig.getDatabasePaths().keySet(); + + for (final String database: databasePaths) { try { - URI uri = new URI(s3Url); - bucketName = uri.getHost(); - bucketPath = removeTrailingSlash(removeLeadingSlash(uri.getPath())); - DBSource.createFolderIfNotExist(tempFolderPath + File.separator + prefixDir); - buildRequestAndDownloadFile(bucketName, bucketPath); + final String s3Uri = maxMindDatabaseConfig.getDatabasePaths().get(database); + final URI uri = new URI(s3Uri); + final String key = uri.getPath().substring(1); + final String bucketName = uri.getHost(); + buildRequestAndDownloadFile(bucketName, key, database); } catch (URISyntaxException ex) { - LOG.info("Initiate Download Exception", ex); + throw new DownloadFailedException("Failed to download database from S3." + ex.getMessage()); } } } - /** - * Removes leading slashes from the input string - * @param str url path - * @return String - */ - public String removeLeadingSlash(String str) { - StringBuilder sb = new StringBuilder(str); - while (sb.length() > 0 && sb.charAt(0) == '/') { - sb.deleteCharAt(0); - } - return sb.toString(); - } - - /** - * Removes trial slashes from the input string - * @param str url path - * @return String - */ - public String removeTrailingSlash(String str) { - StringBuilder sb = new StringBuilder(str); - while (sb.length() > 0 && sb.charAt(sb.length() - 1) == '/') { - sb.setLength(sb.length() - 1); - } - return sb.toString(); - } - /** * Download the mmdb file from the S3 - * @param path path + * @param bucketName Name of the S3 bucket + * @param key Name of S3 object key + * @param fileName Name of the file to save */ - public void buildRequestAndDownloadFile(String... path) { + private void buildRequestAndDownloadFile(final String bucketName, final String key, final String fileName) { try { - S3TransferManager transferManager = createCustomTransferManager(); - DirectoryDownload directoryDownload = - transferManager.downloadDirectory( - DownloadDirectoryRequest.builder() - .destination(Paths.get(tempFolderPath + File.separator + prefixDir)) - .bucket(path[0]) - .listObjectsV2RequestTransformer(l -> l.prefix(path[1])) - .build()); - directoryDownload.completionFuture().join(); + final S3TransferManager transferManager = createCustomTransferManager(); + + DownloadFileRequest downloadFileRequest = DownloadFileRequest.builder() + .getObjectRequest(b -> b.bucket(bucketName).key(key)) + .destination(new File(destinationDirectory + File.separator + fileName + MAXMIND_DATABASE_EXTENSION)) + .build(); + + FileDownload downloadFile = transferManager.downloadFile(downloadFileRequest); + + downloadFile.completionFuture().join(); } catch (Exception ex) { - throw new DownloadFailedException("Download failed: " + ex); + throw new DownloadFailedException("Failed to download database from S3." + ex.getMessage()); } } @@ -112,7 +87,7 @@ public void buildRequestAndDownloadFile(String... path) { * * @return S3TransferManager */ - public S3TransferManager createCustomTransferManager() { + private S3TransferManager createCustomTransferManager() { S3AsyncClient s3AsyncClient = S3AsyncClient.crtBuilder() .region(awsAuthenticationOptionsConfig.getAwsRegion()) diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/loadtype/LoadTypeOptions.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/loadtype/LoadTypeOptions.java deleted file mode 100644 index 38fe132a07..0000000000 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/loadtype/LoadTypeOptions.java +++ /dev/null @@ -1,37 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.loadtype; - -import com.fasterxml.jackson.annotation.JsonCreator; - -import java.util.Arrays; -import java.util.Map; -import java.util.stream.Collectors; - -/** - * LoadTypeOptions enumeration - */ -public enum LoadTypeOptions { - - INMEMORY("memory_map"), - CACHE("cache"); - - private final String option; - private static final Map OPTIONS_MAP = Arrays.stream(LoadTypeOptions.values()) - .collect(Collectors.toMap( - value -> value.option, - value -> value - )); - - LoadTypeOptions(final String option) { - this.option = option; - } - - @JsonCreator - static LoadTypeOptions fromOptionValue(final String option) { - return OPTIONS_MAP.get(option); - } -} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/DatabaseSourceIdentification.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/DatabaseSourceIdentification.java new file mode 100644 index 0000000000..6eb99524b0 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/DatabaseSourceIdentification.java @@ -0,0 +1,120 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.utils; + +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSourceOptions; + +import java.io.File; +import java.net.MalformedURLException; +import java.net.URI; +import java.net.URISyntaxException; +import java.net.URL; +import java.util.List; +import java.util.regex.Pattern; + +/** + * Implementation of class for checking whether URL type is S3 or file path + */ +public class DatabaseSourceIdentification { + + private DatabaseSourceIdentification() { + + } + + private static final String S3_DOMAIN_PATTERN = "[a-zA-Z0-9-]+\\.s3\\.amazonaws\\.com"; + private static final String MANIFEST_ENDPOINT_PATH = "manifest.json"; + + /** + * Check for database path is valid S3 URI or not + * @param uriString uriString + * @return boolean + */ + public static boolean isS3Uri(final String uriString) { + try { + URI uri = new URI(uriString); + if (uri.getScheme() != null && uri.getScheme().equalsIgnoreCase("s3")) { + return true; + } + } catch (Exception e) { + return false; + } + return false; + } + + /** + * Check for database path is valid URL or not + * @param input input + * @return boolean + */ + public static boolean isURL(final String input) { + try { + final URI uri = new URI(input); + final URL url = new URL(input); + return !input.endsWith(MANIFEST_ENDPOINT_PATH) && + !uri.getHost().contains("geoip.maps.opensearch") && + uri.getHost().equals("download.maxmind.com") && + uri.getScheme() != null && + !Pattern.matches(S3_DOMAIN_PATTERN, url.getHost()) && + (uri.getScheme().equals("http") || uri.getScheme().equals("https")); + } catch (URISyntaxException | MalformedURLException e) { + return false; + } + } + + /** + * Check for database path is local file path or not + * @param input input + * @return boolean + */ + public static boolean isFilePath(final String input) { + final File file = new File(input); + return file.exists() && file.isFile(); + } + + /** + * Check for database path is CDN endpoint + * @param input input + * @return boolean + */ + public static boolean isCDNEndpoint(final String input) { + if (input.endsWith(MANIFEST_ENDPOINT_PATH)) { + try { + final URI uri = new URI(input); + return uri.getScheme().equals("http") || uri.getScheme().equals("https"); + } catch (final URISyntaxException e) { + return false; + } + } + return false; + } + + /** + * Get the database path options based on input URL + * @param databasePaths - List of database paths to get databases data from + * @return DBSourceOptions + */ + public static DBSourceOptions getDatabasePathType(final List databasePaths) { + DBSourceOptions downloadSourceOptions = null; + for(final String databasePath : databasePaths) { + + if(DatabaseSourceIdentification.isFilePath(databasePath)) { + return DBSourceOptions.PATH; + } + else if (DatabaseSourceIdentification.isCDNEndpoint(databasePath)) { + return DBSourceOptions.HTTP_MANIFEST; + } + else if(DatabaseSourceIdentification.isURL(databasePath)) + { + return DBSourceOptions.URL; + } + else if(DatabaseSourceIdentification.isS3Uri(databasePath)) + { + return DBSourceOptions.S3; + } + } + return downloadSourceOptions; + } +} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/DbSourceIdentification.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/DbSourceIdentification.java deleted file mode 100644 index e7eaffe75d..0000000000 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/DbSourceIdentification.java +++ /dev/null @@ -1,109 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.utils; - -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSourceOptions; - -import java.net.MalformedURLException; -import java.net.URI; -import java.net.URISyntaxException; -import java.net.URL; -import java.util.List; -import java.util.regex.Pattern; - -/** - * Implementation of class for checking whether URL type is S3 or file path - */ -public class DbSourceIdentification { - - private DbSourceIdentification() { - - } - - private static String s3DomainPattern = "[a-zA-Z0-9-]+\\.s3\\.amazonaws\\.com"; - - /** - * Check for database path is valid S3 URI or not - * @param uriString uriString - * @return boolean - */ - public static boolean isS3Uri(String uriString) { - try { - URI uri = new URI(uriString); - if (uri.getScheme() != null && uri.getScheme().equalsIgnoreCase("s3")) { - return true; - } - } catch (Exception e) { - return false; - } - return false; - } - - /** - * Check for database path is valid S3 URL or not - * @param urlString urlString - * @return boolean - */ - public static boolean isS3Url(String urlString) { - try { - URL url = new URL(urlString); - if (Pattern.matches(s3DomainPattern, url.getHost())) { - return true; - } - } catch (Exception e) { - return false; - } - return false; - } - - /** - * Check for database path is valid URL or not - * @param input input - * @return boolean - */ - public static boolean isURL(String input) { - try { - URI uri = new URI(input); - URL url = new URL(input); - return uri.getScheme() != null && !Pattern.matches(s3DomainPattern, url.getHost()) &&(uri.getScheme().equals("http") || uri.getScheme().equals("https")); - } catch (URISyntaxException | MalformedURLException e) { - return false; - } - } - - /** - * Check for database path is local file path or not - * @param input input - * @return boolean - */ - public static boolean isFilePath(String input) { - return input.startsWith("/") || input.startsWith("./") || input.startsWith("\\") || (input.length() > 1 && input.charAt(1) == ':'); - } - - /** - * Get the database path options based on input URL - * @param databasePaths - List of database paths to get databases data from - * @return DBSourceOptions - */ - public static DBSourceOptions getDatabasePathType(List databasePaths) { - DBSourceOptions downloadSourceOptions = null; - for(final String databasePath : databasePaths) { - - if(DbSourceIdentification.isFilePath(databasePath)) { - return DBSourceOptions.PATH; - } - else if(DbSourceIdentification.isURL(databasePath)) - { - downloadSourceOptions = DBSourceOptions.URL; - } - else if(DbSourceIdentification.isS3Uri(databasePath) || (DbSourceIdentification.isS3Url(databasePath))) - { - downloadSourceOptions = DBSourceOptions.S3; - } - } - return downloadSourceOptions; - } -} diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/IPValidationCheck.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/IPValidationCheck.java index d717ebc711..71e3602b5c 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/IPValidationCheck.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/IPValidationCheck.java @@ -5,6 +5,8 @@ package org.opensearch.dataprepper.plugins.processor.utils; +import org.opensearch.dataprepper.plugins.processor.exception.InvalidIPAddressException; + import java.net.Inet4Address; import java.net.Inet6Address; import java.net.InetAddress; @@ -20,10 +22,15 @@ public class IPValidationCheck { * Check for IP is valid or not * @param ipAddress ipAddress * @return boolean - * @throws UnknownHostException UnknownHostException + * @throws InvalidIPAddressException InvalidIPAddressException */ - public static boolean isPublicIpAddress(final String ipAddress) throws UnknownHostException { - InetAddress address = InetAddress.getByName(ipAddress); + public static boolean isPublicIpAddress(final String ipAddress) { + InetAddress address; + try { + address = InetAddress.getByName(ipAddress); + } catch (final UnknownHostException e) { + return false; + } if (address instanceof Inet6Address || address instanceof Inet4Address) { return !address.isSiteLocalAddress() && !address.isLoopbackAddress(); } diff --git a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/LicenseTypeCheck.java b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/LicenseTypeCheck.java index 1e75d4c4ca..907225a4d8 100644 --- a/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/LicenseTypeCheck.java +++ b/data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/utils/LicenseTypeCheck.java @@ -10,23 +10,19 @@ import java.io.File; import java.nio.file.Path; import java.nio.file.Paths; -import java.util.Arrays; -import java.util.List; -/**cls +/** * * Implementation of class logic to check maxmind database id free or enterprise version */ public class LicenseTypeCheck { - - protected static final String[] DEFAULT_DATABASE_FILENAMES = new String[] { "GeoLite2-ASN.mmdb", "GeoLite2-City.mmdb", "GeoLite2-Country.mmdb" }; - protected static final List geoLite2Database = Arrays.asList(DEFAULT_DATABASE_FILENAMES); - private static final String geoIP2EnterpriseDB = "GeoIP2-Enterprise.mmdb"; + private static final String GEOIP2_DATABASE = "geoip2"; + private static final String GEOLITE2_DATABASE = "geolite2"; + private static final String MMDB = "mmdb"; public LicenseTypeCheck() { - } /** @@ -36,31 +32,30 @@ public LicenseTypeCheck() { * @return license type options */ public LicenseTypeOptions isGeoLite2OrEnterpriseLicense(final String databasePath) { - LicenseTypeOptions licenseTypeOptions = LicenseTypeOptions.ENTERPRISE; - File directory = new File(databasePath); - // list all files present in the directory - File[] files = directory.listFiles(); - - for(File file : files) { - // convert the file name into string - String fileName = file.toString(); - - int index = fileName.lastIndexOf('.'); - if(index > 0) { - String extension = fileName.substring(index + 1); - Path onlyFileName = Paths.get(fileName).getFileName(); - - if((extension.equals("mmdb")) && (geoIP2EnterpriseDB.equals(onlyFileName.toString()))) { - licenseTypeOptions = LicenseTypeOptions.ENTERPRISE; - break; - } - else if((extension.equals("mmdb")) && (geoLite2Database.contains(onlyFileName.toString()))) - { - licenseTypeOptions = LicenseTypeOptions.FREE; + LicenseTypeOptions licenseTypeOptions = null; + final File directory = new File(databasePath); + if (directory.isDirectory()) { + // list all files present in the directory + final File[] files = directory.listFiles(); + + for (final File file : files) { + // convert the file name into string + final String fileName = file.toString(); + + int index = fileName.lastIndexOf('.'); + if (index > 0) { + String extension = fileName.substring(index + 1); + Path onlyFileName = Paths.get(fileName).getFileName(); + + if ((extension.equals(MMDB)) && (onlyFileName.toString().toLowerCase().contains(GEOIP2_DATABASE))) { + licenseTypeOptions = LicenseTypeOptions.ENTERPRISE; + break; + } else if ((extension.equals(MMDB)) && (onlyFileName.toString().toLowerCase().contains(GEOLITE2_DATABASE))) { + licenseTypeOptions = LicenseTypeOptions.FREE; + } } } } return licenseTypeOptions; } } - diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPFieldTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPFieldTest.java new file mode 100644 index 0000000000..4a4896304c --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPFieldTest.java @@ -0,0 +1,26 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor; + +import org.junit.jupiter.api.Test; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; + +class GeoIPFieldTest { + + @Test + void test_findByName_should_return_geoip_field_if_valid() { + final GeoIPField geoIPField = GeoIPField.findByName("city_name"); + assertThat(geoIPField, equalTo(GeoIPField.CITY_NAME)); + } + + @Test + void test_findByName_should_return_null_if_invalid() { + final GeoIPField geoIPField = GeoIPField.findByName("coordinates"); + assertThat(geoIPField, equalTo(null)); + } +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorConfigTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorConfigTest.java index 1709e39234..f8776faa41 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorConfigTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorConfigTest.java @@ -30,22 +30,26 @@ void setUp() { @Test void testDefaultConfig() { assertThat(geoIPProcessorConfig.getEntries(), equalTo(null)); - assertThat(geoIPProcessorConfig.getTagsOnFailure(), equalTo(null)); + assertThat(geoIPProcessorConfig.getTagsOnEngineFailure(), equalTo(null)); + assertThat(geoIPProcessorConfig.getTagsOnIPNotFound(), equalTo(null)); assertThat(geoIPProcessorConfig.getWhenCondition(), equalTo(null)); } @Test void testGetEntries() throws NoSuchFieldException, IllegalAccessException { final List entries = List.of(new EntryConfig()); - final List tagsOnFailure = List.of("tag1", "tag2"); + final List tagsOnEngineFailure = List.of("tag1", "tag2"); + final List tagsOnIPNotFound = List.of("tag3"); final String whenCondition = "/ip == 1.2.3.4"; ReflectivelySetField.setField(GeoIPProcessorConfig.class, geoIPProcessorConfig, "entries", entries); - ReflectivelySetField.setField(GeoIPProcessorConfig.class, geoIPProcessorConfig, "tagsOnFailure", tagsOnFailure); + ReflectivelySetField.setField(GeoIPProcessorConfig.class, geoIPProcessorConfig, "tagsOnEngineFailure", tagsOnEngineFailure); + ReflectivelySetField.setField(GeoIPProcessorConfig.class, geoIPProcessorConfig, "tagsOnIPNotFound", tagsOnIPNotFound); ReflectivelySetField.setField(GeoIPProcessorConfig.class, geoIPProcessorConfig, "whenCondition", whenCondition); assertThat(geoIPProcessorConfig.getEntries(), equalTo(entries)); - assertThat(geoIPProcessorConfig.getTagsOnFailure(), equalTo(tagsOnFailure)); + assertThat(geoIPProcessorConfig.getTagsOnEngineFailure(), equalTo(tagsOnEngineFailure)); + assertThat(geoIPProcessorConfig.getTagsOnIPNotFound(), equalTo(tagsOnIPNotFound)); assertThat(geoIPProcessorConfig.getWhenCondition(), equalTo(whenCondition)); } } diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorTest.java index d6b8a7bf01..6f25dc6f81 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/GeoIPProcessorTest.java @@ -10,7 +10,10 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Captor; import org.mockito.Mock; +import org.mockito.MockedStatic; import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.dataprepper.expression.ExpressionEvaluator; import org.opensearch.dataprepper.metrics.PluginMetrics; @@ -19,31 +22,64 @@ import org.opensearch.dataprepper.model.log.JacksonLog; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.processor.configuration.EntryConfig; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.EnrichFailedException; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader; +import org.opensearch.dataprepper.plugins.processor.exception.EngineFailureException; +import org.opensearch.dataprepper.plugins.processor.exception.EnrichFailedException; import org.opensearch.dataprepper.plugins.processor.extension.GeoIPProcessorService; import org.opensearch.dataprepper.plugins.processor.extension.GeoIpConfigSupplier; -import org.opensearch.dataprepper.test.helper.ReflectivelySetField; +import org.opensearch.dataprepper.plugins.processor.utils.IPValidationCheck; import java.util.ArrayList; import java.util.Collection; +import java.util.Collections; import java.util.HashMap; import java.util.LinkedList; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.stream.Collectors; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsInAnyOrder; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mockStatic; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; -import static org.opensearch.dataprepper.plugins.processor.GeoIPProcessor.GEO_IP_PROCESSING_MATCH; -import static org.opensearch.dataprepper.plugins.processor.GeoIPProcessor.GEO_IP_PROCESSING_MISMATCH; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CITY_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CONTINENT_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CONTINENT_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.COUNTRY_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.IS_COUNTRY_IN_EUROPEAN_UNION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LATITUDE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LEAST_SPECIFIED_SUBDIVISION_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LEAST_SPECIFIED_SUBDIVISION_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LOCATION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LOCATION_ACCURACY_RADIUS; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LONGITUDE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.METRO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.MOST_SPECIFIED_SUBDIVISION_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.MOST_SPECIFIED_SUBDIVISION_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.MOST_SPECIFIED_SUBDIVISION_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.POSTAL_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.POSTAL_CODE_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REGISTERED_COUNTRY_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REGISTERED_COUNTRY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_TYPE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.TIME_ZONE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPProcessor.GEO_IP_EVENTS_FAILED; +import static org.opensearch.dataprepper.plugins.processor.GeoIPProcessor.GEO_IP_EVENTS_FAILED_ENGINE_EXCEPTION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPProcessor.GEO_IP_EVENTS_FAILED_IP_NOT_FOUND; +import static org.opensearch.dataprepper.plugins.processor.GeoIPProcessor.GEO_IP_EVENTS_PROCESSED; +import static org.opensearch.dataprepper.plugins.processor.GeoIPProcessor.GEO_IP_EVENTS_SUCCEEDED; @ExtendWith(MockitoExtension.class) class GeoIPProcessorTest { @@ -62,21 +98,38 @@ class GeoIPProcessorTest { @Mock private PluginMetrics pluginMetrics; @Mock - private Counter geoIpProcessingMatch; + private Counter geoIpEventsProcessed; @Mock - private Counter geoIpProcessingMismatch; + private Counter geoIpEventsFailed; + @Mock + private Counter geoIpEventsSucceeded; + @Mock + private Counter geoIpEventsFailedEngineException; + @Mock + private Counter geoIpEventsFailedIPNotFound; + @Mock + private GeoIPDatabaseReader geoIPDatabaseReader; + @Captor + private ArgumentCaptor> geoIPFieldCaptor; @BeforeEach void setUp() { - when(geoIpConfigSupplier.getGeoIPProcessorService()).thenReturn(geoIPProcessorService); - lenient().when(pluginMetrics.counter(GEO_IP_PROCESSING_MATCH)).thenReturn(geoIpProcessingMatch); - lenient().when(pluginMetrics.counter(GEO_IP_PROCESSING_MISMATCH)).thenReturn(geoIpProcessingMismatch); + when(geoIpConfigSupplier.getGeoIPProcessorService()).thenReturn(Optional.of(geoIPProcessorService)); + lenient().when(geoIPProcessorService.getGeoIPDatabaseReader()).thenReturn(geoIPDatabaseReader); + lenient().when(pluginMetrics.counter(GEO_IP_EVENTS_PROCESSED)).thenReturn(geoIpEventsProcessed); + lenient().when(pluginMetrics.counter(GEO_IP_EVENTS_SUCCEEDED)).thenReturn(geoIpEventsSucceeded); + lenient().when(pluginMetrics.counter(GEO_IP_EVENTS_FAILED)).thenReturn(geoIpEventsFailed); + lenient().when(pluginMetrics.counter(GEO_IP_EVENTS_FAILED_ENGINE_EXCEPTION)).thenReturn(geoIpEventsFailedEngineException); + lenient().when(pluginMetrics.counter(GEO_IP_EVENTS_FAILED_IP_NOT_FOUND)).thenReturn(geoIpEventsFailedIPNotFound); } @AfterEach void tearDown() { - verifyNoMoreInteractions(geoIpProcessingMatch); - verifyNoMoreInteractions(geoIpProcessingMismatch); + verifyNoMoreInteractions(geoIpEventsProcessed); + verifyNoMoreInteractions(geoIpEventsSucceeded); + verifyNoMoreInteractions(geoIpEventsFailed); + verifyNoMoreInteractions(geoIpEventsFailedEngineException); + verifyNoMoreInteractions(geoIpEventsFailedIPNotFound); } private GeoIPProcessor createObjectUnderTest() { @@ -84,88 +137,280 @@ private GeoIPProcessor createObjectUnderTest() { } @Test - void doExecuteTest_with_when_condition_should_only_enrich_events_that_match_when_condition() throws NoSuchFieldException, IllegalAccessException { + void doExecuteTest_with_when_condition_should_enrich_events_that_match_when_condition() { final String whenCondition = "/peer/status == success"; when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); when(geoIPProcessorConfig.getWhenCondition()).thenReturn(whenCondition); when(entry.getSource()).thenReturn("/peer/ip"); when(entry.getTarget()).thenReturn(TARGET); - when(entry.getFields()).thenReturn(setFields()); + when(entry.getIncludeFields()).thenReturn(setFields()); final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); - when(geoIPProcessorService.getGeoData(any(), any())).thenReturn(prepareGeoData()); + when(geoIPDatabaseReader.getGeoData(any(), any(), any())).thenReturn(prepareGeoData()); + + final Record record1 = createCustomRecord("success"); + List> recordsIn = List.of(record1); + + when(expressionEvaluator.evaluateConditional(whenCondition, record1.getData())).thenReturn(true); + + final Collection> records = geoIPProcessor.doExecute(recordsIn); + + assertThat(records.size(), equalTo(1)); + + final Collection> recordsWithLocation = records.stream().filter(record -> record.getData().containsKey(TARGET)) + .collect(Collectors.toList()); - ReflectivelySetField.setField(GeoIPProcessor.class, geoIPProcessor, "geoIPProcessorService", geoIPProcessorService); + assertThat(recordsWithLocation.size(), equalTo(1)); + verify(geoIpEventsSucceeded).increment(); + verify(geoIpEventsProcessed).increment(); + } + + @Test + void doExecuteTest_with_when_condition_should_not_enrich_if_when_condition_is_false() { + final String whenCondition = "/peer/status == success"; + + when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); + when(geoIPProcessorConfig.getWhenCondition()).thenReturn(whenCondition); + + final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); final Record record1 = createCustomRecord("success"); - final Record record2 = createCustomRecord("failed"); List> recordsIn = new ArrayList<>(); recordsIn.add(record1); - recordsIn.add(record2); - when(expressionEvaluator.evaluateConditional(whenCondition, record1.getData())).thenReturn(true); - when(expressionEvaluator.evaluateConditional(whenCondition, record2.getData())).thenReturn(false); + when(expressionEvaluator.evaluateConditional(whenCondition, record1.getData())).thenReturn(false); final Collection> records = geoIPProcessor.doExecute(recordsIn); - assertThat(records.size(), equalTo(2)); + assertThat(records.size(), equalTo(1)); final Collection> recordsWithLocation = records.stream().filter(record -> record.getData().containsKey(TARGET)) .collect(Collectors.toList()); - assertThat(recordsWithLocation.size(), equalTo(1)); + assertThat(recordsWithLocation.size(), equalTo(0)); + } - for (final Record record : recordsWithLocation) { + @Test + void doExecuteTest_should_add_geo_data_to_event_if_source_is_non_null() { + when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); + when(entry.getSource()).thenReturn(SOURCE); + when(entry.getTarget()).thenReturn(TARGET); + when(entry.getIncludeFields()).thenReturn(setFields()); + + final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); + + when(geoIPDatabaseReader.getGeoData(any(), any(), any())).thenReturn(prepareGeoData()); + Collection> records = geoIPProcessor.doExecute(setEventQueue()); + for (final Record record : records) { final Event event = record.getData(); - assertThat(event.get("/peer/status", String.class), equalTo("success")); + assertThat(event.get("/peer/ip", String.class), equalTo("136.226.242.205")); + assertThat(event.containsKey(TARGET), equalTo(true)); + verify(geoIpEventsProcessed).increment(); + verify(geoIpEventsSucceeded).increment(); } - verify(geoIpProcessingMatch).increment(); } @Test - void doExecuteTest() throws NoSuchFieldException, IllegalAccessException { + void doExecuteTest_should_add_geo_data_with_expected_fields_to_event_when_include_fields_is_configured() { when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); when(entry.getSource()).thenReturn(SOURCE); when(entry.getTarget()).thenReturn(TARGET); - when(entry.getFields()).thenReturn(setFields()); + + final List includeFields = List.of("city_name", "asn"); + final List includeFieldsResult = List.of(GeoIPField.CITY_NAME, GeoIPField.ASN); + when(entry.getIncludeFields()).thenReturn(includeFields); final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); - when(geoIPProcessorService.getGeoData(any(), any())).thenReturn(prepareGeoData()); - ReflectivelySetField.setField(GeoIPProcessor.class, geoIPProcessor, - "geoIPProcessorService", geoIPProcessorService); + when(geoIPDatabaseReader.getGeoData(any(), any(), any())).thenReturn(prepareGeoData()); Collection> records = geoIPProcessor.doExecute(setEventQueue()); + verify(geoIPDatabaseReader).getGeoData(any(), geoIPFieldCaptor.capture(), any()); + for (final Record record : records) { final Event event = record.getData(); assertThat(event.get("/peer/ip", String.class), equalTo("136.226.242.205")); - assertThat(event.containsKey("geolocation"), equalTo(true)); - verify(geoIpProcessingMatch).increment(); + assertThat(event.containsKey(TARGET), equalTo(true)); + verify(geoIpEventsProcessed).increment(); + verify(geoIpEventsSucceeded).increment(); } + + final List value = geoIPFieldCaptor.getValue(); + assertThat(value, containsInAnyOrder(includeFieldsResult.toArray())); } @Test - void test_tags_when_enrich_fails() { + void doExecuteTest_should_add_geo_data_with_expected_fields_to_event_when_exclude_fields_is_configured() { + when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); when(entry.getSource()).thenReturn(SOURCE); - when(entry.getFields()).thenReturn(setFields()); + when(entry.getTarget()).thenReturn(TARGET); + + final List excludeFields = List.of("country_name", "country_iso_code", "city_name", "asn", "asn_organization", "network", "ip"); + final List excludeFieldsResult = List.of(CONTINENT_NAME, CONTINENT_CODE, IS_COUNTRY_IN_EUROPEAN_UNION, + REPRESENTED_COUNTRY_NAME, REPRESENTED_COUNTRY_ISO_CODE, REPRESENTED_COUNTRY_TYPE, REGISTERED_COUNTRY_NAME, + REGISTERED_COUNTRY_ISO_CODE, LOCATION, LOCATION_ACCURACY_RADIUS, LATITUDE, LONGITUDE, METRO_CODE, TIME_ZONE, POSTAL_CODE, + MOST_SPECIFIED_SUBDIVISION_NAME, MOST_SPECIFIED_SUBDIVISION_ISO_CODE, LEAST_SPECIFIED_SUBDIVISION_NAME, + LEAST_SPECIFIED_SUBDIVISION_ISO_CODE, COUNTRY_CONFIDENCE, CITY_CONFIDENCE, MOST_SPECIFIED_SUBDIVISION_CONFIDENCE, + LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE, POSTAL_CODE_CONFIDENCE); + when(entry.getExcludeFields()).thenReturn(excludeFields); + + final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); + + when(geoIPDatabaseReader.getGeoData(any(), any(), any())).thenReturn(prepareGeoData()); + Collection> records = geoIPProcessor.doExecute(setEventQueue()); + verify(geoIPDatabaseReader).getGeoData(any(), geoIPFieldCaptor.capture(), any()); + + for (final Record record : records) { + final Event event = record.getData(); + assertThat(event.get("/peer/ip", String.class), equalTo("136.226.242.205")); + assertThat(event.containsKey(TARGET), equalTo(true)); + verify(geoIpEventsProcessed).increment(); + verify(geoIpEventsSucceeded).increment(); + } + + final List value = geoIPFieldCaptor.getValue(); + assertThat(value, containsInAnyOrder(excludeFieldsResult.toArray())); + } + @Test + void doExecuteTest_should_not_add_geo_data_to_event_if_source_is_null() { + when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); + when(entry.getSource()).thenReturn("ip"); + when(entry.getIncludeFields()).thenReturn(setFields()); List testTags = List.of("tag1", "tag2"); - when(geoIPProcessorConfig.getTagsOnFailure()).thenReturn(testTags); + when(geoIPProcessorConfig.getTagsOnIPNotFound()).thenReturn(testTags); + + final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); + + Collection> records = geoIPProcessor.doExecute(setEventQueue()); + + for (final Record record : records) { + final Event event = record.getData(); + assertThat(!event.containsKey("geo"), equalTo(true)); + assertTrue(event.getMetadata().hasTags(testTags)); + verify(geoIpEventsProcessed).increment(); + verify(geoIpEventsFailed).increment(); + verify(geoIpEventsFailedIPNotFound).increment(); + } + } + + @Test + void doExecuteTest_should_not_add_geo_data_to_event_if_returned_data_is_empty() { + when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); + when(entry.getSource()).thenReturn(SOURCE); + when(entry.getIncludeFields()).thenReturn(setFields()); + + final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); + + when(geoIPDatabaseReader.getGeoData(any(), any(), any())).thenReturn(Collections.EMPTY_MAP); + Collection> records = geoIPProcessor.doExecute(setEventQueue()); + for (final Record record : records) { + final Event event = record.getData(); + assertThat(!event.containsKey("geo"), equalTo(true)); + verify(geoIpEventsProcessed).increment(); + verify(geoIpEventsFailed).increment(); + verify(geoIpEventsFailedIPNotFound).increment(); + } + } + + @Test + void doExecuteTest_should_not_add_geodata_if_database_is_expired() { + when(geoIPDatabaseReader.isExpired()).thenReturn(true); + + final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); + + final Collection> records = geoIPProcessor.doExecute(setEventQueue()); + for (final Record record : records) { + final Event event = record.getData(); + assertThat(!event.containsKey("geo"), equalTo(true)); + verify(geoIpEventsProcessed).increment(); + verify(geoIpEventsFailed).increment(); + } + } + + @Test + void doExecuteTest_should_not_add_geodata_if_database_reader_is_null() { + when(geoIPProcessorService.getGeoIPDatabaseReader()).thenReturn(null); + + final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); + + final Collection> records = geoIPProcessor.doExecute(setEventQueue()); + for (final Record record : records) { + final Event event = record.getData(); + assertThat(!event.containsKey("geo"), equalTo(true)); + verify(geoIpEventsProcessed).increment(); + verify(geoIpEventsFailed).increment(); + } + } + + @Test + void doExecuteTest_should_not_add_geodata_if_ip_address_is_not_public() { + try (final MockedStatic ipValidationCheckMockedStatic = mockStatic(IPValidationCheck.class)) { + ipValidationCheckMockedStatic.when(() -> IPValidationCheck.isPublicIpAddress(any())).thenReturn(false); + + when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); + when(entry.getSource()).thenReturn(SOURCE); + when(entry.getIncludeFields()).thenReturn(setFields()); + + final GeoIPProcessor geoIPProcessor = createObjectUnderTest(); + + final Collection> records = geoIPProcessor.doExecute(setEventQueue()); + for (final Record record : records) { + final Event event = record.getData(); + assertThat(!event.containsKey("geo"), equalTo(true)); + verify(geoIpEventsProcessed).increment(); + verify(geoIpEventsFailed).increment(); + verify(geoIpEventsFailedIPNotFound).increment(); + } + } + } + + @Test + void test_ip_not_found_tags_when_EnrichFailedException_is_thrown() { + when(entry.getSource()).thenReturn(SOURCE); + when(entry.getIncludeFields()).thenReturn(setFields()); + + List testTags = List.of("tag1", "tag2"); + when(geoIPProcessorConfig.getTagsOnIPNotFound()).thenReturn(testTags); when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); - when(geoIpConfigSupplier.getGeoIPProcessorService()).thenReturn(geoIPProcessorService); + GeoIPProcessor geoIPProcessor = createObjectUnderTest(); + + doThrow(EnrichFailedException.class).when(geoIPDatabaseReader).getGeoData(any(), any(), any()); + + Collection> records = geoIPProcessor.doExecute(setEventQueue()); + + for (final Record record : records) { + Event event = record.getData(); + assertTrue(event.getMetadata().hasTags(testTags)); + verify(geoIpEventsFailed).increment(); + verify(geoIpEventsProcessed).increment(); + verify(geoIpEventsFailedIPNotFound).increment(); + } + } + + @Test + void test_ip_not_found_tags_when_EngineFailureException_is_thrown() { + when(entry.getSource()).thenReturn(SOURCE); + when(entry.getIncludeFields()).thenReturn(setFields()); + + List testTags = List.of("tag1", "tag2"); + when(geoIPProcessorConfig.getTagsOnEngineFailure()).thenReturn(testTags); + when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry)); GeoIPProcessor geoIPProcessor = createObjectUnderTest(); - doThrow(EnrichFailedException.class).when(geoIPProcessorService).getGeoData(any(), any()); + doThrow(EngineFailureException.class).when(geoIPDatabaseReader).getGeoData(any(), any(), any()); Collection> records = geoIPProcessor.doExecute(setEventQueue()); for (final Record record : records) { Event event = record.getData(); assertTrue(event.getMetadata().hasTags(testTags)); - verify(geoIpProcessingMismatch).increment(); + verify(geoIpEventsProcessed).increment(); + verify(geoIpEventsFailed).increment(); + verify(geoIpEventsFailedEngineException).increment(); } } diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/configuration/EntryConfigTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/configuration/EntryConfigTest.java index 3863663c7a..8b095ffb63 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/configuration/EntryConfigTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/configuration/EntryConfigTest.java @@ -5,7 +5,6 @@ package org.opensearch.dataprepper.plugins.processor.configuration; -import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.test.helper.ReflectivelySetField; @@ -19,31 +18,77 @@ import static org.opensearch.dataprepper.plugins.processor.configuration.EntryConfig.DEFAULT_TARGET; class EntryConfigTest { - public static final String SOURCE_VALUE = "source"; - public static final String TARGET_VALUE = "target"; - public static final List FIELDS_VALUE = List.of("city", "country"); - private EntryConfig entryConfig; - - @BeforeEach - void setUp() { - entryConfig = new EntryConfig(); + + private EntryConfig createObjectUnderTest() { + return new EntryConfig(); } @Test void testDefaultConfig() { + final EntryConfig entryConfig = createObjectUnderTest(); + assertThat(entryConfig.getSource(), is(nullValue())); assertThat(entryConfig.getTarget(), equalTo(DEFAULT_TARGET)); - assertThat(entryConfig.getFields(), is(nullValue())); + assertThat(entryConfig.getIncludeFields(), equalTo(null)); + assertThat(entryConfig.getExcludeFields(), equalTo(null)); } @Test void testCustomConfig() throws NoSuchFieldException, IllegalAccessException { - ReflectivelySetField.setField(EntryConfig.class, entryConfig, "source", SOURCE_VALUE); - ReflectivelySetField.setField(EntryConfig.class, entryConfig, "target", TARGET_VALUE); - ReflectivelySetField.setField(EntryConfig.class, entryConfig, "fields", FIELDS_VALUE); + final EntryConfig entryConfig = createObjectUnderTest(); + + final String sourceValue = "source"; + final String targetValue = "target"; + final List includeFieldsValue = List.of("city_name"); + final List excludeFieldsValue = List.of("asn"); + + ReflectivelySetField.setField(EntryConfig.class, entryConfig, "source", sourceValue); + ReflectivelySetField.setField(EntryConfig.class, entryConfig, "target", targetValue); + ReflectivelySetField.setField(EntryConfig.class, entryConfig, "includeFields", includeFieldsValue); + ReflectivelySetField.setField(EntryConfig.class, entryConfig, "excludeFields", excludeFieldsValue); + + assertThat(entryConfig.getSource(), equalTo(sourceValue)); + assertThat(entryConfig.getTarget(), equalTo(targetValue)); + assertThat(entryConfig.getIncludeFields(), equalTo(includeFieldsValue)); + assertThat(entryConfig.getExcludeFields(), equalTo(excludeFieldsValue)); + } + + @Test + void test_areFieldsValid_should_return_true_if_only_include_fields_is_configured() throws NoSuchFieldException, IllegalAccessException { + final EntryConfig entryConfig = createObjectUnderTest(); + final List includeFields = List.of("city_name", "continent_code"); + + ReflectivelySetField.setField(EntryConfig.class, entryConfig, "includeFields", includeFields); + + assertThat(entryConfig.areFieldsValid(), equalTo(true)); + } + + @Test + void test_areFieldsValid_should_return_true_if_only_exclude_fields_is_configured() throws NoSuchFieldException, IllegalAccessException { + final EntryConfig entryConfig = createObjectUnderTest(); + final List excludeFields = List.of("city_name", "continent_code"); + + ReflectivelySetField.setField(EntryConfig.class, entryConfig, "excludeFields", excludeFields); + + assertThat(entryConfig.areFieldsValid(), equalTo(true)); + } + + @Test + void test_areFieldsValid_should_return_false_if_both_include_and_exclude_fields_are_configured() throws NoSuchFieldException, IllegalAccessException { + final EntryConfig entryConfig = createObjectUnderTest(); + final List includeFields = List.of("city_name", "continent_code"); + final List excludeFields = List.of("asn"); + + ReflectivelySetField.setField(EntryConfig.class, entryConfig, "includeFields", includeFields); + ReflectivelySetField.setField(EntryConfig.class, entryConfig, "excludeFields", excludeFields); + + assertThat(entryConfig.areFieldsValid(), equalTo(false)); + } + + @Test + void test_areFieldsValid_should_return_false_if_both_include_and_exclude_fields_are_not_configured() { + final EntryConfig entryConfig = createObjectUnderTest(); - assertThat(entryConfig.getSource(), equalTo(SOURCE_VALUE)); - assertThat(entryConfig.getTarget(), equalTo(TARGET_VALUE)); - assertThat(entryConfig.getFields(), equalTo(FIELDS_VALUE)); + assertThat(entryConfig.areFieldsValid(), equalTo(false)); } } diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/AutoCountingDatabaseReaderTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/AutoCountingDatabaseReaderTest.java new file mode 100644 index 0000000000..61cf3a4d15 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/AutoCountingDatabaseReaderTest.java @@ -0,0 +1,78 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.databaseenrich; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.plugins.processor.GeoIPDatabase; +import org.opensearch.dataprepper.plugins.processor.GeoIPField; + +import java.net.InetAddress; +import java.util.List; +import java.util.Set; + +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.ASN; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.ASN_ORGANIZATION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.IP; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.NETWORK; + +@ExtendWith(MockitoExtension.class) +class AutoCountingDatabaseReaderTest { + @Mock + private GeoLite2DatabaseReader geoLite2DatabaseReader; + @Mock + private InetAddress inetAddress; + + GeoIPDatabaseReader createObjectUnderTest() { + return new AutoCountingDatabaseReader(geoLite2DatabaseReader); + } + + @Test + void test_database_close_should_not_close_the_reader_if_close_count_is_not_zero() throws Exception { + final GeoIPDatabaseReader objectUnderTest = createObjectUnderTest(); + + objectUnderTest.retain(); + objectUnderTest.close(); + + verifyNoMoreInteractions(geoLite2DatabaseReader); + } + + @Test + void test_database_close_should_close_the_reader_when_close_count_is_zero() throws Exception { + final GeoIPDatabaseReader objectUnderTest = createObjectUnderTest(); + + objectUnderTest.retain(); + objectUnderTest.close(); + objectUnderTest.close(); + + verify(geoLite2DatabaseReader).close(); + } + + @Test + void test_getGeoData_should_call_delegate_reader_getGeoData() { + final GeoIPDatabaseReader objectUnderTest = createObjectUnderTest(); + + final Set databases = Set.of(GeoIPDatabase.ASN); + final List fields = List.of(ASN, ASN_ORGANIZATION, NETWORK, IP); + objectUnderTest.getGeoData(inetAddress, fields, databases); + + verify(geoLite2DatabaseReader).getGeoData(inetAddress, fields, databases); + } + + @Test + void test_isExpired_should_call_delegate_reader_isExpired() { + final GeoIPDatabaseReader objectUnderTest = createObjectUnderTest(); + + objectUnderTest.isExpired(); + + verify(geoLite2DatabaseReader).isExpired(); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIP2DatabaseReaderTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIP2DatabaseReaderTest.java new file mode 100644 index 0000000000..e20820a91a --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoIP2DatabaseReaderTest.java @@ -0,0 +1,341 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.databaseenrich; + +import com.maxmind.db.Metadata; +import com.maxmind.db.Network; +import com.maxmind.geoip2.DatabaseReader; +import com.maxmind.geoip2.exception.GeoIp2Exception; +import com.maxmind.geoip2.model.AsnResponse; +import com.maxmind.geoip2.model.EnterpriseResponse; +import com.maxmind.geoip2.record.City; +import com.maxmind.geoip2.record.Continent; +import com.maxmind.geoip2.record.Country; +import com.maxmind.geoip2.record.Location; +import com.maxmind.geoip2.record.Postal; +import com.maxmind.geoip2.record.RepresentedCountry; +import com.maxmind.geoip2.record.Subdivision; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.plugins.processor.GeoIPDatabase; +import org.opensearch.dataprepper.plugins.processor.GeoIPField; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DatabaseReaderBuilder; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPFileManager; + +import java.io.File; +import java.io.IOException; +import java.net.InetAddress; +import java.nio.file.Path; +import java.util.Collections; +import java.util.Date; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; +import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.ASN; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CITY_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CITY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CONTINENT_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CONTINENT_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.COUNTRY_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.COUNTRY_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.COUNTRY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.IP; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.IS_COUNTRY_IN_EUROPEAN_UNION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LEAST_SPECIFIED_SUBDIVISION_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LEAST_SPECIFIED_SUBDIVISION_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LOCATION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LOCATION_ACCURACY_RADIUS; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.METRO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.MOST_SPECIFIED_SUBDIVISION_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.MOST_SPECIFIED_SUBDIVISION_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.MOST_SPECIFIED_SUBDIVISION_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.NETWORK; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.ASN_ORGANIZATION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.POSTAL_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.POSTAL_CODE_CONFIDENCE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REGISTERED_COUNTRY_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REGISTERED_COUNTRY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_TYPE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.TIME_ZONE; +import static org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader.LAT; +import static org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader.LON; + +@ExtendWith(MockitoExtension.class) +class GeoIP2DatabaseReaderTest { + public static final String GEOIP2_TEST_MMDB_FILES = "./build/resources/test/mmdb-files/geo-ip2"; + public static final long ASN_RESULT = 12345L; + public static final String ASN_ORG_RESULT = "Example Org"; + private static final String NETWORK_RESULT = "1.2.3.0/24"; + private static final String COUNTRY_NAME_RESULT = "United States"; + private static final String COUNTRY_ISO_CODE_RESULT = "US"; + private static final Boolean COUNTRY_IS_IN_EUROPEAN_UNION_RESULT = false; + private static final String CITY_NAME_RESULT = "Los Angeles"; + private static final Double LATITUDE_RESULT = 12.34; + private static final Double LONGITUDE_RESULT = 56.78; + private static final String TIME_ZONE_RESULT = "America/Los_Angeles"; + private static final Integer METRO_CODE_RESULT = 807; + private static final String POSTAL_CODE_RESULT = "90001"; + private static final String LEAST_SPECIFIED_SUBDIVISION_NAME_RESULT = "California"; + private static final String MOST_SPECIFIED_SUBDIVISION_NAME_RESULT = "California"; + private static final String LEAST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT = "CA"; + private static final String MOST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT = "CA"; + private static final String REGISTERED_COUNTRY_NAME_RESULT = "Argentina"; + private static final String REGISTERED_COUNTRY_ISO_CODE_RESULT = "AR"; + private static final String REPRESENTED_COUNTRY_NAME_RESULT = "Belgium"; + private static final String REPRESENTED_COUNTRY_ISO_CODE_RESULT = "BE"; + private static final String REPRESENTED_COUNTRY_TYPE_RESULT = "military"; + private static final String CONTINENT_NAME_RESULT = "North America"; + private static final String CONTINENT_CODE_RESULT = "123456"; + private static final Map LOCATION_RESULT = Map.of(LAT, LATITUDE_RESULT, LON, LONGITUDE_RESULT); + private static final Integer COUNTRY_CONFIDENCE_RESULT = 100; + private static final Integer CITY_CONFIDENCE_RESULT = 90; + private static final Integer LOCATION_ACCURACY_RADIUS_RESULT = 10; + private static final Integer POSTAL_CODE_CONFIDENCE_RESULT = 85; + private static final Integer MOST_SPECIFIED_SUBDIVISION_CONFIDENCE_RESULT = 75; + private static final Integer LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE_RESULT = 60; + private static final String IP_RESULT = "1.2.3.4"; + + @Mock + private DatabaseReaderBuilder databaseReaderBuilder; + @Mock + private DatabaseReader enterpriseDatabaseReader; + @Mock + private Metadata metadata; + @Mock + private EnterpriseResponse enterpriseResponse; + @Mock + private AsnResponse asnResponse; + @Mock + private Continent continent; + @Mock + private Country country; + @Mock + private RepresentedCountry representedCountry; + @Mock + private Country registeredCountry; + @Mock + private City city; + @Mock + private Location location; + @Mock + private Network network; + @Mock + private Postal postal; + @Mock + private Subdivision leastSpecificSubdivision; + @Mock + private Subdivision mostSpecificSubdivision; + @Mock + private InetAddress inetAddress; + @Mock + private GeoIPFileManager geoIPFileManager; + + @BeforeEach + void setUp() throws IOException { + when(databaseReaderBuilder.buildReader(Path.of(GEOIP2_TEST_MMDB_FILES + File.separator + "GeoIP2-Enterprise-Test.mmdb"), 0)) + .thenReturn(enterpriseDatabaseReader); + + when(enterpriseDatabaseReader.getMetadata()).thenReturn(metadata); + final Date date = new Date(9949107436565L); + when(metadata.getBuildDate()).thenReturn(date); + + lenient().when(enterpriseResponse.getContinent()).thenReturn(continent); + lenient().when(continent.getName()).thenReturn(CONTINENT_NAME_RESULT); + lenient().when(continent.getCode()).thenReturn(CONTINENT_CODE_RESULT); + + lenient().when(enterpriseResponse.getCountry()).thenReturn(country); + lenient().when(country.getName()).thenReturn(COUNTRY_NAME_RESULT); + lenient().when(country.getIsoCode()).thenReturn(COUNTRY_ISO_CODE_RESULT); + lenient().when(country.isInEuropeanUnion()).thenReturn(COUNTRY_IS_IN_EUROPEAN_UNION_RESULT); + lenient().when(country.getConfidence()).thenReturn(COUNTRY_CONFIDENCE_RESULT); + + lenient().when(enterpriseResponse.getRegisteredCountry()).thenReturn(registeredCountry); + lenient().when(registeredCountry.getName()).thenReturn(REGISTERED_COUNTRY_NAME_RESULT); + lenient().when(registeredCountry.getIsoCode()).thenReturn(REGISTERED_COUNTRY_ISO_CODE_RESULT); + + lenient().when(enterpriseResponse.getRepresentedCountry()).thenReturn(representedCountry); + lenient().when(representedCountry.getName()).thenReturn(REPRESENTED_COUNTRY_NAME_RESULT); + lenient().when(representedCountry.getIsoCode()).thenReturn(REPRESENTED_COUNTRY_ISO_CODE_RESULT); + lenient().when(representedCountry.getType()).thenReturn(REPRESENTED_COUNTRY_TYPE_RESULT); + + lenient().when(enterpriseResponse.getCity()).thenReturn(city); + lenient().when(city.getName()).thenReturn(CITY_NAME_RESULT); + lenient().when(city.getConfidence()).thenReturn(CITY_CONFIDENCE_RESULT); + + lenient().when(enterpriseResponse.getLocation()).thenReturn(location); + lenient().when(location.getLatitude()).thenReturn(LATITUDE_RESULT); + lenient().when(location.getLongitude()).thenReturn(LONGITUDE_RESULT); + lenient().when(location.getMetroCode()).thenReturn(METRO_CODE_RESULT); + lenient().when(location.getTimeZone()).thenReturn(TIME_ZONE_RESULT); + lenient().when(location.getAccuracyRadius()).thenReturn(LOCATION_ACCURACY_RADIUS_RESULT); + + lenient().when(enterpriseResponse.getPostal()).thenReturn(postal); + lenient().when(postal.getCode()).thenReturn(POSTAL_CODE_RESULT); + lenient().when(postal.getConfidence()).thenReturn(POSTAL_CODE_CONFIDENCE_RESULT); + + lenient().when(enterpriseResponse.getLeastSpecificSubdivision()).thenReturn(leastSpecificSubdivision); + lenient().when(leastSpecificSubdivision.getName()).thenReturn(LEAST_SPECIFIED_SUBDIVISION_NAME_RESULT); + lenient().when(leastSpecificSubdivision.getIsoCode()).thenReturn(LEAST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT); + lenient().when(leastSpecificSubdivision.getConfidence()).thenReturn(LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE_RESULT); + + lenient().when(enterpriseResponse.getMostSpecificSubdivision()).thenReturn(mostSpecificSubdivision); + lenient().when(mostSpecificSubdivision.getName()).thenReturn(MOST_SPECIFIED_SUBDIVISION_NAME_RESULT); + lenient().when(mostSpecificSubdivision.getIsoCode()).thenReturn(MOST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT); + lenient().when(mostSpecificSubdivision.getConfidence()).thenReturn(MOST_SPECIFIED_SUBDIVISION_CONFIDENCE_RESULT); + + lenient().when(asnResponse.getAutonomousSystemNumber()).thenReturn(ASN_RESULT); + lenient().when(asnResponse.getAutonomousSystemOrganization()).thenReturn(ASN_ORG_RESULT); + lenient().when(asnResponse.getNetwork()).thenReturn(network); + lenient().when(asnResponse.getIpAddress()).thenReturn(IP_RESULT); + lenient().when(network.toString()).thenReturn(NETWORK_RESULT); + } + + @AfterEach + void tearDown() { + verifyNoMoreInteractions(enterpriseDatabaseReader); + } + + GeoIP2DatabaseReader createObjectUnderTest() { + return new GeoIP2DatabaseReader(databaseReaderBuilder, geoIPFileManager, GEOIP2_TEST_MMDB_FILES, 0); + } + + @Test + void test_getGeoData_for_all_fields_in_enterprise_database_add_only_fields_configured() throws IOException, GeoIp2Exception { + final GeoIP2DatabaseReader objectUnderTest = createObjectUnderTest(); + + final List fields = List.of(CONTINENT_NAME, CONTINENT_CODE, COUNTRY_NAME, COUNTRY_ISO_CODE, IS_COUNTRY_IN_EUROPEAN_UNION, + COUNTRY_CONFIDENCE, REPRESENTED_COUNTRY_ISO_CODE, REGISTERED_COUNTRY_ISO_CODE, CITY_NAME, CITY_CONFIDENCE, LOCATION, + LOCATION_ACCURACY_RADIUS, METRO_CODE, TIME_ZONE, POSTAL_CODE, POSTAL_CODE_CONFIDENCE); + final Set databases = Set.of(GeoIPDatabase.ENTERPRISE); + + when(enterpriseDatabaseReader.tryEnterprise(inetAddress)).thenReturn(Optional.of(enterpriseResponse)); + + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + + assertThat(geoData.size(), equalTo(fields.size())); + assertThat(geoData.get(CONTINENT_NAME.getFieldName()), equalTo(CONTINENT_NAME_RESULT)); + assertThat(geoData.get(CONTINENT_CODE.getFieldName()), equalTo(CONTINENT_CODE_RESULT)); + assertThat(geoData.get(COUNTRY_NAME.getFieldName()), equalTo(COUNTRY_NAME_RESULT)); + assertThat(geoData.get(COUNTRY_ISO_CODE.getFieldName()), equalTo(COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(IS_COUNTRY_IN_EUROPEAN_UNION.getFieldName()), equalTo(COUNTRY_IS_IN_EUROPEAN_UNION_RESULT)); + assertThat(geoData.get(COUNTRY_CONFIDENCE.getFieldName()), equalTo(COUNTRY_CONFIDENCE_RESULT)); + assertThat(geoData.get(REGISTERED_COUNTRY_ISO_CODE.getFieldName()), equalTo(REGISTERED_COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_ISO_CODE.getFieldName()), equalTo(REPRESENTED_COUNTRY_ISO_CODE_RESULT)); + + assertThat(geoData.get(CITY_NAME.getFieldName()), equalTo(CITY_NAME_RESULT)); + assertThat(geoData.get(CITY_CONFIDENCE.getFieldName()), equalTo(CITY_CONFIDENCE_RESULT)); + assertThat(geoData.get(LOCATION.getFieldName()), equalTo(LOCATION_RESULT)); + assertThat(geoData.get(LOCATION_ACCURACY_RADIUS.getFieldName()), equalTo(LOCATION_ACCURACY_RADIUS_RESULT)); + assertThat(geoData.get(METRO_CODE.getFieldName()), equalTo(METRO_CODE_RESULT)); + assertThat(geoData.get(TIME_ZONE.getFieldName()), equalTo(TIME_ZONE_RESULT)); + assertThat(geoData.get(POSTAL_CODE.getFieldName()), equalTo(POSTAL_CODE_RESULT)); + assertThat(geoData.get(POSTAL_CODE_CONFIDENCE.getFieldName()), equalTo(POSTAL_CODE_CONFIDENCE_RESULT)); + } + + @Test + void test_getGeoData_for_all_fields_in_enterprise_database_when_no_fields_are_configured() throws IOException, GeoIp2Exception { + final GeoIP2DatabaseReader objectUnderTest = createObjectUnderTest(); + final List fields = Collections.emptyList(); + final Set databases = Set.of(GeoIPDatabase.ENTERPRISE); + + when(enterpriseDatabaseReader.tryEnterprise(inetAddress)).thenReturn(Optional.of(enterpriseResponse)); + + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + + assertThat(geoData.get(CONTINENT_NAME.getFieldName()), equalTo(CONTINENT_NAME_RESULT)); + assertThat(geoData.get(CONTINENT_CODE.getFieldName()), equalTo(CONTINENT_CODE_RESULT)); + assertThat(geoData.get(COUNTRY_NAME.getFieldName()), equalTo(COUNTRY_NAME_RESULT)); + assertThat(geoData.get(COUNTRY_CONFIDENCE.getFieldName()), equalTo(COUNTRY_CONFIDENCE_RESULT)); + assertThat(geoData.get(COUNTRY_ISO_CODE.getFieldName()), equalTo(COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(IS_COUNTRY_IN_EUROPEAN_UNION.getFieldName()), equalTo(COUNTRY_IS_IN_EUROPEAN_UNION_RESULT)); + assertThat(geoData.get(REGISTERED_COUNTRY_NAME.getFieldName()), equalTo(REGISTERED_COUNTRY_NAME_RESULT)); + assertThat(geoData.get(REGISTERED_COUNTRY_ISO_CODE.getFieldName()), equalTo(REGISTERED_COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_NAME.getFieldName()), equalTo(REPRESENTED_COUNTRY_NAME_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_ISO_CODE.getFieldName()), equalTo(REPRESENTED_COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_TYPE.getFieldName()), equalTo(REPRESENTED_COUNTRY_TYPE_RESULT)); + + assertThat(geoData.get(CITY_NAME.getFieldName()), equalTo(CITY_NAME_RESULT)); + assertThat(geoData.get(CITY_CONFIDENCE.getFieldName()), equalTo(CITY_CONFIDENCE_RESULT)); + assertThat(geoData.get(LOCATION.getFieldName()), equalTo(LOCATION_RESULT)); + assertThat(geoData.get(LOCATION_ACCURACY_RADIUS.getFieldName()), equalTo(LOCATION_ACCURACY_RADIUS_RESULT)); + assertThat(geoData.get(METRO_CODE.getFieldName()), equalTo(METRO_CODE_RESULT)); + assertThat(geoData.get(TIME_ZONE.getFieldName()), equalTo(TIME_ZONE_RESULT)); + assertThat(geoData.get(POSTAL_CODE.getFieldName()), equalTo(POSTAL_CODE_RESULT)); + assertThat(geoData.get(POSTAL_CODE_CONFIDENCE.getFieldName()), equalTo(POSTAL_CODE_CONFIDENCE_RESULT)); + assertThat(geoData.get(MOST_SPECIFIED_SUBDIVISION_NAME.getFieldName()), equalTo(MOST_SPECIFIED_SUBDIVISION_NAME_RESULT)); + assertThat(geoData.get(MOST_SPECIFIED_SUBDIVISION_ISO_CODE.getFieldName()), equalTo(MOST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT)); + assertThat(geoData.get(MOST_SPECIFIED_SUBDIVISION_CONFIDENCE.getFieldName()), equalTo(MOST_SPECIFIED_SUBDIVISION_CONFIDENCE_RESULT)); + assertThat(geoData.get(LEAST_SPECIFIED_SUBDIVISION_NAME.getFieldName()), equalTo(LEAST_SPECIFIED_SUBDIVISION_NAME_RESULT)); + assertThat(geoData.get(LEAST_SPECIFIED_SUBDIVISION_ISO_CODE.getFieldName()), equalTo(LEAST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT)); + assertThat(geoData.get(LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE.getFieldName()), equalTo(LEAST_SPECIFIED_SUBDIVISION_CONFIDENCE_RESULT)); + } + + @Test + void test_getGeoData_for_asn_fields_in_enterprise_database_when_no_fields_are_configured() throws IOException, GeoIp2Exception { + final GeoIP2DatabaseReader objectUnderTest = createObjectUnderTest(); + final Set databases = Set.of(GeoIPDatabase.ASN); + final List fields = List.of(ASN, ASN_ORGANIZATION, NETWORK, IP); + + when(enterpriseDatabaseReader.tryAsn(inetAddress)).thenReturn(Optional.of(asnResponse)); + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + + assertThat(geoData.size(), equalTo(fields.size())); + assertThat(geoData.get(ASN.getFieldName()), equalTo(ASN_RESULT)); + assertThat(geoData.get(ASN_ORGANIZATION.getFieldName()), equalTo(ASN_ORG_RESULT)); + assertThat(geoData.get(NETWORK.getFieldName()), equalTo(NETWORK_RESULT)); + assertThat(geoData.get(IP.getFieldName()), equalTo(IP_RESULT)); + } + + @Test + void test_database_expired_should_return_false_when_expiry_date_is_in_future() { + final Date date = new Date(9949107436565L); + when(metadata.getBuildDate()).thenReturn(date); + final GeoIP2DatabaseReader objectUnderTest = createObjectUnderTest(); + + assertThat(objectUnderTest.isExpired(), equalTo(false)); + } + + @Test + void test_database_expired_should_return_true_when_expiry_date_is_in_past() throws IOException { + final Date date = new Date(91911199999L); + when(metadata.getBuildDate()).thenReturn(date); + doNothing().when(geoIPFileManager).deleteDirectory(any()); + + final GeoIP2DatabaseReader objectUnderTest = createObjectUnderTest(); + + assertThat(objectUnderTest.isExpired(), equalTo(true)); + verify(enterpriseDatabaseReader).close(); + } + + @Test + void test_database_close_should_close_the_reader() throws IOException { + doNothing().when(geoIPFileManager).deleteDirectory(any()); + final GeoIP2DatabaseReader objectUnderTest = createObjectUnderTest(); + + objectUnderTest.close(); + + assertThat(objectUnderTest.isExpired(), equalTo(false)); + verify(enterpriseDatabaseReader).close(); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoLite2DatabaseReaderTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoLite2DatabaseReaderTest.java new file mode 100644 index 0000000000..54f6bb573c --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GeoLite2DatabaseReaderTest.java @@ -0,0 +1,384 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.databaseenrich; + +import com.maxmind.db.Metadata; +import com.maxmind.db.Network; +import com.maxmind.geoip2.DatabaseReader; +import com.maxmind.geoip2.exception.GeoIp2Exception; +import com.maxmind.geoip2.model.AsnResponse; +import com.maxmind.geoip2.model.CityResponse; +import com.maxmind.geoip2.model.CountryResponse; +import com.maxmind.geoip2.record.City; +import com.maxmind.geoip2.record.Continent; +import com.maxmind.geoip2.record.Country; +import com.maxmind.geoip2.record.Location; +import com.maxmind.geoip2.record.Postal; +import com.maxmind.geoip2.record.RepresentedCountry; +import com.maxmind.geoip2.record.Subdivision; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.plugins.processor.GeoIPDatabase; +import org.opensearch.dataprepper.plugins.processor.GeoIPField; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPFileManager; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DatabaseReaderBuilder; + +import java.io.File; +import java.io.IOException; +import java.net.InetAddress; +import java.nio.file.Path; +import java.util.Date; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; +import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.plugins.processor.GeoIPDatabase.CITY; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.ASN; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CITY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CONTINENT_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.CONTINENT_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.COUNTRY_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.COUNTRY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.IP; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.IS_COUNTRY_IN_EUROPEAN_UNION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LATITUDE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LEAST_SPECIFIED_SUBDIVISION_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LEAST_SPECIFIED_SUBDIVISION_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LOCATION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LOCATION_ACCURACY_RADIUS; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.LONGITUDE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.METRO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.MOST_SPECIFIED_SUBDIVISION_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.MOST_SPECIFIED_SUBDIVISION_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.NETWORK; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.ASN_ORGANIZATION; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.POSTAL_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REGISTERED_COUNTRY_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REGISTERED_COUNTRY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_ISO_CODE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_NAME; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.REPRESENTED_COUNTRY_TYPE; +import static org.opensearch.dataprepper.plugins.processor.GeoIPField.TIME_ZONE; +import static org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader.LAT; +import static org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader.LON; + + +@ExtendWith(MockitoExtension.class) +class GeoLite2DatabaseReaderTest { + public static final String GEOLITE2_TEST_MMDB_FILES = "./build/resources/test/mmdb-files/geo-lite2"; + public static final long ASN_RESULT = 12345L; + public static final String ASN_ORG_RESULT = "Example Org"; + private static final String NETWORK_RESULT = "1.2.3.0/24"; + private static final String COUNTRY_NAME_RESULT = "United States"; + private static final String COUNTRY_ISO_CODE_RESULT = "US"; + private static final Boolean COUNTRY_IS_IN_EUROPEAN_UNION_RESULT = false; + private static final String CITY_NAME_RESULT = "Los Angeles"; + private static final Double LATITUDE_RESULT = 12.34; + private static final Double LONGITUDE_RESULT = 56.78; + private static final Integer LOCATION_ACCURACY_RADIUS_RESULT = 10; + private static final String TIME_ZONE_RESULT = "America/Los_Angeles"; + private static final Integer METRO_CODE_RESULT = 807; + private static final String POSTAL_CODE_RESULT = "90001"; + private static final String LEAST_SPECIFIED_SUBDIVISION_NAME_RESULT = "California"; + private static final String MOST_SPECIFIED_SUBDIVISION_NAME_RESULT = "California"; + private static final String LEAST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT = "CA"; + private static final String MOST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT = "CA"; + private static final String REGISTERED_COUNTRY_NAME_RESULT = "Argentina"; + private static final String REGISTERED_COUNTRY_ISO_CODE_RESULT = "AR"; + private static final String REPRESENTED_COUNTRY_NAME_RESULT = "Belgium"; + private static final String REPRESENTED_COUNTRY_ISO_CODE_RESULT = "BE"; + private static final String REPRESENTED_COUNTRY_TYPE_RESULT = "military"; + private static final String CONTINENT_NAME_RESULT = "North America"; + private static final String CONTINENT_CODE_RESULT = "123456"; + private static final Map LOCATION_RESULT = Map.of(LAT, LATITUDE_RESULT, LON, LONGITUDE_RESULT); + private static final String IP_RESULT = "1.2.3.4"; + + @Mock + private DatabaseReaderBuilder databaseReaderBuilder; + @Mock + private DatabaseReader countryDatabaseReader; + @Mock + private DatabaseReader cityDatabaseReader; + @Mock + private DatabaseReader asnDatabaseReader; + @Mock + private Metadata metadata; + @Mock + private CountryResponse countryResponse; + @Mock + private CityResponse cityResponse; + @Mock + private AsnResponse asnResponse; + @Mock + private Continent continent; + @Mock + private Country country; + @Mock + private RepresentedCountry representedCountry; + @Mock + private Country registeredCountry; + @Mock + private City city; + @Mock + private Location location; + @Mock + private Network network; + @Mock + private Postal postal; + @Mock + private Subdivision leastSpecificSubdivision; + @Mock + private Subdivision mostSpecificSubdivision; + @Mock + private InetAddress inetAddress; + @Mock + private GeoIPFileManager geoIPFileManager; + + @BeforeEach + void setUp() throws IOException { + when(databaseReaderBuilder.buildReader(Path.of(GEOLITE2_TEST_MMDB_FILES + File.separator + "GeoLite2-Country-Test.mmdb"), 0)) + .thenReturn(countryDatabaseReader); + when(databaseReaderBuilder.buildReader(Path.of(GEOLITE2_TEST_MMDB_FILES + File.separator + "GeoLite2-City-Test.mmdb"), 0)) + .thenReturn(cityDatabaseReader); + when(databaseReaderBuilder.buildReader(Path.of(GEOLITE2_TEST_MMDB_FILES + File.separator + "GeoLite2-ASN-Test.mmdb"), 0)) + .thenReturn(asnDatabaseReader); + + when(countryDatabaseReader.getMetadata()).thenReturn(metadata); + when(cityDatabaseReader.getMetadata()).thenReturn(metadata); + when(asnDatabaseReader.getMetadata()).thenReturn(metadata); + final Date date = new Date(9949107436565L); + when(metadata.getBuildDate()).thenReturn(date); + + lenient().when(countryResponse.getContinent()).thenReturn(continent); + lenient().when(continent.getName()).thenReturn(CONTINENT_NAME_RESULT); + lenient().when(continent.getCode()).thenReturn(CONTINENT_CODE_RESULT); + + lenient().when(countryResponse.getCountry()).thenReturn(country); + lenient().when(country.getName()).thenReturn(COUNTRY_NAME_RESULT); + lenient().when(country.getIsoCode()).thenReturn(COUNTRY_ISO_CODE_RESULT); + lenient().when(country.isInEuropeanUnion()).thenReturn(COUNTRY_IS_IN_EUROPEAN_UNION_RESULT); + + lenient().when(countryResponse.getRegisteredCountry()).thenReturn(registeredCountry); + lenient().when(registeredCountry.getName()).thenReturn(REGISTERED_COUNTRY_NAME_RESULT); + lenient().when(registeredCountry.getIsoCode()).thenReturn(REGISTERED_COUNTRY_ISO_CODE_RESULT); + + lenient().when(countryResponse.getRepresentedCountry()).thenReturn(representedCountry); + lenient().when(representedCountry.getName()).thenReturn(REPRESENTED_COUNTRY_NAME_RESULT); + lenient().when(representedCountry.getIsoCode()).thenReturn(REPRESENTED_COUNTRY_ISO_CODE_RESULT); + lenient().when(representedCountry.getType()).thenReturn(REPRESENTED_COUNTRY_TYPE_RESULT); + + lenient().when(cityResponse.getCity()).thenReturn(city); + lenient().when(city.getName()).thenReturn(CITY_NAME_RESULT); + + lenient().when(cityResponse.getContinent()).thenReturn(continent); + lenient().when(cityResponse.getCountry()).thenReturn(country); + lenient().when(cityResponse.getRegisteredCountry()).thenReturn(registeredCountry); + lenient().when(cityResponse.getRepresentedCountry()).thenReturn(representedCountry); + + lenient().when(cityResponse.getLocation()).thenReturn(location); + lenient().when(location.getLatitude()).thenReturn(LATITUDE_RESULT); + lenient().when(location.getLongitude()).thenReturn(LONGITUDE_RESULT); + lenient().when(location.getMetroCode()).thenReturn(METRO_CODE_RESULT); + lenient().when(location.getTimeZone()).thenReturn(TIME_ZONE_RESULT); + lenient().when(location.getAccuracyRadius()).thenReturn(LOCATION_ACCURACY_RADIUS_RESULT); + + lenient().when(cityResponse.getPostal()).thenReturn(postal); + lenient().when(postal.getCode()).thenReturn(POSTAL_CODE_RESULT); + + lenient().when(cityResponse.getLeastSpecificSubdivision()).thenReturn(leastSpecificSubdivision); + lenient().when(leastSpecificSubdivision.getName()).thenReturn(LEAST_SPECIFIED_SUBDIVISION_NAME_RESULT); + lenient().when(leastSpecificSubdivision.getIsoCode()).thenReturn(LEAST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT); + + lenient().when(cityResponse.getMostSpecificSubdivision()).thenReturn(mostSpecificSubdivision); + lenient().when(mostSpecificSubdivision.getName()).thenReturn(MOST_SPECIFIED_SUBDIVISION_NAME_RESULT); + lenient().when(mostSpecificSubdivision.getIsoCode()).thenReturn(MOST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT); + + lenient().when(asnResponse.getAutonomousSystemNumber()).thenReturn(ASN_RESULT); + lenient().when(asnResponse.getAutonomousSystemOrganization()).thenReturn(ASN_ORG_RESULT); + lenient().when(asnResponse.getNetwork()).thenReturn(network); + lenient().when(asnResponse.getIpAddress()).thenReturn(IP_RESULT); + lenient().when(network.toString()).thenReturn(NETWORK_RESULT); + } + + @AfterEach + void tearDown() { + verifyNoMoreInteractions(countryDatabaseReader); + verifyNoMoreInteractions(cityDatabaseReader); + verifyNoMoreInteractions(asnDatabaseReader); + } + + GeoLite2DatabaseReader createObjectUnderTest() { + return new GeoLite2DatabaseReader(databaseReaderBuilder, geoIPFileManager, GEOLITE2_TEST_MMDB_FILES, 0); + } + + @Test + void test_getGeoData_for_all_fields_in_country_database() throws IOException, GeoIp2Exception { + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + final List fields = List.of(CONTINENT_NAME, CONTINENT_CODE, COUNTRY_NAME, COUNTRY_ISO_CODE, IS_COUNTRY_IN_EUROPEAN_UNION, + REPRESENTED_COUNTRY_NAME, REPRESENTED_COUNTRY_ISO_CODE, REPRESENTED_COUNTRY_TYPE, REGISTERED_COUNTRY_NAME, + REGISTERED_COUNTRY_ISO_CODE); + final Set databases = Set.of(GeoIPDatabase.COUNTRY); + + when(countryDatabaseReader.tryCountry(inetAddress)).thenReturn(Optional.of(countryResponse)); + + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + assertThat(geoData.size(), equalTo(fields.size())); + assertThat(geoData.get(CONTINENT_NAME.getFieldName()), equalTo(CONTINENT_NAME_RESULT)); + assertThat(geoData.get(CONTINENT_CODE.getFieldName()), equalTo(CONTINENT_CODE_RESULT)); + assertThat(geoData.get(COUNTRY_NAME.getFieldName()), equalTo(COUNTRY_NAME_RESULT)); + assertThat(geoData.get(COUNTRY_ISO_CODE.getFieldName()), equalTo(COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(IS_COUNTRY_IN_EUROPEAN_UNION.getFieldName()), equalTo(COUNTRY_IS_IN_EUROPEAN_UNION_RESULT)); + assertThat(geoData.get(REGISTERED_COUNTRY_NAME.getFieldName()), equalTo(REGISTERED_COUNTRY_NAME_RESULT)); + assertThat(geoData.get(REGISTERED_COUNTRY_ISO_CODE.getFieldName()), equalTo(REGISTERED_COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_NAME.getFieldName()), equalTo(REPRESENTED_COUNTRY_NAME_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_ISO_CODE.getFieldName()), equalTo(REPRESENTED_COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_TYPE.getFieldName()), equalTo(REPRESENTED_COUNTRY_TYPE_RESULT)); + } + + @Test + void test_getGeoData_for_all_fields_in_city_database() throws IOException, GeoIp2Exception { + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + final List fields = List.of(CONTINENT_NAME, CONTINENT_CODE, COUNTRY_NAME, COUNTRY_ISO_CODE, IS_COUNTRY_IN_EUROPEAN_UNION, + REPRESENTED_COUNTRY_NAME, REPRESENTED_COUNTRY_ISO_CODE, REPRESENTED_COUNTRY_TYPE, REGISTERED_COUNTRY_NAME, + REGISTERED_COUNTRY_ISO_CODE, CITY_NAME, LOCATION, LOCATION_ACCURACY_RADIUS, LATITUDE, LONGITUDE, METRO_CODE, TIME_ZONE, POSTAL_CODE, + MOST_SPECIFIED_SUBDIVISION_NAME, MOST_SPECIFIED_SUBDIVISION_ISO_CODE, LEAST_SPECIFIED_SUBDIVISION_NAME, + LEAST_SPECIFIED_SUBDIVISION_ISO_CODE); + final Set databases = Set.of(CITY); + + when(cityDatabaseReader.tryCity(inetAddress)).thenReturn(Optional.of(cityResponse)); + + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + + assertThat(geoData.size(), equalTo(fields.size())); + assertThat(geoData.get(CONTINENT_NAME.getFieldName()), equalTo(CONTINENT_NAME_RESULT)); + assertThat(geoData.get(CONTINENT_CODE.getFieldName()), equalTo(CONTINENT_CODE_RESULT)); + assertThat(geoData.get(COUNTRY_NAME.getFieldName()), equalTo(COUNTRY_NAME_RESULT)); + assertThat(geoData.get(COUNTRY_ISO_CODE.getFieldName()), equalTo(COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(IS_COUNTRY_IN_EUROPEAN_UNION.getFieldName()), equalTo(COUNTRY_IS_IN_EUROPEAN_UNION_RESULT)); + assertThat(geoData.get(REGISTERED_COUNTRY_NAME.getFieldName()), equalTo(REGISTERED_COUNTRY_NAME_RESULT)); + assertThat(geoData.get(REGISTERED_COUNTRY_ISO_CODE.getFieldName()), equalTo(REGISTERED_COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_NAME.getFieldName()), equalTo(REPRESENTED_COUNTRY_NAME_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_ISO_CODE.getFieldName()), equalTo(REPRESENTED_COUNTRY_ISO_CODE_RESULT)); + assertThat(geoData.get(REPRESENTED_COUNTRY_TYPE.getFieldName()), equalTo(REPRESENTED_COUNTRY_TYPE_RESULT)); + + assertThat(geoData.get(CITY_NAME.getFieldName()), equalTo(CITY_NAME_RESULT)); + assertThat(geoData.get(LOCATION.getFieldName()), equalTo(LOCATION_RESULT)); + assertThat(geoData.get(LATITUDE.getFieldName()), equalTo(LATITUDE_RESULT)); + assertThat(geoData.get(LONGITUDE.getFieldName()), equalTo(LONGITUDE_RESULT)); + assertThat(geoData.get(LOCATION_ACCURACY_RADIUS.getFieldName()), equalTo(LOCATION_ACCURACY_RADIUS_RESULT)); + assertThat(geoData.get(METRO_CODE.getFieldName()), equalTo(METRO_CODE_RESULT)); + assertThat(geoData.get(TIME_ZONE.getFieldName()), equalTo(TIME_ZONE_RESULT)); + assertThat(geoData.get(POSTAL_CODE.getFieldName()), equalTo(POSTAL_CODE_RESULT)); + assertThat(geoData.get(MOST_SPECIFIED_SUBDIVISION_NAME.getFieldName()), equalTo(MOST_SPECIFIED_SUBDIVISION_NAME_RESULT)); + assertThat(geoData.get(MOST_SPECIFIED_SUBDIVISION_ISO_CODE.getFieldName()), equalTo(MOST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT)); + assertThat(geoData.get(LEAST_SPECIFIED_SUBDIVISION_NAME.getFieldName()), equalTo(LEAST_SPECIFIED_SUBDIVISION_NAME_RESULT)); + assertThat(geoData.get(LEAST_SPECIFIED_SUBDIVISION_ISO_CODE.getFieldName()), equalTo(LEAST_SPECIFIED_SUBDIVISION_ISO_CODE_RESULT)); + } + + @Test + void test_getGeoData_for_all_fields_in_asn_database() throws IOException, GeoIp2Exception { + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + final List fields = List.of(ASN, ASN_ORGANIZATION, NETWORK, IP); + final Set databases = Set.of(GeoIPDatabase.ASN); + + when(asnDatabaseReader.tryAsn(inetAddress)).thenReturn(Optional.of(asnResponse)); + + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + + assertThat(geoData.size(), equalTo(fields.size())); + assertThat(geoData.get(ASN.getFieldName()), equalTo(ASN_RESULT)); + assertThat(geoData.get(ASN_ORGANIZATION.getFieldName()), equalTo(ASN_ORG_RESULT)); + assertThat(geoData.get(NETWORK.getFieldName()), equalTo(NETWORK_RESULT)); + assertThat(geoData.get(IP.getFieldName()), equalTo(IP_RESULT)); + } + + @Test + void test_getGeoData_for_country_database_should_not_add_any_fields_if_country_is_not_required() { + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + final List fields = List.of(CONTINENT_NAME, CONTINENT_CODE, COUNTRY_NAME); + final Set databases = Set.of(); + + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + assertThat(geoData.size(), equalTo(0)); + } + + @Test + void test_getGeoData_for_city_database_should_not_add_any_fields_if_city_is_not_required() { + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + final List fields = List.of(CONTINENT_NAME, CONTINENT_CODE, COUNTRY_NAME); + final Set databases = Set.of(); + + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + + assertThat(geoData.size(), equalTo(0)); + } + + @Test + void test_getGeoData_for_asn_database_should_not_add_any_fields_if_asn_is_not_required() { + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + final List fields = List.of(ASN, ASN_ORGANIZATION, NETWORK); + final Set databases = Set.of(); + + final Map geoData = objectUnderTest.getGeoData(inetAddress, fields, databases); + + assertThat(geoData.size(), equalTo(0)); + } + + @Test + void test_database_expired_should_return_false_when_expiry_date_is_in_future() { + final Date date = new Date(9949107436565L); + when(metadata.getBuildDate()).thenReturn(date); + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + assertThat(objectUnderTest.isExpired(), equalTo(false)); + } + + @Test + void test_database_expired_should_return_true_when_expiry_date_is_in_past() throws IOException { + final Date date = new Date(91911199999L); + when(metadata.getBuildDate()).thenReturn(date); + doNothing().when(geoIPFileManager).deleteFile(any()); + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + assertThat(objectUnderTest.isExpired(), equalTo(true)); + verify(countryDatabaseReader).close(); + verify(cityDatabaseReader).close(); + verify(asnDatabaseReader).close(); + } + + @Test + void test_database_close_should_close_the_reader() throws IOException { + doNothing().when(geoIPFileManager).deleteDirectory(any()); + final GeoLite2DatabaseReader objectUnderTest = createObjectUnderTest(); + + objectUnderTest.close(); + + assertThat(objectUnderTest.isExpired(), equalTo(false)); + verify(countryDatabaseReader).close(); + verify(cityDatabaseReader).close(); + verify(asnDatabaseReader).close(); + } + +} diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoIP2DataTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoIP2DataTest.java deleted file mode 100644 index 5493082d08..0000000000 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoIP2DataTest.java +++ /dev/null @@ -1,102 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.databaseenrich; - -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Disabled; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.plugins.processor.GeoIPProcessorConfig; -import org.opensearch.dataprepper.plugins.processor.extension.GeoIPProcessorService; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSource; -import org.opensearch.dataprepper.plugins.processor.extension.GeoIpServiceConfig; -import org.opensearch.dataprepper.plugins.processor.extension.MaxMindConfig; - -import java.io.File; -import java.net.InetAddress; -import java.net.UnknownHostException; -import java.util.List; -import java.util.Map; - -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.Mockito.when; - -@ExtendWith(MockitoExtension.class) -class GetGeoIP2DataTest { - - private static final String PATH = "./src/test/resources/mmdb-file/geo-enterprise"; - private String tempFolderPath = System.getProperty("java.io.tmpdir") + File.separator + "GeoIP"; - private static final String PREFIX_DIR = "first_database"; - public static final int REFRESH_SCHEDULE = 10; - public static final String IP = "2001:4860:4860::8888"; - @Mock - private GeoIPProcessorConfig geoIPProcessorConfig; - @Mock - private GeoIpServiceConfig geoIpServiceConfig; - @Mock - private MaxMindConfig maxMindConfig; - @Mock - private DBSource downloadSource; - private GetGeoIP2Data getGeoIP2Data; - private final int cacheSize = 4068; - - @BeforeEach - void setUp() { - when(geoIpServiceConfig.getMaxMindConfig()) - .thenReturn(maxMindConfig); - String dbPath = "./src/test/resources/mmdb-file/geo-enterprise"; - getGeoIP2Data = new GetGeoIP2Data(dbPath, cacheSize); - } - - @Disabled("Doesn't have valid GeoIP2-Enterprise.mmdb") - @Test - void getGeoDataTest() throws UnknownHostException { - - List attributes = List.of("city_name", "country_name"); - InetAddress inetAddress = InetAddress.getByName(IP); - GeoIPProcessorService.downloadReady = false; - Map geoData = getGeoIP2Data.getGeoData(inetAddress, attributes, PREFIX_DIR); - assertThat(geoData.get("country_iso_code"), equalTo("US")); - assertThat(geoData.get("ip"), equalTo("2001:4860:4860:0:0:0:0:8888")); - } - - @Test - @Disabled - void switchDatabaseReaderTest() { - tempFolderPath = System.getProperty("java.io.tmpdir") + File.separator + "GeoIPMaxmind"; - DBSource.createFolderIfNotExist(tempFolderPath); - getGeoIP2Data = new GetGeoIP2Data(tempFolderPath, cacheSize); - assertDoesNotThrow(() -> { - getGeoIP2Data.switchDatabaseReader(); - }); - assertDoesNotThrow(() -> { - getGeoIP2Data.closeReader(); - }); - } - - @Test - @Disabled - void getGeoDataTest_cover_EnrichFailedException() throws UnknownHostException { - List attributes = List.of("city_name", "country_name"); - InetAddress inetAddress = InetAddress.getByName(IP); - GeoIPProcessorService.downloadReady = false; - assertThrows(EnrichFailedException.class, () -> getGeoIP2Data.getGeoData(inetAddress, attributes, PREFIX_DIR)); - } - - @Test - @Disabled - void closeReaderTest() throws UnknownHostException { - getGeoIP2Data.closeReader(); - List attributes = List.of("city_name", "country_name"); - InetAddress inetAddress = InetAddress.getByName(IP); - assertThrows(EnrichFailedException.class, () -> getGeoIP2Data.getGeoData(inetAddress, attributes, PREFIX_DIR)); - } -} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoLite2DataTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoLite2DataTest.java deleted file mode 100644 index 7a58317a5d..0000000000 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/databaseenrich/GetGeoLite2DataTest.java +++ /dev/null @@ -1,129 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.databaseenrich; - -import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Disabled; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.plugins.processor.extension.GeoIPProcessorService; -import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.DBSource; -import org.opensearch.dataprepper.plugins.processor.extension.GeoIpServiceConfig; -import org.opensearch.dataprepper.plugins.processor.extension.MaxMindConfig; - -import java.io.File; -import java.net.InetAddress; -import java.net.UnknownHostException; -import java.util.List; -import java.util.Map; - -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.Mockito.when; - -@ExtendWith(MockitoExtension.class) -class GetGeoLite2DataTest { - - private static final String PATH = "./src/test/resources/mmdb-file/geo-lite2"; - public static final String IP = "2a02:ec00:0:0:0:0:0:0"; - private static final String PREFIX_DIR = "first_database"; - private String tempFolderPath = System.getProperty("java.io.tmpdir") + File.separator + "GeoIP"; - @Mock - private GeoIpServiceConfig geoIpServiceConfig; - @Mock - private MaxMindConfig maxMindConfig; - private GetGeoLite2Data getGeoLite2Data; - private final int cacheSize = 4068; - - @BeforeEach - void setUp() { - when(geoIpServiceConfig.getMaxMindConfig()).thenReturn(maxMindConfig); - when(maxMindConfig.getCacheSize()).thenReturn(cacheSize); - - final String dbPath = "./src/test/resources/mmdb-file/geo-lite2"; - getGeoLite2Data = new GetGeoLite2Data(dbPath, cacheSize); - } - - @Test - @Disabled - void getGeoDataTest_without_attributes() throws UnknownHostException { - List attributes = List.of(); - InetAddress inetAddress = InetAddress.getByName(IP); - GeoIPProcessorService.downloadReady = false; - Map geoData = getGeoLite2Data.getGeoData(inetAddress, attributes, tempFolderPath); - Assertions.assertNotNull(geoData); - assertThat(geoData.get("country_iso_code"), equalTo("FR")); - assertThat(geoData.get("ip"), equalTo(IP)); - assertDoesNotThrow(() -> { - getGeoLite2Data.closeReader(); - }); - } - - @Test - @Disabled - void getGeoDataTest_with_attributes() throws UnknownHostException { - List attributes = List.of("city_name", - "country_name", - "ip", - "country_iso_code", - "continent_name", - "region_iso_code", - "region_name", - "timezone", - "location", - "asn", - "organization_name", "network"); - InetAddress inetAddress = InetAddress.getByName(IP); - GeoIPProcessorService.downloadReady = false; - Map geoData = getGeoLite2Data.getGeoData(inetAddress, attributes, tempFolderPath); - Assertions.assertNotNull(geoData); - assertThat(geoData.get("country_name"), equalTo("France")); - assertThat(geoData.get("ip"), equalTo(IP)); - assertDoesNotThrow(() -> { - getGeoLite2Data.closeReader(); - }); - } - - @Test - @Disabled - void getGeoDataTest_cover_EnrichFailedException() throws UnknownHostException { - List attributes = List.of(); - InetAddress inetAddress = InetAddress.getByName(IP); - String dbPath = "/src/test/resources/mmdb-file/geo-enterprise"; - getGeoLite2Data = new GetGeoLite2Data(dbPath, cacheSize); - GeoIPProcessorService.downloadReady = false; - assertThrows(EnrichFailedException.class, () -> getGeoLite2Data.getGeoData(inetAddress, attributes, tempFolderPath)); - } - - @Test - @Disabled - void switchDatabaseReaderTest() { - tempFolderPath = System.getProperty("java.io.tmpdir") + File.separator + "GeoIPMaxmind"; - DBSource.createFolderIfNotExist(tempFolderPath); - getGeoLite2Data = new GetGeoLite2Data(tempFolderPath, cacheSize); - assertDoesNotThrow(() -> { - getGeoLite2Data.switchDatabaseReader(); - }); - assertDoesNotThrow(() -> { - getGeoLite2Data.closeReader(); - }); - } - - @Test - @Disabled - void closeReaderTest() throws UnknownHostException { - // While closing the readers, all the readers reset to null. - getGeoLite2Data.closeReader(); - List attributes = List.of("city_name", "country_name"); - InetAddress inetAddress = InetAddress.getByName(IP); - assertThrows(NullPointerException.class, () -> getGeoLite2Data.getGeoData(inetAddress, attributes, PREFIX_DIR)); - } -} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/DatabaseReaderInitializationExceptionTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/DatabaseReaderInitializationExceptionTest.java new file mode 100644 index 0000000000..d5a3edbabc --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/DatabaseReaderInitializationExceptionTest.java @@ -0,0 +1,33 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.UUID; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +class DatabaseReaderInitializationExceptionTest { + private String message; + + @BeforeEach + void setUp() { + message = UUID.randomUUID().toString(); + } + + private DatabaseReaderInitializationException createObjectUnderTest() { + return new DatabaseReaderInitializationException(message); + } + + @Test + void getMessage_returns_message() { + assertThat(createObjectUnderTest().getMessage(), equalTo(message)); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/DownloadFailedExceptionTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/DownloadFailedExceptionTest.java new file mode 100644 index 0000000000..ca12b8a8d0 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/DownloadFailedExceptionTest.java @@ -0,0 +1,33 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.UUID; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +class DownloadFailedExceptionTest { + private String message; + + @BeforeEach + void setUp() { + message = UUID.randomUUID().toString(); + } + + private DownloadFailedException createObjectUnderTest() { + return new DownloadFailedException(message); + } + + @Test + void getMessage_returns_message() { + assertThat(createObjectUnderTest().getMessage(), equalTo(message)); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/EngineFailureExceptionTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/EngineFailureExceptionTest.java new file mode 100644 index 0000000000..b1a6f4b822 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/EngineFailureExceptionTest.java @@ -0,0 +1,33 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.UUID; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +class EngineFailureExceptionTest { + private String message; + @BeforeEach + void setUp() { + message = UUID.randomUUID().toString(); + } + + private EngineFailureException createObjectUnderTest() { + return new EngineFailureException(message); + } + + @Test + void getMessage_returns_message() { + assertThat(createObjectUnderTest().getMessage(), equalTo(message)); + } + + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/EnrichFailedExceptionTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/EnrichFailedExceptionTest.java new file mode 100644 index 0000000000..af31a9532e --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/EnrichFailedExceptionTest.java @@ -0,0 +1,33 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.UUID; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +class EnrichFailedExceptionTest { + private String message; + + @BeforeEach + void setUp() { + message = UUID.randomUUID().toString(); + } + + private EnrichFailedException createObjectUnderTest() { + return new EnrichFailedException(message); + } + + @Test + void getMessage_returns_message() { + assertThat(createObjectUnderTest().getMessage(), equalTo(message)); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/InvalidIPAddressExceptionTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/InvalidIPAddressExceptionTest.java new file mode 100644 index 0000000000..d1f9e0c517 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/InvalidIPAddressExceptionTest.java @@ -0,0 +1,32 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.UUID; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +class InvalidIPAddressExceptionTest { + private String message; + @BeforeEach + void setUp() { + message = UUID.randomUUID().toString(); + } + + private InvalidIPAddressException createObjectUnderTest() { + return new InvalidIPAddressException(message); + } + + @Test + void getMessage_returns_message() { + assertThat(createObjectUnderTest().getMessage(), equalTo(message)); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/NoValidDatabaseFoundExceptionTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/NoValidDatabaseFoundExceptionTest.java new file mode 100644 index 0000000000..4463d10141 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/exception/NoValidDatabaseFoundExceptionTest.java @@ -0,0 +1,32 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.exception; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.UUID; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +class NoValidDatabaseFoundExceptionTest { + private String message; + + @BeforeEach + void setUp() { + message = UUID.randomUUID().toString(); + } + + private NoValidDatabaseFoundException createObjectUnderTest() { + return new NoValidDatabaseFoundException(message); + } + + @Test + void getMessage_returns_message() { + assertThat(createObjectUnderTest().getMessage(), equalTo(message)); + } +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/DefaultGeoIpConfigSupplierTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/DefaultGeoIpConfigSupplierTest.java index 9d42f5b6f2..c5f484a769 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/DefaultGeoIpConfigSupplierTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/DefaultGeoIpConfigSupplierTest.java @@ -10,10 +10,15 @@ import org.mockito.Mock; import org.mockito.MockedConstruction; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPDatabaseManager; + +import java.util.Optional; +import java.util.concurrent.locks.ReentrantReadWriteLock; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.instanceOf; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.mockConstruction; @ExtendWith(MockitoExtension.class) @@ -21,8 +26,14 @@ class DefaultGeoIpConfigSupplierTest { @Mock private GeoIpServiceConfig geoIpServiceConfig; + @Mock + private GeoIPDatabaseManager geoIPDatabaseManager; + + @Mock + private ReentrantReadWriteLock.ReadLock readLock; + private DefaultGeoIpConfigSupplier createObjectUnderTest() { - return new DefaultGeoIpConfigSupplier(geoIpServiceConfig); + return new DefaultGeoIpConfigSupplier(geoIpServiceConfig, geoIPDatabaseManager, readLock); } @Test @@ -30,11 +41,24 @@ void getGeoIpProcessorService_returns_geoIPProcessorService() { try (final MockedConstruction mockedConstruction = mockConstruction(GeoIPProcessorService.class)) { final DefaultGeoIpConfigSupplier objectUnderTest = createObjectUnderTest(); - final GeoIPProcessorService geoIPProcessorService = objectUnderTest.getGeoIPProcessorService(); + final Optional geoIPProcessorService = objectUnderTest.getGeoIPProcessorService(); assertThat(mockedConstruction.constructed().size(), equalTo(1)); - assertThat(geoIPProcessorService, instanceOf(GeoIPProcessorService.class)); - assertThat(geoIPProcessorService, equalTo(mockedConstruction.constructed().get(0))); + assertTrue(geoIPProcessorService.isPresent()); + assertThat(geoIPProcessorService.get(), instanceOf(GeoIPProcessorService.class)); + assertThat(geoIPProcessorService.get(), equalTo(mockedConstruction.constructed().get(0))); + } + } + + @Test + void getGeoIpProcessorService_returns_empty_if_service_config_is_null() { + try (final MockedConstruction mockedConstruction = + mockConstruction(GeoIPProcessorService.class)) { + final DefaultGeoIpConfigSupplier objectUnderTest = new DefaultGeoIpConfigSupplier(null, geoIPDatabaseManager, readLock); + final Optional geoIPProcessorService = objectUnderTest.getGeoIPProcessorService(); + + assertThat(mockedConstruction.constructed().size(), equalTo(0)); + assertTrue(geoIPProcessorService.isEmpty()); } } } diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIPProcessorServiceTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIPProcessorServiceTest.java index 232f98edec..6a360fe8d3 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIPProcessorServiceTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIPProcessorServiceTest.java @@ -6,89 +6,107 @@ package org.opensearch.dataprepper.plugins.processor.extension; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.plugins.processor.GeoIPProcessorConfig; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.GetGeoData; -import org.opensearch.dataprepper.test.helper.ReflectivelySetField; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader; +import org.opensearch.dataprepper.plugins.processor.extension.databasedownload.GeoIPDatabaseManager; -import java.io.File; -import java.net.InetAddress; -import java.net.UnknownHostException; -import java.util.HashMap; -import java.util.List; -import java.util.Map; +import java.time.Duration; +import java.time.Instant; +import java.util.concurrent.locks.ReentrantReadWriteLock; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; @ExtendWith(MockitoExtension.class) -@Disabled class GeoIPProcessorServiceTest { - - private static final String S3_URL = "https://mybucket10012023.s3.amazonaws.com/GeoLite2/"; - private static final String URL = "https://download.maxmind.com/app/geoip_download?edition_" + - "id=GeoLite2-ASN&suffix=tar.gz"; - private static final String PATH = "./src/test/resources/mmdb-file/geo-lite2"; - private static final String S3_REGION = "us-east-1"; - public static final int REFRESH_SCHEDULE = 10; - public static final String IP = "2001:4860:4860::8888"; - String tempFolderPath = System.getProperty("java.io.tmpdir") + File.separator + "GeoIP"; @Mock - private GeoIPProcessorConfig geoIPProcessorConfig; + private GeoIPDatabaseReader geoIPDatabaseReaderMock; @Mock - private GetGeoData geoData; + private GeoIPDatabaseReader newGeoIPDatabaseReaderMock; @Mock private MaxMindConfig maxMindConfig; - private GeoIPProcessorService geoIPProcessorService; + @Mock + private GeoIpServiceConfig geoIpServiceConfig; + @Mock + private GeoIPDatabaseManager geoIPDatabaseManager; + @Mock + private ReentrantReadWriteLock.ReadLock readLock; @BeforeEach void setUp() { + when(geoIpServiceConfig.getMaxMindConfig()).thenReturn(maxMindConfig); + doNothing().when(geoIPDatabaseManager).initiateDatabaseDownload(); + when(geoIPDatabaseManager.getGeoIPDatabaseReader()).thenReturn(geoIPDatabaseReaderMock, newGeoIPDatabaseReaderMock); + } + GeoIPProcessorService createObjectUnderTest() { + return new GeoIPProcessorService(geoIpServiceConfig, geoIPDatabaseManager, readLock); } @Test - void getGeoDataTest_PATH() throws UnknownHostException, NoSuchFieldException, IllegalAccessException { - // TODO: pass in geoIpServiceConfig object - geoIPProcessorService = new GeoIPProcessorService(null); - - List attributes = List.of(); - InetAddress inetAddress = InetAddress.getByName(IP); - ReflectivelySetField.setField(GeoIPProcessorService.class, - geoIPProcessorService, "geoData", geoData); - when(geoData.getGeoData(any(), any(), anyString())).thenReturn(prepareGeoData()); - Map geoData = geoIPProcessorService.getGeoData(inetAddress, attributes); - assertThat(geoData.get("country_iso_code"), equalTo("US")); - assertThat(geoData.get("continent_name"), equalTo("North America")); + void test_getGeoIPDatabaseReader_should_not_trigger_update_when_refresh_interval_is_high() { + when(geoIPDatabaseManager.getNextUpdateAt()).thenReturn(Instant.now().plus(Duration.ofHours(1))); + + final GeoIPProcessorService objectUnderTest = createObjectUnderTest(); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReaderMock, equalTo(geoIPDatabaseReader)); + verify(geoIPDatabaseManager).getGeoIPDatabaseReader(); + verifyNoMoreInteractions(geoIPDatabaseManager); } @Test - void getGeoDataTest_URL() throws UnknownHostException, NoSuchFieldException, IllegalAccessException { - // TODO: pass in geoIpServiceConfig object - geoIPProcessorService = new GeoIPProcessorService(null); - - List attributes = List.of(); - InetAddress inetAddress = InetAddress.getByName(IP); - ReflectivelySetField.setField(GeoIPProcessorService.class, - geoIPProcessorService, "geoData", geoData); - when(geoData.getGeoData(any(), any(), anyString())).thenReturn(prepareGeoData()); - Map geoData = geoIPProcessorService.getGeoData(inetAddress, attributes); - assertThat(geoData.get("country_iso_code"), equalTo("US")); - assertThat(geoData.get("continent_name"), equalTo("North America")); + void test_getGeoIPDatabaseReader_should_trigger_update_when_refresh_interval_is_met() throws InterruptedException { + when(geoIPDatabaseManager.getNextUpdateAt()).thenReturn(Instant.now().plus(Duration.ofNanos(1))); + doNothing().when(geoIPDatabaseManager).updateDatabaseReader(); + + final GeoIPProcessorService objectUnderTest = createObjectUnderTest(); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReaderMock, equalTo(geoIPDatabaseReader)); + + // Wait for update to be called by ExecutorService + Thread.sleep(1000); + verify(geoIPDatabaseManager).updateDatabaseReader(); + verify(geoIPDatabaseManager).getGeoIPDatabaseReader(); + verify(geoIPDatabaseManager).getNextUpdateAt(); + verify(geoIPDatabaseManager).setNextUpdateAt(any()); + verifyNoMoreInteractions(geoIPDatabaseManager); + + final GeoIPDatabaseReader newGeoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(newGeoIPDatabaseReaderMock, equalTo(newGeoIPDatabaseReader)); } - private Map prepareGeoData() { - Map geoDataMap = new HashMap<>(); - geoDataMap.put("country_iso_code", "US"); - geoDataMap.put("continent_name", "North America"); - geoDataMap.put("timezone", "America/Chicago"); - geoDataMap.put("country_name", "United States"); - return geoDataMap; + @Test + void test_shutdown() throws InterruptedException { + when(geoIPDatabaseManager.getNextUpdateAt()).thenReturn(Instant.now().plus(Duration.ofNanos(1))); + doNothing().when(geoIPDatabaseManager).updateDatabaseReader(); + + final GeoIPProcessorService objectUnderTest = createObjectUnderTest(); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReaderMock, equalTo(geoIPDatabaseReader)); + + // Wait for update to be called by ExecutorService + Thread.sleep(1000); + verify(geoIPDatabaseManager).updateDatabaseReader(); + verify(geoIPDatabaseManager).getGeoIPDatabaseReader(); + verify(geoIPDatabaseManager).getNextUpdateAt(); + verify(geoIPDatabaseManager).setNextUpdateAt(any()); + verifyNoMoreInteractions(geoIPDatabaseManager); + + final GeoIPDatabaseReader newGeoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(newGeoIPDatabaseReaderMock, equalTo(newGeoIPDatabaseReader)); + + assertDoesNotThrow(objectUnderTest::shutdown); } -} \ No newline at end of file +} diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigExtensionTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigExtensionTest.java index 91c48322a2..11a6cefccc 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigExtensionTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpConfigExtensionTest.java @@ -5,6 +5,7 @@ package org.opensearch.dataprepper.plugins.processor.extension; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.ArgumentCaptor; @@ -19,17 +20,25 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.mockito.Mockito.mockConstruction; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; @ExtendWith(MockitoExtension.class) class GeoIpConfigExtensionTest { @Mock private ExtensionPoints extensionPoints; - @Mock private GeoIpServiceConfig geoIpServiceConfig; - @Mock - private ExtensionProvider.Context context; + private MaxMindConfig maxMindConfig; + @Mock + private MaxMindDatabaseConfig maxMindDatabaseConfig; + + + @BeforeEach + void setUp() { + when(geoIpServiceConfig.getMaxMindConfig()).thenReturn(maxMindConfig); + when(maxMindConfig.getMaxMindDatabaseConfig()).thenReturn(maxMindDatabaseConfig); + } private GeoIpConfigExtension createObjectUnderTest() { return new GeoIpConfigExtension(geoIpServiceConfig); @@ -48,17 +57,6 @@ void apply_should_addExtensionProvider() { assertThat(actualExtensionProvider, instanceOf(GeoIpConfigProvider.class)); } - @Test - void extension_should_create_supplier_with_default_config_if_not_configured() { - try (final MockedConstruction mockedConstruction = - mockConstruction(GeoIpServiceConfig.class)) { - final GeoIpConfigExtension geoIpConfigExtension = new GeoIpConfigExtension(null); - - assertThat(mockedConstruction.constructed().size(), equalTo(1)); - assertThat(geoIpConfigExtension, instanceOf(GeoIpConfigExtension.class)); - } - } - @Test void extension_should_create_supplier_with_provided_config() { try (final MockedConstruction mockedConstruction = diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpServiceConfigTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpServiceConfigTest.java index 6376c0a935..aaedb3158c 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpServiceConfigTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/GeoIpServiceConfigTest.java @@ -13,6 +13,7 @@ import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.parser.DataPrepperDurationDeserializer; import org.opensearch.dataprepper.parser.model.DataPrepperConfiguration; +import software.amazon.awssdk.regions.Region; import java.io.File; import java.io.IOException; @@ -49,8 +50,11 @@ void testGeoIpServiceConfig() throws IOException { assertThat(maxMindConfig, notNullValue()); assertThat(maxMindConfig.getDatabaseRefreshInterval(), equalTo(Duration.ofDays(10))); - assertThat(maxMindConfig.getDatabasePaths().size(), equalTo(2)); + assertThat(maxMindConfig.getMaxMindDatabaseConfig().getDatabasePaths().size(), equalTo(2)); assertThat(maxMindConfig.getCacheSize(), equalTo(2048)); + assertThat(maxMindConfig.getDatabaseDestination(), equalTo("/tst/resources/geoip")); + assertThat(maxMindConfig.getAwsAuthenticationOptionsConfig(), notNullValue()); + assertThat(maxMindConfig.getAwsAuthenticationOptionsConfig().getAwsRegion(), equalTo(Region.of("us-east-1"))); } private GeoIpServiceConfig makeConfig(final String filePath) throws IOException { diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindConfigTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindConfigTest.java index 6b22d14bf4..c8b5b190a2 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindConfigTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindConfigTest.java @@ -7,16 +7,29 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.dataprepper.test.helper.ReflectivelySetField; +import java.io.File; +import java.net.URISyntaxException; import java.time.Duration; -import java.util.List; +import java.util.Map; import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.instanceOf; import static org.hamcrest.MatcherAssert.assertThat; +import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.plugins.processor.extension.MaxMindConfig.DEFAULT_DATABASE_DESTINATION; +@ExtendWith(MockitoExtension.class) class MaxMindConfigTest { private MaxMindConfig maxMindConfig; + @Mock + private MaxMindDatabaseConfig maxMindDatabaseConfig; @BeforeEach void setup() { @@ -25,25 +38,60 @@ void setup() { @Test void testDefaultConfig() { - assertThat(maxMindConfig.getDatabasePaths().size(), equalTo(0)); assertThat(maxMindConfig.getDatabaseRefreshInterval(), equalTo(Duration.ofDays(7))); assertThat(maxMindConfig.getCacheSize(), equalTo(4096)); assertThat(maxMindConfig.getAwsAuthenticationOptionsConfig(), equalTo(null)); + assertThat(maxMindConfig.getDatabaseDestination(), equalTo(DEFAULT_DATABASE_DESTINATION + File.separator + "geoip")); + assertThat(maxMindConfig.getMaxMindDatabaseConfig(), instanceOf(MaxMindDatabaseConfig.class)); } @Test void testCustomConfig() throws NoSuchFieldException, IllegalAccessException { ReflectivelySetField.setField(MaxMindConfig.class, maxMindConfig, "databaseRefreshInterval", Duration.ofDays(10)); ReflectivelySetField.setField(MaxMindConfig.class, maxMindConfig, "cacheSize", 2048); - ReflectivelySetField.setField(MaxMindConfig.class, maxMindConfig, "databasePaths", List.of("path1", "path2", "path3")); + ReflectivelySetField.setField(MaxMindConfig.class, maxMindConfig, "maxMindDatabaseConfig", maxMindDatabaseConfig); + ReflectivelySetField.setField(MaxMindConfig.class, maxMindConfig, "databaseDestination", "/data"); final AwsAuthenticationOptionsConfig awsAuthenticationOptionsConfig = new AwsAuthenticationOptionsConfig(); ReflectivelySetField.setField(MaxMindConfig.class, maxMindConfig, "awsAuthenticationOptionsConfig", awsAuthenticationOptionsConfig); assertThat(maxMindConfig.getDatabaseRefreshInterval(), equalTo(Duration.ofDays(10))); assertThat(maxMindConfig.getCacheSize(), equalTo(2048)); - assertThat(maxMindConfig.getDatabasePaths().size(), equalTo(3)); + assertThat(maxMindConfig.getDatabaseDestination(), equalTo("/data/geoip")); assertThat(maxMindConfig.getAwsAuthenticationOptionsConfig(), equalTo(awsAuthenticationOptionsConfig)); + assertThat(maxMindConfig.getMaxMindDatabaseConfig(), equalTo(maxMindDatabaseConfig)); } + @ParameterizedTest + @CsvSource({ + "https://download.maxmind.com/, false, true", + "http://download.maxmind.com/, false, false", + "https://download.maxmind.com/, true, true", + "http://download.maxmind.com/, true, true"}) + void testSecureEndpoint(final String databasePath, final boolean insecure, final boolean result) + throws NoSuchFieldException, IllegalAccessException, URISyntaxException { + when(maxMindDatabaseConfig.getDatabasePaths()).thenReturn(Map.of("name", databasePath)); + ReflectivelySetField.setField(MaxMindConfig.class, maxMindConfig, "maxMindDatabaseConfig", maxMindDatabaseConfig); + ReflectivelySetField.setField(MaxMindConfig.class, maxMindConfig, "insecure", insecure); + + assertThat(maxMindConfig.getMaxMindDatabaseConfig().getDatabasePaths().size(), equalTo(1)); + assertThat(maxMindConfig.isHttpsEndpointOrInsecure(), equalTo(result)); + } + + @ParameterizedTest + @CsvSource({ + "s3://geoip/data, false, false", + "s3://geoip/data, true, true"}) + void testValidPaths(final String databasePath, final boolean awsConfig, final boolean result) + throws NoSuchFieldException, IllegalAccessException { + when(maxMindDatabaseConfig.getDatabasePaths()).thenReturn(Map.of("name", databasePath)); + ReflectivelySetField.setField(MaxMindConfig.class, maxMindConfig, "maxMindDatabaseConfig", maxMindDatabaseConfig); + if (awsConfig) { + final AwsAuthenticationOptionsConfig awsAuthenticationOptionsConfig = new AwsAuthenticationOptionsConfig(); + ReflectivelySetField.setField(MaxMindConfig.class, maxMindConfig, "awsAuthenticationOptionsConfig", awsAuthenticationOptionsConfig); + } + + assertThat(maxMindConfig.getMaxMindDatabaseConfig().getDatabasePaths().size(), equalTo(1)); + assertThat(maxMindConfig.isAwsAuthenticationOptionsValid(), equalTo(result)); + } } \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindDatabaseConfigTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindDatabaseConfigTest.java new file mode 100644 index 0000000000..1e00db0bbf --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/MaxMindDatabaseConfigTest.java @@ -0,0 +1,97 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.opensearch.dataprepper.test.helper.ReflectivelySetField; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig.DEFAULT_ASN_ENDPOINT; +import static org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig.DEFAULT_CITY_ENDPOINT; +import static org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig.DEFAULT_COUNTRY_ENDPOINT; +import static org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig.GEOIP2_ENTERPRISE; +import static org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig.GEOLITE2_ASN; +import static org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig.GEOLITE2_CITY; +import static org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig.GEOLITE2_COUNTRY; + +class MaxMindDatabaseConfigTest { + private MaxMindDatabaseConfig maxMindDatabaseConfig; + @BeforeEach + void setup() { + maxMindDatabaseConfig = new MaxMindDatabaseConfig(); + } + + @Test + void test_default_values() { + assertThat(maxMindDatabaseConfig.getDatabasePaths().get(GEOLITE2_CITY), equalTo(DEFAULT_CITY_ENDPOINT)); + assertThat(maxMindDatabaseConfig.getDatabasePaths().get(GEOLITE2_COUNTRY), equalTo(DEFAULT_COUNTRY_ENDPOINT)); + assertThat(maxMindDatabaseConfig.getDatabasePaths().get(GEOLITE2_ASN), equalTo(DEFAULT_ASN_ENDPOINT)); + assertThat(maxMindDatabaseConfig.getDatabasePaths().get(GEOIP2_ENTERPRISE), equalTo(null)); + assertThat(maxMindDatabaseConfig.getDatabasePaths().size(), equalTo(3)); + } + + @Test + void test_getDatabasePaths_should_not_use_default_endpoints_if_enterprise_database_is_configured() throws NoSuchFieldException, IllegalAccessException { + ReflectivelySetField.setField(MaxMindDatabaseConfig.class, maxMindDatabaseConfig, "enterpriseDatabase", "enterprise_database_path"); + + assertThat(maxMindDatabaseConfig.getDatabasePaths().size(), equalTo(1)); + assertThat(maxMindDatabaseConfig.getDatabasePaths().get(GEOLITE2_CITY), equalTo(null)); + assertThat(maxMindDatabaseConfig.getDatabasePaths().get(GEOLITE2_COUNTRY), equalTo(null)); + assertThat(maxMindDatabaseConfig.getDatabasePaths().get(GEOLITE2_ASN), equalTo(null)); + assertThat(maxMindDatabaseConfig.getDatabasePaths().get(GEOIP2_ENTERPRISE), equalTo("enterprise_database_path")); + } + + @Test + void test_getDatabasePaths_should_use_only_configured_endpoints() throws NoSuchFieldException, IllegalAccessException { + ReflectivelySetField.setField(MaxMindDatabaseConfig.class, maxMindDatabaseConfig, "cityDatabase", "city_database_path"); + ReflectivelySetField.setField(MaxMindDatabaseConfig.class, maxMindDatabaseConfig, "countryDatabase", "country_database_path"); + + assertThat(maxMindDatabaseConfig.getDatabasePaths().size(), equalTo(2)); + assertThat(maxMindDatabaseConfig.getDatabasePaths().get(GEOLITE2_CITY), equalTo("city_database_path")); + assertThat(maxMindDatabaseConfig.getDatabasePaths().get(GEOLITE2_COUNTRY), equalTo("country_database_path")); + assertThat(maxMindDatabaseConfig.getDatabasePaths().get(GEOLITE2_ASN), equalTo(null)); + assertThat(maxMindDatabaseConfig.getDatabasePaths().get(GEOIP2_ENTERPRISE), equalTo(null)); + } + + @Test + void test_custom_endpoints() throws NoSuchFieldException, IllegalAccessException { + ReflectivelySetField.setField(MaxMindDatabaseConfig.class, maxMindDatabaseConfig, "cityDatabase", "city_database_path"); + ReflectivelySetField.setField(MaxMindDatabaseConfig.class, maxMindDatabaseConfig, "countryDatabase", "country_database_path"); + ReflectivelySetField.setField(MaxMindDatabaseConfig.class, maxMindDatabaseConfig, "asnDatabase", "asn_database_path"); + ReflectivelySetField.setField(MaxMindDatabaseConfig.class, maxMindDatabaseConfig, "enterpriseDatabase", "enterprise_database_path"); + + assertThat(maxMindDatabaseConfig.getDatabasePaths().size(), equalTo(4)); + assertThat(maxMindDatabaseConfig.getDatabasePaths().get(GEOLITE2_CITY), equalTo("city_database_path")); + assertThat(maxMindDatabaseConfig.getDatabasePaths().get(GEOLITE2_COUNTRY), equalTo("country_database_path")); + assertThat(maxMindDatabaseConfig.getDatabasePaths().get(GEOLITE2_ASN), equalTo("asn_database_path")); + assertThat(maxMindDatabaseConfig.getDatabasePaths().get(GEOIP2_ENTERPRISE), equalTo("enterprise_database_path")); + } + + @Test + void test_invalid_config_when_geolite2_and_enterprise_databases_are_configured() throws NoSuchFieldException, IllegalAccessException { + ReflectivelySetField.setField(MaxMindDatabaseConfig.class, maxMindDatabaseConfig, "asnDatabase", "asn_database_path"); + ReflectivelySetField.setField(MaxMindDatabaseConfig.class, maxMindDatabaseConfig, "enterpriseDatabase", "enterprise_database_path"); + + assertThat(maxMindDatabaseConfig.isDatabasesValid(), equalTo(false)); + } + + @ParameterizedTest + @CsvSource({"https://download.maxmind.com/, true", + "https://example.com/, false", + "s3://geoip/data, true", + "https://geoip.maps.opensearch.org/v1/mmdb/geolite2-city/manifest.json, true", + "https://geoip.maps.opensearch.org/v1/mmdb/geolite2/input.json, false"}) + void test_isPathsValid(final String path, final boolean result) throws NoSuchFieldException, IllegalAccessException { + ReflectivelySetField.setField(MaxMindDatabaseConfig.class, maxMindDatabaseConfig, "asnDatabase", path); + + assertThat(maxMindDatabaseConfig.isPathsValid(), equalTo(result)); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSourceTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSourceTest.java deleted file mode 100644 index 3b7ed49d87..0000000000 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DBSourceTest.java +++ /dev/null @@ -1,42 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; - -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.MockedStatic; -import org.mockito.junit.jupiter.MockitoExtension; - -import java.io.File; - -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.mockito.Mockito.mockStatic; - -@ExtendWith(MockitoExtension.class) -class DBSourceTest { - - private final String outputFilePath = System.getProperty("java.io.tmpdir") + "GeoTest"; - - @Test - void createFolderIfNotExistTest() { - try (MockedStatic mockedStatic = mockStatic(DBSource.class)) { - mockedStatic.when(() -> DBSource.createFolderIfNotExist(outputFilePath)).thenReturn(new File(outputFilePath)); - File actualFile = new File(outputFilePath); - assertEquals(actualFile, DBSource.createFolderIfNotExist(outputFilePath)); - } - } - - @Test - void deleteDirectoryTest() { - DBSource.createFolderIfNotExist(outputFilePath); - DBSource.createFolderIfNotExist(outputFilePath + File.separator + "GeoIPz"); - DBSource.createFolderIfNotExist(outputFilePath + File.separator + "GeoIPx"); - assertDoesNotThrow(() -> { - DBSource.deleteDirectory(new File(outputFilePath)); - }); - } -} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderBuilderTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderBuilderTest.java new file mode 100644 index 0000000000..1b435fb1f2 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderBuilderTest.java @@ -0,0 +1,71 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; + +import com.maxmind.geoip2.DatabaseReader; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.mockito.junit.jupiter.MockitoExtension; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Path; + +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +@ExtendWith(MockitoExtension.class) +class DatabaseReaderBuilderTest { + public static final String GEOLITE2_TEST_MMDB_FILES = "./build/resources/test/mmdb-files/geo-lite2"; + public static final String GEOIP2_TEST_MMDB_FILES = "./build/resources/test/mmdb-files/geo-ip2"; + + @ParameterizedTest + @ValueSource(strings = {"geolite2-asn", "geolite2-city", "geolite2-country"}) + void createLoaderTest_for_geolite2_databases(final String databaseName) throws IOException { + String databaseToUse = null; + final File directory = new File(GEOLITE2_TEST_MMDB_FILES); + final String[] list = directory.list(); + for (String fileName: list) { + final String lowerCaseFileName = fileName.toLowerCase(); + if (fileName.endsWith(".mmdb") + && lowerCaseFileName.contains(databaseName)) { + databaseToUse = fileName; + } + } + + final Path path = Path.of(GEOLITE2_TEST_MMDB_FILES + File.separator + databaseToUse); + + final DatabaseReaderBuilder databaseReaderBuilder = new DatabaseReaderBuilder(); + + final DatabaseReader databaseReader = databaseReaderBuilder.buildReader(path, 4096); + assertNotNull(databaseReader); + assertTrue(databaseToUse.toLowerCase().contains(databaseReader.getMetadata().getDatabaseType().toLowerCase())); + } + + @ParameterizedTest + @ValueSource(strings = {"geoip2-enterprise"}) + void createLoaderTest_for_geoip2_databases(final String databaseName) throws IOException { + String databaseToUse = null; + final File directory = new File(GEOIP2_TEST_MMDB_FILES); + final String[] list = directory.list(); + for (String fileName: list) { + final String lowerCaseFileName = fileName.toLowerCase(); + if (fileName.endsWith(".mmdb") + && lowerCaseFileName.contains(databaseName)) { + databaseToUse = fileName; + } + } + + final Path path = Path.of(GEOIP2_TEST_MMDB_FILES + File.separator + databaseToUse); + + final DatabaseReaderBuilder databaseReaderBuilder = new DatabaseReaderBuilder(); + + final DatabaseReader databaseReader = databaseReaderBuilder.buildReader(path, 4096); + assertNotNull(databaseReader); + assertTrue(databaseToUse.toLowerCase().contains(databaseReader.getMetadata().getDatabaseType().toLowerCase())); + } +} diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderCreateTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderCreateTest.java deleted file mode 100644 index 621498e8e2..0000000000 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/DatabaseReaderCreateTest.java +++ /dev/null @@ -1,43 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; - -import com.maxmind.geoip2.DatabaseReader; -import org.apache.commons.io.FileUtils; -import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; - -import java.io.BufferedInputStream; -import java.io.File; -import java.io.IOException; -import java.net.URL; -import java.nio.file.Path; - -import static org.mockito.Mockito.when; - -@ExtendWith(MockitoExtension.class) -class DatabaseReaderCreateTest { - @Mock - private Path path; - - @Test - void createLoaderTest() throws IOException { - final String testFileURL = "https://github.com/maxmind/MaxMind-DB/raw/main/test-data/GeoLite2-City-Test.mmdb"; - final File file = File.createTempFile( "GeoIP2-City-Test", ".mmdb"); - - final BufferedInputStream in = new BufferedInputStream(new URL(testFileURL).openStream()); - FileUtils.copyInputStreamToFile(in, file); - when(path.toFile()).thenReturn(file); - - DatabaseReader databaseReader = DatabaseReaderCreate.createLoader(path, 4096); - Assertions.assertNotNull(databaseReader); - in.close(); - file.deleteOnExit(); - } -} diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoDataFactoryTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoDataFactoryTest.java deleted file mode 100644 index 83be77f774..0000000000 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoDataFactoryTest.java +++ /dev/null @@ -1,51 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; - -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.GetGeoData; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.GetGeoIP2Data; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.GetGeoLite2Data; -import org.opensearch.dataprepper.plugins.processor.extension.MaxMindConfig; -import org.opensearch.dataprepper.plugins.processor.utils.LicenseTypeCheck; - -import static org.junit.jupiter.api.Assertions.assertInstanceOf; -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.Mockito.when; - -@ExtendWith(MockitoExtension.class) -class GeoDataFactoryTest { - - @Mock - private MaxMindConfig maxMindConfig; - - @Mock - private LicenseTypeCheck licenseTypeCheck; - - @Test - void testCreateWithFreeLicense() { - when(licenseTypeCheck.isGeoLite2OrEnterpriseLicense(anyString())).thenReturn(LicenseTypeOptions.FREE); - final GeoDataFactory geoDataFactory = new GeoDataFactory(maxMindConfig, licenseTypeCheck); - final String databasePath = "testPath"; - - final GetGeoData getGeoData = geoDataFactory.create(databasePath); - assertInstanceOf(GetGeoLite2Data.class, getGeoData); - } - - @Test - void testCreateWithEnterpriseLicense() { - when(licenseTypeCheck.isGeoLite2OrEnterpriseLicense(anyString())).thenReturn(LicenseTypeOptions.ENTERPRISE); - final GeoDataFactory geoDataFactory = new GeoDataFactory(maxMindConfig, licenseTypeCheck); - final String databasePath = "testPath"; - - final GetGeoData getGeoData = geoDataFactory.create(databasePath); - assertInstanceOf(GetGeoIP2Data.class, getGeoData); - } - -} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPDatabaseManagerTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPDatabaseManagerTest.java new file mode 100644 index 0000000000..c5c9ea4510 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPDatabaseManagerTest.java @@ -0,0 +1,349 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.MockedConstruction; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.AutoCountingDatabaseReader; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIP2DatabaseReader; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoIPDatabaseReader; +import org.opensearch.dataprepper.plugins.processor.databaseenrich.GeoLite2DatabaseReader; +import org.opensearch.dataprepper.plugins.processor.exception.DownloadFailedException; +import org.opensearch.dataprepper.plugins.processor.extension.MaxMindConfig; +import org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig; +import org.opensearch.dataprepper.plugins.processor.utils.LicenseTypeCheck; +import org.opensearch.dataprepper.test.helper.ReflectivelySetField; + +import java.time.Duration; +import java.time.Instant; +import java.util.HashMap; +import java.util.concurrent.locks.ReentrantReadWriteLock; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mockConstruction; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +class GeoIPDatabaseManagerTest { + private static final String S3_URI = "s3://geoip/data/GeoLite2-Country-Test.mmdb"; + private static final String HTTPS_URL = "https://download.maxmind.com/app/geoip_download?editionid=GeoLite2-ASN&suffix=tar.gz"; + private static final String PATH = "./build/resources/test/mmdb-files/geo-lite2/GeoLite2-Country-Test.mmdb"; + private static final String CDN_ENDPOINT = "https://devo.geoip.maps.opensearch.org/v1/mmdb/geolite2/manifest.json"; + + @Mock + private MaxMindConfig maxMindConfig; + + @Mock + private LicenseTypeCheck licenseTypeCheck; + + @Mock + private DatabaseReaderBuilder databaseReaderBuilder; + + @Mock + private ReentrantReadWriteLock.WriteLock writeLock; + + @Mock + private GeoIPFileManager databaseFileManager; + + @Mock + private MaxMindDatabaseConfig maxMindDatabaseConfig; + + @BeforeEach + void setUp() { + lenient().when(licenseTypeCheck.isGeoLite2OrEnterpriseLicense(any())).thenReturn(LicenseTypeOptions.FREE); + lenient().when(maxMindConfig.getMaxMindDatabaseConfig()).thenReturn(maxMindDatabaseConfig); + } + + private GeoIPDatabaseManager createObjectUnderTest() { + return new GeoIPDatabaseManager(maxMindConfig, licenseTypeCheck, databaseReaderBuilder, databaseFileManager, writeLock); + } + + @Test + void test_initiateDatabaseDownload_with_geolite2_file_path_should_use_local_download_service_and_geolite2_reader() throws Exception { + try (final MockedConstruction localDBDownloadServiceMockedConstruction = mockConstruction(LocalDBDownloadService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload()); + final MockedConstruction geoLite2DatabaseReaderMockedConstruction = mockConstruction(GeoLite2DatabaseReader.class); + final MockedConstruction autoCountingDatabaseReaderMockedConstruction = mockConstruction(AutoCountingDatabaseReader.class) + ) { + + final HashMap databases = new HashMap<>(); + databases.put("geolite2_country", PATH); + when(maxMindDatabaseConfig.getDatabasePaths()).thenReturn(databases); + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + objectUnderTest.initiateDatabaseDownload(); + + assertThat(localDBDownloadServiceMockedConstruction.constructed().size(), equalTo(1)); + verify(localDBDownloadServiceMockedConstruction.constructed().get(0)).initiateDownload(); + assertThat(geoLite2DatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + assertThat(autoCountingDatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(autoCountingDatabaseReaderMockedConstruction.constructed().get(0))); + } + } + + @Test + void test_initiateDatabaseDownload_with_geoip2_file_path_should_use_local_download_service_and_geoip2_reader() throws Exception { + try (final MockedConstruction localDBDownloadServiceMockedConstruction = mockConstruction(LocalDBDownloadService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload()); + final MockedConstruction geoIP2DatabaseReaderMockedConstruction = mockConstruction(GeoIP2DatabaseReader.class); + final MockedConstruction autoCountingDatabaseReaderMockedConstruction = mockConstruction(AutoCountingDatabaseReader.class) + ) { + when(licenseTypeCheck.isGeoLite2OrEnterpriseLicense(any())).thenReturn(LicenseTypeOptions.ENTERPRISE); + final HashMap databases = new HashMap<>(); + databases.put("geolite2_country", PATH); + when(maxMindDatabaseConfig.getDatabasePaths()).thenReturn(databases); + + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + objectUnderTest.initiateDatabaseDownload(); + + assertThat(localDBDownloadServiceMockedConstruction.constructed().size(), equalTo(1)); + verify(localDBDownloadServiceMockedConstruction.constructed().get(0)).initiateDownload(); + assertThat(geoIP2DatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + assertThat(autoCountingDatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(autoCountingDatabaseReaderMockedConstruction.constructed().get(0))); + } + } + + @Test + void test_initiateDatabaseDownload_with_geolite2_s3_uri_should_use_s3_download_service_and_geolite2_reader() { + try (final MockedConstruction s3DBServiceMockedConstruction = mockConstruction(S3DBService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload()); + final MockedConstruction geoLite2DatabaseReaderMockedConstruction = mockConstruction(GeoLite2DatabaseReader.class); + final MockedConstruction autoCountingDatabaseReaderMockedConstruction = mockConstruction(AutoCountingDatabaseReader.class) + ) { + + final HashMap databases = new HashMap<>(); + databases.put("geolite2_country", S3_URI); + when(maxMindDatabaseConfig.getDatabasePaths()).thenReturn(databases); + when(maxMindConfig.getMaxMindDatabaseConfig()).thenReturn(maxMindDatabaseConfig); + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + objectUnderTest.initiateDatabaseDownload(); + + assertThat(s3DBServiceMockedConstruction.constructed().size(), equalTo(1)); + verify(s3DBServiceMockedConstruction.constructed().get(0)).initiateDownload(); + assertThat(geoLite2DatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + assertThat(autoCountingDatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(autoCountingDatabaseReaderMockedConstruction.constructed().get(0))); + } + } + + @Test + void test_initiateDatabaseDownload_with_geoip2_s3_uri_should_use_s3_download_service_and_geoip2_reader() { + try (final MockedConstruction s3DBServiceMockedConstruction = mockConstruction(S3DBService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload()); + final MockedConstruction geoIP2DatabaseReaderMockedConstruction = mockConstruction(GeoIP2DatabaseReader.class); + final MockedConstruction autoCountingDatabaseReaderMockedConstruction = mockConstruction(AutoCountingDatabaseReader.class) + ) { + when(licenseTypeCheck.isGeoLite2OrEnterpriseLicense(any())).thenReturn(LicenseTypeOptions.ENTERPRISE); + final HashMap databases = new HashMap<>(); + databases.put("geolite2_country", S3_URI); + when(maxMindDatabaseConfig.getDatabasePaths()).thenReturn(databases); + when(maxMindConfig.getMaxMindDatabaseConfig()).thenReturn(maxMindDatabaseConfig); + + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + objectUnderTest.initiateDatabaseDownload(); + + assertThat(s3DBServiceMockedConstruction.constructed().size(), equalTo(1)); + verify(s3DBServiceMockedConstruction.constructed().get(0)).initiateDownload(); + assertThat(geoIP2DatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + assertThat(autoCountingDatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(autoCountingDatabaseReaderMockedConstruction.constructed().get(0))); + } + } + + @Test + void test_initiateDatabaseDownload_with_geolite2_url_should_use_http_download_service_and_geolite2_reader() { + try (final MockedConstruction httpDBDownloadServiceMockedConstruction = mockConstruction(HttpDBDownloadService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload()); + final MockedConstruction geoLite2DatabaseReaderMockedConstruction = mockConstruction(GeoLite2DatabaseReader.class); + final MockedConstruction autoCountingDatabaseReaderMockedConstruction = mockConstruction(AutoCountingDatabaseReader.class) + ) { + + final HashMap databases = new HashMap<>(); + databases.put("geolite2_country", HTTPS_URL); + when(maxMindDatabaseConfig.getDatabasePaths()).thenReturn(databases); + when(maxMindConfig.getMaxMindDatabaseConfig()).thenReturn(maxMindDatabaseConfig); + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + objectUnderTest.initiateDatabaseDownload(); + + assertThat(httpDBDownloadServiceMockedConstruction.constructed().size(), equalTo(1)); + verify(httpDBDownloadServiceMockedConstruction.constructed().get(0)).initiateDownload(); + assertThat(geoLite2DatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + assertThat(autoCountingDatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(autoCountingDatabaseReaderMockedConstruction.constructed().get(0))); + } + } + + @Test + void test_initiateDatabaseDownload_with_geoip2_url_should_use_http_download_service_and_geoip2_reader() { + try (final MockedConstruction httpDBDownloadServiceMockedConstruction = mockConstruction(HttpDBDownloadService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload()); + final MockedConstruction geoIP2DatabaseReaderMockedConstruction = mockConstruction(GeoIP2DatabaseReader.class); + final MockedConstruction autoCountingDatabaseReaderMockedConstruction = mockConstruction(AutoCountingDatabaseReader.class) + ) { + when(licenseTypeCheck.isGeoLite2OrEnterpriseLicense(any())).thenReturn(LicenseTypeOptions.ENTERPRISE); + final HashMap databases = new HashMap<>(); + databases.put("geolite2_country", HTTPS_URL); + when(maxMindDatabaseConfig.getDatabasePaths()).thenReturn(databases); + when(maxMindConfig.getMaxMindDatabaseConfig()).thenReturn(maxMindDatabaseConfig); + + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + objectUnderTest.initiateDatabaseDownload(); + + assertThat(httpDBDownloadServiceMockedConstruction.constructed().size(), equalTo(1)); + verify(httpDBDownloadServiceMockedConstruction.constructed().get(0)).initiateDownload(); + assertThat(geoIP2DatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + assertThat(autoCountingDatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(autoCountingDatabaseReaderMockedConstruction.constructed().get(0))); + } + } + + @Test + void test_initiateDatabaseDownload_with_geolite2_cdn_should_use_cdn_download_service_and_geolite2_reader() { + try (final MockedConstruction cdnDownloadServiceMockedConstruction = mockConstruction(ManifestDownloadService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload()); + final MockedConstruction geoLite2DatabaseReaderMockedConstruction = mockConstruction(GeoLite2DatabaseReader.class); + final MockedConstruction autoCountingDatabaseReaderMockedConstruction = mockConstruction(AutoCountingDatabaseReader.class) + ) { + final HashMap databases = new HashMap<>(); + databases.put("geolite2_country", CDN_ENDPOINT); + when(maxMindDatabaseConfig.getDatabasePaths()).thenReturn(databases); + when(maxMindConfig.getMaxMindDatabaseConfig()).thenReturn(maxMindDatabaseConfig); + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + objectUnderTest.initiateDatabaseDownload(); + + assertThat(cdnDownloadServiceMockedConstruction.constructed().size(), equalTo(1)); + verify(cdnDownloadServiceMockedConstruction.constructed().get(0)).initiateDownload(); + assertThat(geoLite2DatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + assertThat(autoCountingDatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(autoCountingDatabaseReaderMockedConstruction.constructed().get(0))); + } + } + + @Test + void test_updateDatabaseReader_with_geolite2_cdn_should_use_cdn_download_service_and_geolite2_reader_and_get_new_reader() throws Exception { + try (final MockedConstruction cdnDownloadServiceMockedConstruction = mockConstruction(ManifestDownloadService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload()); + final MockedConstruction geoLite2DatabaseReaderMockedConstruction = mockConstruction(GeoLite2DatabaseReader.class); + final MockedConstruction autoCountingDatabaseReaderMockedConstruction = mockConstruction(AutoCountingDatabaseReader.class) + ) { + final HashMap databases = new HashMap<>(); + databases.put("geolite2_country", CDN_ENDPOINT); + when(maxMindDatabaseConfig.getDatabasePaths()).thenReturn(databases); + when(maxMindConfig.getMaxMindDatabaseConfig()).thenReturn(maxMindDatabaseConfig); + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + + objectUnderTest.initiateDatabaseDownload(); + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(autoCountingDatabaseReaderMockedConstruction.constructed().get(0))); + + objectUnderTest.updateDatabaseReader(); + + assertThat(cdnDownloadServiceMockedConstruction.constructed().size(), equalTo(2)); + for (ManifestDownloadService manifestDownloadService : cdnDownloadServiceMockedConstruction.constructed()) { + verify(manifestDownloadService).initiateDownload(); + } + assertThat(geoLite2DatabaseReaderMockedConstruction.constructed().size(), equalTo(2)); + assertThat(autoCountingDatabaseReaderMockedConstruction.constructed().size(), equalTo(2)); + // verify if first instance is closed + verify(autoCountingDatabaseReaderMockedConstruction.constructed().get(0)).close(); + final GeoIPDatabaseReader updatedGeoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(updatedGeoIPDatabaseReader, equalTo(autoCountingDatabaseReaderMockedConstruction.constructed().get(1))); + } + } + + @Test + void test_initiateDatabaseDownload_with_exception_should_update_nextUpdateAt_correctly_with_backoff() { + try (final MockedConstruction cdnDownloadServiceMockedConstruction = mockConstruction(ManifestDownloadService.class, + (mock2, context2)-> doThrow(DownloadFailedException.class).when(mock2).initiateDownload()); + final MockedConstruction geoLite2DatabaseReaderMockedConstruction = mockConstruction(GeoLite2DatabaseReader.class); + final MockedConstruction autoCountingDatabaseReaderMockedConstruction = mockConstruction(AutoCountingDatabaseReader.class) + ) { + final HashMap databases = new HashMap<>(); + databases.put("geolite2_country", CDN_ENDPOINT); + when(maxMindDatabaseConfig.getDatabasePaths()).thenReturn(databases); + when(maxMindConfig.getMaxMindDatabaseConfig()).thenReturn(maxMindDatabaseConfig); + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + + assertThrows(DownloadFailedException.class, objectUnderTest::initiateDatabaseDownload); + + assertThat(cdnDownloadServiceMockedConstruction.constructed().size(), equalTo(1)); + verify(cdnDownloadServiceMockedConstruction.constructed().get(0)).initiateDownload(); + assertThat(geoLite2DatabaseReaderMockedConstruction.constructed().size(), equalTo(0)); + assertThat(autoCountingDatabaseReaderMockedConstruction.constructed().size(), equalTo(0)); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(null)); + + assertTrue(Instant.now().plus(Duration.ofMinutes(5)).isAfter(objectUnderTest.getNextUpdateAt())); + assertTrue(Instant.now().minus(Duration.ofMinutes(5)).isBefore(objectUnderTest.getNextUpdateAt())); + } + } + + @Test + void test_initiateDatabaseDownload_without_exception_should_update_databases_as_configured() { + try (final MockedConstruction cdnDownloadServiceMockedConstruction = mockConstruction(ManifestDownloadService.class, + (mock2, context2)-> doNothing().when(mock2).initiateDownload()); + final MockedConstruction geoLite2DatabaseReaderMockedConstruction = mockConstruction(GeoLite2DatabaseReader.class); + final MockedConstruction autoCountingDatabaseReaderMockedConstruction = mockConstruction(AutoCountingDatabaseReader.class) + ) { + final HashMap databases = new HashMap<>(); + databases.put("geolite2_country", CDN_ENDPOINT); + when(maxMindDatabaseConfig.getDatabasePaths()).thenReturn(databases); + when(maxMindConfig.getMaxMindDatabaseConfig()).thenReturn(maxMindDatabaseConfig); + when(maxMindConfig.getDatabaseRefreshInterval()).thenReturn(Duration.ofDays(3)); + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + + objectUnderTest.initiateDatabaseDownload(); + + assertThat(cdnDownloadServiceMockedConstruction.constructed().size(), equalTo(1)); + verify(cdnDownloadServiceMockedConstruction.constructed().get(0)).initiateDownload(); + assertThat(geoLite2DatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + assertThat(autoCountingDatabaseReaderMockedConstruction.constructed().size(), equalTo(1)); + + final GeoIPDatabaseReader geoIPDatabaseReader = objectUnderTest.getGeoIPDatabaseReader(); + assertThat(geoIPDatabaseReader, equalTo(autoCountingDatabaseReaderMockedConstruction.constructed().get(0))); + + assertTrue(Instant.now().plus(Duration.ofDays(4)).isAfter(objectUnderTest.getNextUpdateAt())); + assertTrue(Instant.now().minus(Duration.ofDays(2)).isBefore(objectUnderTest.getNextUpdateAt())); + } + } + + @Test + void test_getNextUpdateAt() throws NoSuchFieldException, IllegalAccessException { + final GeoIPDatabaseManager objectUnderTest = createObjectUnderTest(); + + ReflectivelySetField.setField(GeoIPDatabaseManager.class, objectUnderTest, "nextUpdateAt", Instant.now()); + + assertTrue(objectUnderTest.getNextUpdateAt().isBefore(Instant.now())); + assertTrue(objectUnderTest.getNextUpdateAt().plus(Duration.ofMinutes(1)).isAfter(Instant.now())); + } +} diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPFileManagerTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPFileManagerTest.java new file mode 100644 index 0000000000..c7c74acf8d --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/GeoIPFileManagerTest.java @@ -0,0 +1,36 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; + +import org.junit.jupiter.api.Test; + +import java.io.File; + +import static org.junit.jupiter.api.Assertions.assertTrue; + +class GeoIPFileManagerTest { + private final String outputFilePath = "./src/test/resources/geoip/test"; + + @Test + void createFolderIfNotExistTest() { + final GeoIPFileManager geoIPFileManager = new GeoIPFileManager(); + geoIPFileManager.createDirectoryIfNotExist(outputFilePath); + + final File file = new File(outputFilePath); + assertTrue(file.exists()); + } + + @Test + void deleteDirectoryTest() { + final GeoIPFileManager geoIPFileManager = new GeoIPFileManager(); + geoIPFileManager.createDirectoryIfNotExist(outputFilePath); + + final File file = new File(outputFilePath); + assertTrue(file.isDirectory()); + geoIPFileManager.deleteDirectory(file); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/HttpDBDownloadServiceTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/HttpDBDownloadServiceTest.java index 8b2bbf47cd..930c7eb651 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/HttpDBDownloadServiceTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/HttpDBDownloadServiceTest.java @@ -7,9 +7,9 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; - -import java.util.List; +import org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; @@ -18,6 +18,10 @@ class HttpDBDownloadServiceTest { private static final String PREFIX_DIR = "first_database"; private HttpDBDownloadService downloadThroughUrl; + @Mock + private GeoIPFileManager geoIPFileManager; + @Mock + private MaxMindDatabaseConfig maxMindDatabaseConfig; @Test void initiateDownloadTest() { @@ -25,11 +29,11 @@ void initiateDownloadTest() { downloadThroughUrl = createObjectUnderTest(); assertDoesNotThrow(() -> { - downloadThroughUrl.initiateDownload(List.of(databasePath)); + downloadThroughUrl.initiateDownload(); }); } private HttpDBDownloadService createObjectUnderTest() { - return new HttpDBDownloadService(PREFIX_DIR); + return new HttpDBDownloadService(PREFIX_DIR, geoIPFileManager, maxMindDatabaseConfig); } } \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/LocalDBDownloadServiceTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/LocalDBDownloadServiceTest.java index eaee24774c..b6097ab47f 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/LocalDBDownloadServiceTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/LocalDBDownloadServiceTest.java @@ -5,55 +5,80 @@ package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig; import java.io.File; import java.io.FileWriter; import java.io.IOException; -import java.util.List; +import java.util.Map; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.when; @ExtendWith(MockitoExtension.class) class LocalDBDownloadServiceTest { - - private static final String PREFIX_DIR = "geo-lite2"; - String tempFolderPath = System.getProperty("java.io.tmpdir") + File.separator + "GeoIP"; - String srcDir = System.getProperty("java.io.tmpdir") + File.separator + "Maxmind"; + private final String destinationDirectory = "./src/test/resources/dest/"; + private final String sourceDirectory = "./src/test/resources/src/"; private LocalDBDownloadService downloadThroughLocalPath; + @Mock + private MaxMindDatabaseConfig maxMindDatabaseConfig; @Test void initiateDownloadTest() throws IOException { - createFolder(System.getProperty("java.io.tmpdir") + File.separator + "Maxmind"); - generateSampleFiles(srcDir, 5); - downloadThroughLocalPath = createObjectUnderTest(); + generateSampleFiles(); assertDoesNotThrow(() -> { - downloadThroughLocalPath.initiateDownload(List.of(srcDir)); + downloadThroughLocalPath.initiateDownload(); }); + + assertTrue(new File(destinationDirectory + File.separator + "filename.mmdb").exists()); } private LocalDBDownloadService createObjectUnderTest() { - return new LocalDBDownloadService(PREFIX_DIR); + when(maxMindDatabaseConfig.getDatabasePaths()).thenReturn(Map.of("filename", sourceDirectory + File.separator + "SampleFile.mmdb")); + createFolder(destinationDirectory); + return new LocalDBDownloadService(destinationDirectory, maxMindDatabaseConfig); } - private static void createFolder(String folderName) { + private void createFolder(String folderName) { File folder = new File(folderName); if (!folder.exists()) { folder.mkdir(); } } - private static void generateSampleFiles(String folderName, int numFiles) throws IOException { - for (int i = 1; i <= numFiles; i++) { - String fileName = "SampleFile" + i + ".txt"; - String content = "This is sample file " + i; + private void generateSampleFiles() throws IOException { + String fileName = "SampleFile.mmdb"; + String content = "This is sample file"; + + createFolder(sourceDirectory); + new File(sourceDirectory + File.separator + fileName); + try (FileWriter writer = new FileWriter(sourceDirectory + File.separator + fileName)) { + writer.write(content); + } + } + + @AfterEach + void cleanUp() { + deleteDirectory(new File(sourceDirectory)); + deleteDirectory(new File(destinationDirectory)); + } - try (FileWriter writer = new FileWriter(folderName + File.separator + fileName)) { - writer.write(content); + public void deleteDirectory(final File file) { + if (file.exists()) { + for (final File subFile : file.listFiles()) { + if (subFile.isDirectory()) { + deleteDirectory(subFile); + } + subFile.delete(); } + file.delete(); } } } \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/ManifestDownloadServiceTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/ManifestDownloadServiceTest.java new file mode 100644 index 0000000000..4ad90e99e2 --- /dev/null +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/ManifestDownloadServiceTest.java @@ -0,0 +1,56 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.extension.databasedownload; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.plugins.processor.exception.DownloadFailedException; +import org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig; + +import java.io.File; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +class ManifestDownloadServiceTest { + private static final String OUTPUT_DIR = "./src/test/resources/geoip"; + private String path = "https://geoip.maps.opensearch.org/v1/mmdb/geolite2-city/manifest.json"; + private String invalid_path = "https://geoip.maps.opensearch.org/v1/mmdb/geolite2-test/manifest.json"; + + @Mock + private MaxMindDatabaseConfig maxMindDatabaseConfig; + + private ManifestDownloadService createObjectUnderTest() { + return new ManifestDownloadService(OUTPUT_DIR, maxMindDatabaseConfig); + } + + @Test + void test_with_valid_endpoint_should_download_file() { + when(maxMindDatabaseConfig.getDatabasePaths()).thenReturn(Map.of("geolite2-city", path)); + final ManifestDownloadService objectUnderTest = createObjectUnderTest(); + objectUnderTest.initiateDownload(); + + final File file = new File(OUTPUT_DIR + File.separator + "geolite2-city.mmdb"); + assertTrue(file.exists()); + + file.deleteOnExit(); + final File directory = new File(OUTPUT_DIR); + directory.deleteOnExit(); + } + + @Test + void test_with_invalid_endpoint_should_throw_exception() { + when(maxMindDatabaseConfig.getDatabasePaths()).thenReturn(Map.of("geolite2-city", invalid_path)); + final ManifestDownloadService objectUnderTest = createObjectUnderTest(); + assertThrows(DownloadFailedException.class, () -> objectUnderTest.initiateDownload()); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/S3DBServiceTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/S3DBServiceTest.java index 815d2eb624..9fbc28acc0 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/S3DBServiceTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/extension/databasedownload/S3DBServiceTest.java @@ -10,34 +10,37 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.plugins.processor.GeoIPProcessorConfig; -import org.opensearch.dataprepper.plugins.processor.databaseenrich.DownloadFailedException; +import org.opensearch.dataprepper.plugins.processor.exception.DownloadFailedException; +import org.opensearch.dataprepper.plugins.processor.extension.AwsAuthenticationOptionsConfig; +import org.opensearch.dataprepper.plugins.processor.extension.MaxMindDatabaseConfig; -import java.util.List; +import java.util.Map; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.when; @ExtendWith(MockitoExtension.class) class S3DBServiceTest { private static final String S3_URI = "s3://mybucket10012023/GeoLite2/"; - private static final String PREFIX_DIR = "first_database"; - private static final String S3_REGION = "us-east-1"; + private static final String DATABASE_DIR = "first_database"; @Mock - private GeoIPProcessorConfig geoIPProcessorConfig; + private MaxMindDatabaseConfig maxMindDatabaseConfig; + @Mock + private AwsAuthenticationOptionsConfig awsAuthenticationOptionsConfig; @BeforeEach void setUp() { - + when(maxMindDatabaseConfig.getDatabasePaths()).thenReturn(Map.of("database-name", S3_URI)); } @Test void initiateDownloadTest_DownloadFailedException() { S3DBService downloadThroughS3 = createObjectUnderTest(); - assertThrows(DownloadFailedException.class, () -> downloadThroughS3.initiateDownload(List.of(S3_URI))); + assertThrows(DownloadFailedException.class, () -> downloadThroughS3.initiateDownload()); } private S3DBService createObjectUnderTest() { - return new S3DBService(null, PREFIX_DIR); + return new S3DBService(awsAuthenticationOptionsConfig, DATABASE_DIR, maxMindDatabaseConfig); } } \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/loadtype/LoadTypeOptionsTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/loadtype/LoadTypeOptionsTest.java deleted file mode 100644 index 999a970482..0000000000 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/loadtype/LoadTypeOptionsTest.java +++ /dev/null @@ -1,30 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.processor.loadtype; - -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.junit.jupiter.MockitoExtension; - -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.junit.jupiter.api.Assertions.assertNotNull; - -@ExtendWith(MockitoExtension.class) -class LoadTypeOptionsTest { - - @Test - void notNull_test() { - assertNotNull(LoadTypeOptions.INMEMORY); - } - - @Test - void fromOptionValue_test() { - LoadTypeOptions loadTypeOptions = LoadTypeOptions.fromOptionValue("memory_map"); - assertNotNull(loadTypeOptions); - assertThat(loadTypeOptions.toString(), equalTo("INMEMORY")); - } -} diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/DbSourceIdentificationTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/DbSourceIdentificationTest.java index 7be0b00db9..684e1d66b3 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/DbSourceIdentificationTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/DbSourceIdentificationTest.java @@ -20,65 +20,81 @@ class DbSourceIdentificationTest { private static final String S3_URI = "s3://dataprepper/logdata/22833bd46b8e0.mmdb"; - private static final String S3_URL = "https://dataprepper.s3.amazonaws.com/logdata/22833bd46b8e0.json"; - private static final String URL = "https://www.dataprepper.com"; - private static final String PATH = "./src/test/resources/mmdb-file/geo-lite2"; + private static final String URL = "https://download.maxmind.com/"; + private static final String DIRECTORY_PATH = "./build/resources/test/mmdb-files/geo-lite2"; + private static final String FILE_PATH = "./build/resources/test/mmdb-files/geo-lite2/GeoLite2-ASN-Test.mmdb"; + private static final String CDN_ENDPOINT_HOST = "https://devo.geoip.maps.opensearch.org/v1/mmdb/geolite2/manifest.json"; @Test void test_positive_case() { - assertTrue(DbSourceIdentification.isS3Uri(S3_URI)); - assertTrue(DbSourceIdentification.isS3Url(S3_URL)); - assertTrue(DbSourceIdentification.isURL(URL)); - assertTrue(DbSourceIdentification.isFilePath(PATH)); + assertTrue(DatabaseSourceIdentification.isS3Uri(S3_URI)); + assertTrue(DatabaseSourceIdentification.isURL(URL)); + assertTrue(DatabaseSourceIdentification.isFilePath(FILE_PATH)); + assertTrue(DatabaseSourceIdentification.isCDNEndpoint(CDN_ENDPOINT_HOST)); } @Test void test_negative_case() { - assertFalse(DbSourceIdentification.isS3Uri(S3_URL)); - assertFalse(DbSourceIdentification.isS3Uri(URL)); - assertFalse(DbSourceIdentification.isS3Uri(PATH)); + assertFalse(DatabaseSourceIdentification.isS3Uri(CDN_ENDPOINT_HOST)); + assertFalse(DatabaseSourceIdentification.isS3Uri(URL)); + assertFalse(DatabaseSourceIdentification.isS3Uri(DIRECTORY_PATH)); - assertFalse(DbSourceIdentification.isS3Url(S3_URI)); - assertFalse(DbSourceIdentification.isS3Url(URL)); - assertFalse(DbSourceIdentification.isS3Url(PATH)); + assertFalse(DatabaseSourceIdentification.isURL(S3_URI)); + assertFalse(DatabaseSourceIdentification.isURL(CDN_ENDPOINT_HOST)); + assertFalse(DatabaseSourceIdentification.isURL(DIRECTORY_PATH)); - assertFalse(DbSourceIdentification.isURL(S3_URI)); - assertFalse(DbSourceIdentification.isURL(S3_URL)); - assertFalse(DbSourceIdentification.isURL(PATH)); + assertFalse(DatabaseSourceIdentification.isFilePath(S3_URI)); + assertFalse(DatabaseSourceIdentification.isFilePath(CDN_ENDPOINT_HOST)); + assertFalse(DatabaseSourceIdentification.isFilePath(URL)); - assertFalse(DbSourceIdentification.isFilePath(S3_URI)); - assertFalse(DbSourceIdentification.isFilePath(S3_URL)); - assertFalse(DbSourceIdentification.isFilePath(URL)); + assertFalse(DatabaseSourceIdentification.isCDNEndpoint(S3_URI)); + assertFalse(DatabaseSourceIdentification.isCDNEndpoint(DIRECTORY_PATH)); + assertFalse(DatabaseSourceIdentification.isCDNEndpoint(URL)); } @Test - void getDatabasePathTypeTest_PATH() throws NoSuchFieldException, IllegalAccessException { - List databasePath = List.of("./src/test/resources/mmdb-file/geo-lite2"); - DBSourceOptions dbSourceOptions = DbSourceIdentification.getDatabasePathType(databasePath); + void getDatabasePathTypeTest_should_return_PATH_if_file() { + List databasePath = List.of(FILE_PATH); + DBSourceOptions dbSourceOptions = DatabaseSourceIdentification.getDatabasePathType(databasePath); Assertions.assertNotNull(dbSourceOptions); assertThat(dbSourceOptions, equalTo(DBSourceOptions.PATH)); } @Test - void getDatabasePathTypeTest_URL() throws NoSuchFieldException, IllegalAccessException { + void getDatabasePathTypeTest_should_null_if_directory() { + List databasePath = List.of(DIRECTORY_PATH); + DBSourceOptions dbSourceOptions = DatabaseSourceIdentification.getDatabasePathType(databasePath); + Assertions.assertNull(dbSourceOptions); + } + + @Test + void getDatabasePathTypeTest_URL() { List databasePath = List.of("https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-ASN&suffix=tar.gz"); - DBSourceOptions dbSourceOptions = DbSourceIdentification.getDatabasePathType(databasePath); + DBSourceOptions dbSourceOptions = DatabaseSourceIdentification.getDatabasePathType(databasePath); Assertions.assertNotNull(dbSourceOptions); assertThat(dbSourceOptions, equalTo(DBSourceOptions.URL)); } @Test - void getDatabasePathTypeTest_S3() throws NoSuchFieldException, IllegalAccessException { + void getDatabasePathTypeTest_S3() { List databasePath = List.of("s3://mybucket10012023/GeoLite2/"); - DBSourceOptions dbSourceOptions = DbSourceIdentification.getDatabasePathType(databasePath); + DBSourceOptions dbSourceOptions = DatabaseSourceIdentification.getDatabasePathType(databasePath); Assertions.assertNotNull(dbSourceOptions); assertThat(dbSourceOptions, equalTo(DBSourceOptions.S3)); } + @Test + void getDatabasePathTypeTest_CDN() { + List databasePath = List.of(CDN_ENDPOINT_HOST); + DBSourceOptions dbSourceOptions = DatabaseSourceIdentification.getDatabasePathType(databasePath); + Assertions.assertNotNull(dbSourceOptions); + assertThat(dbSourceOptions, equalTo(DBSourceOptions.HTTP_MANIFEST)); + } + @Test void isS3Uri_NullPointerException_test() { assertDoesNotThrow(() -> { - DbSourceIdentification.isS3Uri(null); + DatabaseSourceIdentification.isS3Uri(null); }); } } \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/IPValidationCheckTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/IPValidationCheckTest.java index 487f131f81..f18d3019b6 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/IPValidationCheckTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/IPValidationCheckTest.java @@ -10,27 +10,25 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.junit.jupiter.MockitoExtension; -import java.net.UnknownHostException; - @ExtendWith(MockitoExtension.class) class IPValidationCheckTest { private static final String PRIVATE_IP_ADDRESS = "192.168.29.233"; private static final String PUBLIC_IP_ADDRESS = "2001:4860:4860::8888"; - private static final String INVALID_IP_ADDRESS = "255.255.255.0"; + private static final String INVALID_IP_ADDRESS = "255.255.255.999"; @Test - void ipValidationcheckTest_positive() throws UnknownHostException { + void ipValidationcheckTest_public() { Assertions.assertTrue(IPValidationCheck.isPublicIpAddress(PUBLIC_IP_ADDRESS)); } @Test - void ipValidationcheckTest_negative() throws UnknownHostException { + void ipValidationcheckTest_negative() { Assertions.assertFalse(IPValidationCheck.isPublicIpAddress(PRIVATE_IP_ADDRESS)); } @Test - void ipValidationcheckTest_invalid() throws UnknownHostException { - Assertions.assertTrue(IPValidationCheck.isPublicIpAddress(INVALID_IP_ADDRESS)); + void ipValidationcheckTest_invalid() { + Assertions.assertFalse(IPValidationCheck.isPublicIpAddress(INVALID_IP_ADDRESS)); } } \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/LicenseTypeCheckTest.java b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/LicenseTypeCheckTest.java index 02d40cdf1e..5530367d7c 100644 --- a/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/LicenseTypeCheckTest.java +++ b/data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/utils/LicenseTypeCheckTest.java @@ -17,30 +17,24 @@ @ExtendWith(MockitoExtension.class) class LicenseTypeCheckTest { - private static final String FOLDER_PATH_GEO_LITE2 = "./src/test/resources/mmdb-file/geo-lite2"; - private static final String FOLDER_PATH_GEO_ENTERPRISE = "./src/test/resources/mmdb-file/geo-enterprise"; + private static final String FOLDER_PATH_GEO_LITE2 = "./build/resources/test/mmdb-files/geo-lite2"; + private static final String FOLDER_PATH_GEO_ENTERPRISE = "./build/resources/test/mmdb-files/geo-enterprise"; private LicenseTypeCheck createObjectUnderTest() { return new LicenseTypeCheck(); } @Test - void isGeoLite2OrEnterpriseLicenseTest_positive() { + void test_isGeoLite2OrEnterpriseLicenseTest_should_return_free_when_geolite2_databases_are_used() { final LicenseTypeCheck objectUnderTest = createObjectUnderTest(); LicenseTypeOptions licenseTypeOptionsFree = objectUnderTest.isGeoLite2OrEnterpriseLicense(FOLDER_PATH_GEO_LITE2); assertThat(licenseTypeOptionsFree, equalTo(LicenseTypeOptions.FREE)); - - LicenseTypeOptions licenseTypeOptionsEnterprise = objectUnderTest.isGeoLite2OrEnterpriseLicense(FOLDER_PATH_GEO_ENTERPRISE); - assertThat(licenseTypeOptionsEnterprise, equalTo(LicenseTypeOptions.ENTERPRISE)); } @Test - void isGeoLite2OrEnterpriseLicenseTest_negative() { + void test_isGeoLite2OrEnterpriseLicenseTest_should_return_enterprise_when_geoip2_databases_are_used() { final LicenseTypeCheck objectUnderTest = createObjectUnderTest(); LicenseTypeOptions licenseTypeOptionsFree = objectUnderTest.isGeoLite2OrEnterpriseLicense(FOLDER_PATH_GEO_ENTERPRISE); - assertThat(licenseTypeOptionsFree, not(equalTo(LicenseTypeOptions.FREE))); - - LicenseTypeOptions licenseTypeOptionsEnterprise = objectUnderTest.isGeoLite2OrEnterpriseLicense(FOLDER_PATH_GEO_LITE2); - assertThat(licenseTypeOptionsEnterprise, not(equalTo(LicenseTypeOptions.ENTERPRISE))); + assertThat(licenseTypeOptionsFree, not(equalTo(LicenseTypeOptions.ENTERPRISE))); } } \ No newline at end of file diff --git a/data-prepper-plugins/geoip-processor/src/test/resources/geoip_service_config.yaml b/data-prepper-plugins/geoip-processor/src/test/resources/geoip_service_config.yaml index 54d3ebb818..40114779f0 100644 --- a/data-prepper-plugins/geoip-processor/src/test/resources/geoip_service_config.yaml +++ b/data-prepper-plugins/geoip-processor/src/test/resources/geoip_service_config.yaml @@ -1,9 +1,13 @@ extensions: geoip_service: maxmind: - database_paths: ["path1", "path2"] + databases: + city: "city_path" + country: "country_path" database_refresh_interval: "P10D" - cache_size: 2048 + cache_count: 2048 + insecure: true + database_destination: "/tst/resources" aws: region: "us-east-1" sts_role_arn: "arn:aws:iam::123456789:role/data-prepper-execution-role" \ No newline at end of file diff --git a/data-prepper-plugins/grok-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessor.java b/data-prepper-plugins/grok-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessor.java index 1c3c9e6d15..391d50f603 100644 --- a/data-prepper-plugins/grok-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessor.java +++ b/data-prepper-plugins/grok-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessor.java @@ -33,6 +33,7 @@ import java.nio.file.Files; import java.nio.file.NotDirectoryException; import java.nio.file.Path; +import java.time.Instant; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; @@ -53,6 +54,8 @@ import java.util.stream.Collectors; import static org.opensearch.dataprepper.logging.DataPrepperMarkers.EVENT; +import static org.opensearch.dataprepper.plugins.processor.grok.GrokProcessorConfig.TOTAL_PATTERNS_ATTEMPTED_METADATA_KEY; +import static org.opensearch.dataprepper.plugins.processor.grok.GrokProcessorConfig.TOTAL_TIME_SPENT_IN_GROK_METADATA_KEY; @SingleThread @@ -120,6 +123,8 @@ public GrokProcessor(final PluginSetting pluginSetting, final ExpressionEvaluato @Override public Collection> doExecute(final Collection> records) { for (final Record record : records) { + + final Instant startTime = Instant.now(); final Event event = record.getData(); try { if (Objects.nonNull(grokProcessorConfig.getGrokWhen()) && !expressionEvaluator.evaluateConditional(grokProcessorConfig.getGrokWhen(), event)) { @@ -141,6 +146,16 @@ public Collection> doExecute(final Collection> recor LOG.error(EVENT, "An exception occurred when matching record [{}]", record.getData(), e); grokProcessingErrorsCounter.increment(); } + + final Instant endTime = Instant.now(); + + Long totalEventTimeInGrok = (Long) event.getMetadata().getAttribute(TOTAL_TIME_SPENT_IN_GROK_METADATA_KEY); + if (totalEventTimeInGrok == null) { + totalEventTimeInGrok = 0L; + } + + final long timeSpentInThisGrok = endTime.toEpochMilli() - startTime.toEpochMilli(); + event.getMetadata().setAttribute(TOTAL_TIME_SPENT_IN_GROK_METADATA_KEY, totalEventTimeInGrok + timeSpentInThisGrok); } return records; } @@ -228,6 +243,8 @@ private void compileMatchPatterns() { private void matchAndMerge(final Event event) { final Map grokkedCaptures = new HashMap<>(); + int patternsAttempted = 0; + for (final Map.Entry> entry : fieldToGrok.entrySet()) { for (final Grok grok : entry.getValue()) { final String value = event.get(entry.getKey(), String.class); @@ -238,6 +255,8 @@ private void matchAndMerge(final Event event) { final Map captures = match.capture(); mergeCaptures(grokkedCaptures, captures); + patternsAttempted++; + if (shouldBreakOnMatch(grokkedCaptures)) { break; } @@ -262,6 +281,15 @@ private void matchAndMerge(final Event event) { } else { grokProcessingMatchCounter.increment(); } + + if (grokProcessorConfig.getIncludePerformanceMetadata()) { + Integer totalPatternsAttemptedForEvent = (Integer) event.getMetadata().getAttribute(TOTAL_PATTERNS_ATTEMPTED_METADATA_KEY); + if (totalPatternsAttemptedForEvent == null) { + totalPatternsAttemptedForEvent = 0; + } + + event.getMetadata().setAttribute(TOTAL_PATTERNS_ATTEMPTED_METADATA_KEY, totalPatternsAttemptedForEvent + patternsAttempted); + } } private void mergeCaptures(final Map original, final Map updates) { diff --git a/data-prepper-plugins/grok-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessorConfig.java b/data-prepper-plugins/grok-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessorConfig.java index cea6f58c04..de9daf91d5 100644 --- a/data-prepper-plugins/grok-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessorConfig.java +++ b/data-prepper-plugins/grok-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessorConfig.java @@ -11,6 +11,11 @@ import java.util.Map; public class GrokProcessorConfig { + + static final String TOTAL_PATTERNS_ATTEMPTED_METADATA_KEY = "_total_grok_patterns_attempted"; + + static final String TOTAL_TIME_SPENT_IN_GROK_METADATA_KEY = "_total_grok_processing_time"; + static final String BREAK_ON_MATCH = "break_on_match"; static final String KEEP_EMPTY_CAPTURES = "keep_empty_captures"; static final String MATCH = "match"; @@ -25,6 +30,8 @@ public class GrokProcessorConfig { static final String TAGS_ON_MATCH_FAILURE = "tags_on_match_failure"; static final String TAGS_ON_TIMEOUT = "tags_on_timeout"; + static final String INCLUDE_PERFORMANCE_METADATA = "performance_metadata"; + static final boolean DEFAULT_BREAK_ON_MATCH = true; static final boolean DEFAULT_KEEP_EMPTY_CAPTURES = false; static final boolean DEFAULT_NAMED_CAPTURES_ONLY = true; @@ -46,6 +53,8 @@ public class GrokProcessorConfig { private final List tagsOnMatchFailure; private final List tagsOnTimeout; + private final boolean includePerformanceMetadata; + private GrokProcessorConfig(final boolean breakOnMatch, final boolean keepEmptyCaptures, final Map> match, @@ -58,7 +67,8 @@ private GrokProcessorConfig(final boolean breakOnMatch, final String targetKey, final String grokWhen, final List tagsOnMatchFailure, - final List tagsOnTimeout) { + final List tagsOnTimeout, + final boolean includePerformanceMetadata) { this.breakOnMatch = breakOnMatch; this.keepEmptyCaptures = keepEmptyCaptures; @@ -73,6 +83,7 @@ private GrokProcessorConfig(final boolean breakOnMatch, this.grokWhen = grokWhen; this.tagsOnMatchFailure = tagsOnMatchFailure; this.tagsOnTimeout = tagsOnTimeout.isEmpty() ? tagsOnMatchFailure : tagsOnTimeout; + this.includePerformanceMetadata = includePerformanceMetadata; } public static GrokProcessorConfig buildConfig(final PluginSetting pluginSetting) { @@ -88,7 +99,8 @@ public static GrokProcessorConfig buildConfig(final PluginSetting pluginSetting) pluginSetting.getStringOrDefault(TARGET_KEY, DEFAULT_TARGET_KEY), pluginSetting.getStringOrDefault(GROK_WHEN, null), pluginSetting.getTypedList(TAGS_ON_MATCH_FAILURE, String.class), - pluginSetting.getTypedList(TAGS_ON_TIMEOUT, String.class)); + pluginSetting.getTypedList(TAGS_ON_TIMEOUT, String.class), + pluginSetting.getBooleanOrDefault(INCLUDE_PERFORMANCE_METADATA, false)); } public boolean isBreakOnMatch() { @@ -140,4 +152,6 @@ public List getTagsOnMatchFailure() { public List getTagsOnTimeout() { return tagsOnTimeout; } + + public boolean getIncludePerformanceMetadata() { return includePerformanceMetadata; } } diff --git a/data-prepper-plugins/grok-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessorConfigTests.java b/data-prepper-plugins/grok-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessorConfigTests.java index 7bd3ac662e..eb69968a96 100644 --- a/data-prepper-plugins/grok-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessorConfigTests.java +++ b/data-prepper-plugins/grok-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessorConfigTests.java @@ -77,6 +77,7 @@ public void testDefault() { assertThat(grokProcessorConfig.getGrokWhen(), equalTo(null)); assertThat(grokProcessorConfig.getTagsOnMatchFailure(), equalTo(Collections.emptyList())); assertThat(grokProcessorConfig.getTagsOnTimeout(), equalTo(Collections.emptyList())); + assertThat(grokProcessorConfig.getIncludePerformanceMetadata(), equalTo(false)); } @Test @@ -91,7 +92,8 @@ public void testValidConfig() { TEST_PATTERNS_FILES_GLOB, TEST_PATTERN_DEFINITIONS, TEST_TIMEOUT_MILLIS, - TEST_TARGET_KEY); + TEST_TARGET_KEY, + true); final GrokProcessorConfig grokProcessorConfig = GrokProcessorConfig.buildConfig(validPluginSetting); @@ -105,6 +107,7 @@ public void testValidConfig() { assertThat(grokProcessorConfig.getTargetKey(), equalTo(TEST_TARGET_KEY)); assertThat(grokProcessorConfig.isNamedCapturesOnly(), equalTo(false)); assertThat(grokProcessorConfig.getTimeoutMillis(), equalTo(TEST_TIMEOUT_MILLIS)); + assertThat(grokProcessorConfig.getIncludePerformanceMetadata(), equalTo(true)); } @Test @@ -119,7 +122,8 @@ public void testInvalidConfig() { TEST_PATTERNS_FILES_GLOB, TEST_PATTERN_DEFINITIONS, TEST_TIMEOUT_MILLIS, - TEST_TARGET_KEY); + TEST_TARGET_KEY, + false); invalidPluginSetting.getSettings().put(GrokProcessorConfig.MATCH, TEST_INVALID_MATCH); @@ -135,7 +139,8 @@ private PluginSetting completePluginSettingForGrokProcessor(final boolean breakO final String patternsFilesGlob, final Map patternDefinitions, final int timeoutMillis, - final String targetKey) { + final String targetKey, + final boolean includePerformanceMetadata) { final Map settings = new HashMap<>(); settings.put(GrokProcessorConfig.BREAK_ON_MATCH, breakOnMatch); settings.put(GrokProcessorConfig.NAMED_CAPTURES_ONLY, namedCapturesOnly); @@ -147,6 +152,7 @@ private PluginSetting completePluginSettingForGrokProcessor(final boolean breakO settings.put(GrokProcessorConfig.PATTERNS_FILES_GLOB, patternsFilesGlob); settings.put(GrokProcessorConfig.TIMEOUT_MILLIS, timeoutMillis); settings.put(GrokProcessorConfig.TARGET_KEY, targetKey); + settings.put(GrokProcessorConfig.INCLUDE_PERFORMANCE_METADATA, includePerformanceMetadata); return new PluginSetting(PLUGIN_NAME, settings); } diff --git a/data-prepper-plugins/grok-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessorTests.java b/data-prepper-plugins/grok-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessorTests.java index bb3b6857c7..d2ff11c3e0 100644 --- a/data-prepper-plugins/grok-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessorTests.java +++ b/data-prepper-plugins/grok-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/grok/GrokProcessorTests.java @@ -46,6 +46,7 @@ import static org.hamcrest.CoreMatchers.hasItem; import static org.hamcrest.CoreMatchers.notNullValue; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.greaterThan; import static org.hamcrest.Matchers.not; import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.eq; @@ -58,6 +59,8 @@ import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; import static org.opensearch.dataprepper.plugins.processor.grok.GrokProcessor.EXECUTOR_SERVICE_SHUTDOWN_TIMEOUT; +import static org.opensearch.dataprepper.plugins.processor.grok.GrokProcessorConfig.TOTAL_PATTERNS_ATTEMPTED_METADATA_KEY; +import static org.opensearch.dataprepper.plugins.processor.grok.GrokProcessorConfig.TOTAL_TIME_SPENT_IN_GROK_METADATA_KEY; import static org.opensearch.dataprepper.test.matcher.MapEquals.isEqualWithoutTimestamp; @@ -160,6 +163,8 @@ private GrokProcessor createObjectUnderTest() { @Test public void testMatchMerge() throws JsonProcessingException, ExecutionException, InterruptedException, TimeoutException { + pluginSetting.getSettings().put(GrokProcessorConfig.INCLUDE_PERFORMANCE_METADATA, false); + grokProcessor = createObjectUnderTest(); capture.put("key_capture_1", "value_capture_1"); @@ -182,7 +187,11 @@ public void testMatchMerge() throws JsonProcessingException, ExecutionException, assertThat(grokkedRecords.size(), equalTo(1)); assertThat(grokkedRecords.get(0), notNullValue()); + assertThat(grokkedRecords.get(0).getData(), notNullValue()); + assertThat(grokkedRecords.get(0).getData().getMetadata(), notNullValue()); + assertThat(grokkedRecords.get(0).getData().getMetadata().getAttribute(TOTAL_PATTERNS_ATTEMPTED_METADATA_KEY), equalTo(null)); assertRecordsAreEqual(grokkedRecords.get(0), resultRecord); + verify(grokProcessingMatchCounter, times(1)).increment(); verify(grokProcessingTime, times(1)).record(any(Runnable.class)); verifyNoInteractions(grokProcessingErrorsCounter, grokProcessingMismatchCounter, grokProcessingTimeoutsCounter); @@ -516,6 +525,65 @@ public void testNoCaptures() throws JsonProcessingException { verifyNoInteractions(grokProcessingErrorsCounter, grokProcessingMatchCounter, grokProcessingTimeoutsCounter); } + @Test + public void testMatchOnSecondPattern() throws JsonProcessingException { + pluginSetting.getSettings().put(GrokProcessorConfig.INCLUDE_PERFORMANCE_METADATA, true); + + when(match.capture()).thenReturn(Collections.emptyMap()); + when(grokSecondMatch.match(messageInput)).thenReturn(secondMatch); + when(secondMatch.capture()).thenReturn(capture); + + grokProcessor = createObjectUnderTest(); + + final Map testData = new HashMap(); + testData.put("message", messageInput); + final Record record = buildRecordWithEvent(testData); + + final List> grokkedRecords = (List>) grokProcessor.doExecute(Collections.singletonList(record)); + + assertThat(grokkedRecords.size(), equalTo(1)); + assertThat(grokkedRecords.get(0), notNullValue()); + assertThat(grokkedRecords.get(0).getData(), notNullValue()); + assertThat(grokkedRecords.get(0).getData().getMetadata(), notNullValue()); + assertThat(grokkedRecords.get(0).getData().getMetadata().getAttribute(TOTAL_PATTERNS_ATTEMPTED_METADATA_KEY), equalTo(2)); + assertThat((Long) grokkedRecords.get(0).getData().getMetadata().getAttribute(TOTAL_TIME_SPENT_IN_GROK_METADATA_KEY), greaterThan(0L)); + assertRecordsAreEqual(grokkedRecords.get(0), record); + verify(grokProcessingMismatchCounter, times(1)).increment(); + verify(grokProcessingTime, times(1)).record(any(Runnable.class)); + verifyNoInteractions(grokProcessingErrorsCounter, grokProcessingMatchCounter, grokProcessingTimeoutsCounter); + } + + @Test + public void testMatchOnSecondPatternWithExistingMetadataForTotalPatternMatches() throws JsonProcessingException { + pluginSetting.getSettings().put(GrokProcessorConfig.INCLUDE_PERFORMANCE_METADATA, true); + + when(match.capture()).thenReturn(Collections.emptyMap()); + when(grokSecondMatch.match(messageInput)).thenReturn(secondMatch); + when(secondMatch.capture()).thenReturn(capture); + + grokProcessor = createObjectUnderTest(); + + final Map testData = new HashMap(); + testData.put("message", messageInput); + final Record record = buildRecordWithEvent(testData); + + record.getData().getMetadata().setAttribute(TOTAL_PATTERNS_ATTEMPTED_METADATA_KEY, 1); + record.getData().getMetadata().setAttribute(TOTAL_TIME_SPENT_IN_GROK_METADATA_KEY, 300L); + + final List> grokkedRecords = (List>) grokProcessor.doExecute(Collections.singletonList(record)); + + assertThat(grokkedRecords.size(), equalTo(1)); + assertThat(grokkedRecords.get(0), notNullValue()); + assertThat(grokkedRecords.get(0).getData(), notNullValue()); + assertThat(grokkedRecords.get(0).getData().getMetadata(), notNullValue()); + assertThat(grokkedRecords.get(0).getData().getMetadata().getAttribute(TOTAL_PATTERNS_ATTEMPTED_METADATA_KEY), equalTo(3)); + assertThat((Long) grokkedRecords.get(0).getData().getMetadata().getAttribute(TOTAL_TIME_SPENT_IN_GROK_METADATA_KEY), greaterThan(300L)); + assertRecordsAreEqual(grokkedRecords.get(0), record); + verify(grokProcessingMismatchCounter, times(1)).increment(); + verify(grokProcessingTime, times(1)).record(any(Runnable.class)); + verifyNoInteractions(grokProcessingErrorsCounter, grokProcessingMatchCounter, grokProcessingTimeoutsCounter); + } + @Nested class WithTags { private String tagOnMatchFailure1; diff --git a/data-prepper-plugins/http-common/build.gradle b/data-prepper-plugins/http-common/build.gradle new file mode 100644 index 0000000000..54fa5d346d --- /dev/null +++ b/data-prepper-plugins/http-common/build.gradle @@ -0,0 +1,20 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +dependencies { + implementation 'org.apache.httpcomponents:httpcore:4.4.16' + testImplementation testLibs.bundles.junit +} + +jacocoTestCoverageVerification { + dependsOn jacocoTestReport + violationRules { + rule { //in addition to core projects rule + limit { + minimum = 0.90 + } + } + } +} \ No newline at end of file diff --git a/data-prepper-plugins/http-common/data/certificate/test_cert.crt b/data-prepper-plugins/http-common/data/certificate/test_cert.crt new file mode 100644 index 0000000000..26c78d1411 --- /dev/null +++ b/data-prepper-plugins/http-common/data/certificate/test_cert.crt @@ -0,0 +1,14 @@ +-----BEGIN CERTIFICATE----- +MIICHTCCAYYCCQD4hqYeYDQZADANBgkqhkiG9w0BAQUFADBSMQswCQYDVQQGEwJV +UzELMAkGA1UECAwCVFgxDzANBgNVBAcMBkF1c3RpbjEPMA0GA1UECgwGQW1hem9u +MRQwEgYDVQQLDAtEYXRhcHJlcHBlcjAgFw0yMTA2MjUxOTIzMTBaGA8yMTIxMDYw +MTE5MjMxMFowUjELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAlRYMQ8wDQYDVQQHDAZB +dXN0aW4xDzANBgNVBAoMBkFtYXpvbjEUMBIGA1UECwwLRGF0YXByZXBwZXIwgZ8w +DQYJKoZIhvcNAQEBBQADgY0AMIGJAoGBAKrb3YhdKbQ5PtLHall10iLZC9ZdDVrq +HOvqVSM8NHlL8f82gJ3l0n9k7hYc5eKisutaS9eDTmJ+Dnn8xn/qPSKTIq9Wh+OZ +O+e9YEEpI/G4F9KpGULgMyRg9sJK0GlZdEt9o5GJNJIJUkptJU5eiLuE0IV+jyJo +Nvm8OE6EJPqxAgMBAAEwDQYJKoZIhvcNAQEFBQADgYEAjgnX5n/Tt7eo9uakIGAb +uBhvYdR8JqKXqF9rjFJ/MIK7FdQSF/gCdjnvBhzLlZFK/Nb6MGKoSKm5Lcr75LgC +FyhIwp3WlqQksiMFnOypYVY71vqDgj6UKdMaOBgthsYhngj8lC+wsVzWqQvkJ2Qg +/GAIzJwiZfXiaevQHRk79qI= +-----END CERTIFICATE----- diff --git a/data-prepper-plugins/http-common/src/main/java/org/opensearch/dataprepper/plugins/truststore/TrustStoreProvider.java b/data-prepper-plugins/http-common/src/main/java/org/opensearch/dataprepper/plugins/truststore/TrustStoreProvider.java new file mode 100644 index 0000000000..2979a6d7a1 --- /dev/null +++ b/data-prepper-plugins/http-common/src/main/java/org/opensearch/dataprepper/plugins/truststore/TrustStoreProvider.java @@ -0,0 +1,113 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.truststore; + +import org.apache.http.ssl.SSLContextBuilder; +import org.apache.http.ssl.SSLContexts; +import org.apache.http.ssl.TrustStrategy; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.net.ssl.SSLContext; +import javax.net.ssl.TrustManager; +import javax.net.ssl.TrustManagerFactory; +import java.io.ByteArrayInputStream; +import java.io.InputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.security.KeyStore; +import java.security.cert.Certificate; +import java.security.cert.CertificateFactory; +import java.security.cert.X509Certificate; + +public class TrustStoreProvider { + private static final Logger LOG = LoggerFactory.getLogger(TrustStoreProvider.class); + + public static TrustManager[] createTrustManager(final Path certificatePath) { + LOG.info("Using the certificate path {} to create trust manager.", certificatePath.toString()); + try { + final KeyStore keyStore = createKeyStore(certificatePath); + final TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance("X509"); + trustManagerFactory.init(keyStore); + return trustManagerFactory.getTrustManagers(); + } catch (Exception ex) { + throw new RuntimeException(ex.getMessage(), ex); + } + } + + public static TrustManager[] createTrustManager(final String certificateContent) { + LOG.info("Using the certificate content to create trust manager."); + try (InputStream certificateInputStream = new ByteArrayInputStream(certificateContent.getBytes())) { + final KeyStore keyStore = createKeyStore(certificateInputStream); + final TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance("X509"); + trustManagerFactory.init(keyStore); + return trustManagerFactory.getTrustManagers(); + } catch (Exception ex) { + throw new RuntimeException(ex.getMessage(), ex); + } + } + + public static TrustManager[] createTrustAllManager() { + LOG.info("Using the trust all manager to create trust manager."); + return new TrustManager[]{ + new X509TrustAllManager() + }; + } + + private static KeyStore createKeyStore(final Path certificatePath) throws Exception { + try (InputStream certificateInputStream = Files.newInputStream(certificatePath)) { + return createKeyStore(certificateInputStream); + } + } + + private static KeyStore createKeyStore(final InputStream certificateInputStream) throws Exception { + final CertificateFactory factory = CertificateFactory.getInstance("X.509"); + final Certificate trustedCa = factory.generateCertificate(certificateInputStream); + final KeyStore trustStore = KeyStore.getInstance("pkcs12"); + trustStore.load(null, null); + trustStore.setCertificateEntry("ca", trustedCa); + return trustStore; + } + + public static SSLContext createSSLContext(final Path certificatePath) { + LOG.info("Using the certificate path to create SSL context."); + try (InputStream is = Files.newInputStream(certificatePath)) { + return createSSLContext(is); + } catch (Exception ex) { + throw new RuntimeException(ex.getMessage(), ex); + } + } + + public static SSLContext createSSLContext(final String certificateContent) { + LOG.info("Using the certificate content to create SSL context."); + try (InputStream certificateInputStream = new ByteArrayInputStream(certificateContent.getBytes())) { + return createSSLContext(certificateInputStream); + } catch (Exception ex) { + throw new RuntimeException(ex.getMessage(), ex); + } + } + + private static SSLContext createSSLContext(final InputStream certificateInputStream) throws Exception { + KeyStore trustStore = createKeyStore(certificateInputStream); + SSLContextBuilder sslContextBuilder = SSLContexts.custom() + .loadTrustMaterial(trustStore, null); + return sslContextBuilder.build(); + } + + public static SSLContext createSSLContextWithTrustAllStrategy() { + LOG.info("Using the trust all strategy to create SSL context."); + try { + return SSLContexts.custom().loadTrustMaterial(null, new TrustStrategy() { + @Override + public boolean isTrusted(X509Certificate[] chain, String authType) { + return true; + } + }).build(); + } catch (Exception ex) { + throw new RuntimeException(ex.getMessage(), ex); + } + } +} diff --git a/data-prepper-plugins/http-common/src/main/java/org/opensearch/dataprepper/plugins/truststore/X509TrustAllManager.java b/data-prepper-plugins/http-common/src/main/java/org/opensearch/dataprepper/plugins/truststore/X509TrustAllManager.java new file mode 100644 index 0000000000..378c29dcb8 --- /dev/null +++ b/data-prepper-plugins/http-common/src/main/java/org/opensearch/dataprepper/plugins/truststore/X509TrustAllManager.java @@ -0,0 +1,22 @@ +package org.opensearch.dataprepper.plugins.truststore; + +import javax.net.ssl.X509TrustManager; +import java.security.cert.CertificateException; +import java.security.cert.X509Certificate; + +public class X509TrustAllManager implements X509TrustManager { + @Override + public void checkClientTrusted(X509Certificate[] chain, String authType) throws CertificateException { + + } + + @Override + public void checkServerTrusted(X509Certificate[] chain, String authType) throws CertificateException { + + } + + @Override + public X509Certificate[] getAcceptedIssuers() { + return null; + } +} diff --git a/data-prepper-plugins/http-common/src/test/java/org/opensearch/dataprepper/plugins/truststore/TrustStoreProviderTest.java b/data-prepper-plugins/http-common/src/test/java/org/opensearch/dataprepper/plugins/truststore/TrustStoreProviderTest.java new file mode 100644 index 0000000000..c01b36c6a8 --- /dev/null +++ b/data-prepper-plugins/http-common/src/test/java/org/opensearch/dataprepper/plugins/truststore/TrustStoreProviderTest.java @@ -0,0 +1,85 @@ +package org.opensearch.dataprepper.plugins.truststore; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.junit.jupiter.MockitoExtension; + +import javax.net.ssl.SSLContext; +import javax.net.ssl.TrustManager; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; + +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.notNullValue; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.arrayWithSize; +import static org.hamcrest.Matchers.instanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; + +@ExtendWith(MockitoExtension.class) +class TrustStoreProviderTest { + + private TrustStoreProvider trustStoreProvider; + + @BeforeEach + void setUp() { + trustStoreProvider = new TrustStoreProvider(); + } + + @Test + void createTrustManagerWithCertificatePath() { + final Path certFilePath = Path.of("data/certificate/test_cert.crt"); + final TrustManager[] trustManagers = trustStoreProvider.createTrustManager(certFilePath); + assertThat(trustManagers, is(notNullValue())); + } + + @Test + void createTrustManagerWithInvalidCertificatePath() { + final Path certFilePath = Path.of("data/certificate/cert_doesnt_exist.crt"); + assertThrows(RuntimeException.class, () -> trustStoreProvider.createTrustManager(certFilePath)); + } + + @Test + void createTrustManagerWithCertificateContent() throws IOException { + final Path certFilePath = Path.of("data/certificate/test_cert.crt"); + final String certificateContent = Files.readString(certFilePath); + final TrustManager[] trustManagers = trustStoreProvider.createTrustManager(certificateContent); + assertThat(trustManagers, is(notNullValue())); + } + + @Test + void createTrustManagerWithInvalidCertificateContent() { + assertThrows(RuntimeException.class, () -> trustStoreProvider.createTrustManager("invalid certificate content")); + } + + @Test + void createTrustAllManager() { + final TrustManager[] trustManagers = trustStoreProvider.createTrustAllManager(); + assertThat(trustManagers, is(notNullValue())); + assertThat(trustManagers, is(arrayWithSize(1))); + assertThat(trustManagers[0], is(instanceOf(X509TrustAllManager.class))); + } + + @Test + void createSSLContextWithCertificatePath() { + final Path certFilePath = Path.of("data/certificate/test_cert.crt"); + final SSLContext sslContext = trustStoreProvider.createSSLContext(certFilePath); + assertThat(sslContext, is(notNullValue())); + } + + @Test + void createSSLContextWithCertificateContent() throws IOException { + final Path certFilePath = Path.of("data/certificate/test_cert.crt"); + final String certificateContent = Files.readString(certFilePath); + final SSLContext sslContext = trustStoreProvider.createSSLContext(certificateContent); + assertThat(sslContext, is(notNullValue())); + } + + @Test + void createSSLContextWithTrustAllStrategy() { + final SSLContext sslContext = trustStoreProvider.createSSLContextWithTrustAllStrategy(); + assertThat(sslContext, is(notNullValue())); + } +} \ No newline at end of file diff --git a/data-prepper-plugins/kafka-connect-plugins/README.md b/data-prepper-plugins/kafka-connect-plugins/README.md index e4bd43dd68..e3517b2885 100644 --- a/data-prepper-plugins/kafka-connect-plugins/README.md +++ b/data-prepper-plugins/kafka-connect-plugins/README.md @@ -1,5 +1,12 @@ # Kafka Connect Source +The `kafka-connect-plugins` project has not been released. And the maintainers have no plans to pick this work up. +If you wish to pick this project up, please open a GitHub issue to discuss. +The original code is available in Git history if needed. + + +## Old README information + This is a source plugin that start a Kafka Connect and Connectors. Please note that the Kafka Connect Source has to work with Kafka Buffer. ## Usages diff --git a/data-prepper-plugins/kafka-connect-plugins/build.gradle b/data-prepper-plugins/kafka-connect-plugins/build.gradle deleted file mode 100644 index 8a7d8f20fa..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/build.gradle +++ /dev/null @@ -1,68 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -plugins { - id 'java' -} - -configurations.all { - exclude group: 'org.apache.zookeeper', module: 'zookeeper' -} - -dependencies { - implementation project(':data-prepper-plugins:aws-plugin-api') - implementation project(':data-prepper-plugins:common') - implementation project(':data-prepper-plugins:kafka-plugins') - implementation 'org.apache.kafka:connect-runtime:3.5.1' - implementation 'software.amazon.awssdk:sts' - implementation 'software.amazon.awssdk:secretsmanager' - implementation 'javax.validation:validation-api:2.0.1.Final' - implementation libs.reflections.core - implementation 'io.micrometer:micrometer-core' - implementation ('io.confluent:kafka-schema-registry:7.5.0') { - exclude group: 'org.glassfish.jersey.containers', module: 'jersey-container-servlet' - exclude group: 'org.glassfish.jersey.inject', module: 'jersey-hk2' - exclude group: 'org.glassfish.jersey.ext', module: 'jersey-bean-validation' - } - // Common Debezium Connector - implementation 'io.debezium:debezium-api:2.3.0.Final' - implementation 'io.debezium:debezium-core:2.3.0.Final' - implementation 'io.debezium:debezium-storage-kafka:2.3.0.Final' - implementation 'io.debezium:debezium-storage-file:2.3.0.Final' - // Debezium MySQL Connector - implementation 'org.antlr:antlr4-runtime:4.10.1' - implementation 'io.debezium:debezium-connector-mysql:2.3.0.Final' - implementation 'io.debezium:debezium-ddl-parser:2.3.0.Final' - implementation 'com.zendesk:mysql-binlog-connector-java:0.28.1' - implementation 'com.mysql:mysql-connector-j:8.0.33' - implementation 'com.github.luben:zstd-jni:1.5.0-2' - // Debezium Postgres connector - implementation 'io.debezium:debezium-connector-postgres:2.3.0.Final' - implementation 'org.postgresql:postgresql:42.5.1' - implementation 'com.google.protobuf:protobuf-java:3.19.6' - // Debezium Mongodb connector - implementation 'io.debezium:debezium-connector-mongodb:2.3.0.Final' - implementation 'org.mongodb:mongodb-driver-core:4.7.1' - implementation 'org.mongodb:mongodb-driver-sync:4.7.1' - implementation 'org.mongodb:bson:4.7.1' - runtimeOnly 'org.mongodb:bson-record-codec:4.7.1' - // test - testImplementation project(':data-prepper-test-common') - testImplementation project(':data-prepper-core') - testImplementation 'org.yaml:snakeyaml:2.0' - testImplementation 'com.fasterxml.jackson.datatype:jackson-datatype-jsr310' - testImplementation testLibs.mockito.inline -} - -jacocoTestCoverageVerification { - dependsOn jacocoTestReport - violationRules { - rule { //in addition to core projects rule - limit { - minimum = 0.90 - } - } - } -} \ No newline at end of file diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/ConnectorConfig.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/ConnectorConfig.java deleted file mode 100644 index b92e78cdd1..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/ConnectorConfig.java +++ /dev/null @@ -1,37 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.configuration; - -import com.fasterxml.jackson.annotation.JsonProperty; -import org.opensearch.dataprepper.plugins.kafkaconnect.util.Connector; - -import java.util.List; -import java.util.Properties; - -public abstract class ConnectorConfig { - @JsonProperty("force_update") - public Boolean forceUpdate = false; - private String bootstrapServers; - private Properties authProperties; - - public abstract List buildConnectors(); - - public Properties getAuthProperties() { - return this.authProperties; - } - - public void setAuthProperties(Properties authProperties) { - this.authProperties = authProperties; - } - - public String getBootstrapServers() { - return this.bootstrapServers; - } - - public void setBootstrapServers(String bootStrapServers) { - this.bootstrapServers = bootStrapServers; - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/CredentialsConfig.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/CredentialsConfig.java deleted file mode 100644 index 714a681a97..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/CredentialsConfig.java +++ /dev/null @@ -1,91 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.configuration; - -import com.fasterxml.jackson.annotation.JsonCreator; -import com.fasterxml.jackson.annotation.JsonProperty; -import com.fasterxml.jackson.core.type.TypeReference; -import com.fasterxml.jackson.databind.ObjectMapper; -import org.opensearch.dataprepper.plugins.kafkaconnect.util.SecretManagerHelper; - -import java.util.Map; - -public class CredentialsConfig { - private final String username; - private final String password; - - @JsonCreator - public CredentialsConfig(@JsonProperty("plaintext") final PlainText plainText, - @JsonProperty("secret_manager") final SecretManager secretManager) { - if (plainText != null && secretManager != null) { - throw new IllegalArgumentException("plaintext and secret_manager cannot both be set"); - } - if (plainText != null) { - if (plainText.username == null || plainText.password == null) { - throw new IllegalArgumentException("user and password must be set for plaintext credentials"); - } - this.username = plainText.username; - this.password = plainText.password; - } else if (secretManager != null) { - if (secretManager.secretId == null || secretManager.region == null) { - throw new IllegalArgumentException("secretId and region must be set for aws credential type"); - } - final Map secretMap = this.getSecretValueMap(secretManager.stsRoleArn, secretManager.region, secretManager.secretId); - if (!secretMap.containsKey("username") || !secretMap.containsKey("password")) { - throw new RuntimeException("username or password missing in secret manager."); - } - this.username = secretMap.get("username"); - this.password = secretMap.get("password"); - } else { - throw new IllegalArgumentException("plaintext or secret_manager must be set"); - } - } - - private Map getSecretValueMap(String stsRoleArn, String region, String secretId) { - ObjectMapper objectMapper = new ObjectMapper(); - try { - final String secretValue = SecretManagerHelper.getSecretValue(stsRoleArn, region, secretId); - return objectMapper.readValue(secretValue, new TypeReference<>() {}); - } catch (Exception e) { - throw new RuntimeException("Failed to get credentials.", e); - } - } - - public String getUsername() { - return username; - } - - public String getPassword() { - return password; - } - - public static class PlainText { - private String username; - private String password; - - @JsonCreator - public PlainText(@JsonProperty("username") String username, - @JsonProperty("password") String password) { - this.username = username; - this.password = password; - } - } - - public static class SecretManager { - private String region; - private String secretId; - private String stsRoleArn; - - @JsonCreator - public SecretManager(@JsonProperty("sts_role_arn") String stsRoleArn, - @JsonProperty("region") String region, - @JsonProperty("secretId") String secretId) { - this.stsRoleArn = stsRoleArn; - this.region = region; - this.secretId = secretId; - } - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/MongoDBConfig.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/MongoDBConfig.java deleted file mode 100644 index 02da26c9d4..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/MongoDBConfig.java +++ /dev/null @@ -1,173 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.configuration; - -import com.fasterxml.jackson.annotation.JsonCreator; -import com.fasterxml.jackson.annotation.JsonProperty; -import jakarta.validation.constraints.NotNull; -import org.opensearch.dataprepper.plugins.kafkaconnect.util.Connector; - -import java.util.ArrayList; -import java.util.Arrays; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.stream.Collectors; - -public class MongoDBConfig extends ConnectorConfig { - public static final String CONNECTOR_CLASS = "io.debezium.connector.mongodb.MongoDbConnector"; - private static final String MONGODB_CONNECTION_STRING_FORMAT = "mongodb://%s:%s/?replicaSet=rs0&directConnection=true"; - private static final String DEFAULT_PORT = "27017"; - private static final String DEFAULT_SNAPSHOT_MODE = "never"; - private static final Boolean SSL_ENABLED = false; - private static final Boolean SSL_INVALID_HOST_ALLOWED = false; - private static final String DEFAULT_SNAPSHOT_FETCH_SIZE = "1000"; - @JsonProperty("hostname") - @NotNull - private String hostname; - @JsonProperty("port") - private String port = DEFAULT_PORT; - @JsonProperty("credentials") - private CredentialsConfig credentialsConfig; - @JsonProperty("ingestion_mode") - private IngestionMode ingestionMode = IngestionMode.EXPORT_STREAM; - @JsonProperty("export_config") - private ExportConfig exportConfig = new ExportConfig(); - @JsonProperty("snapshot_fetch_size") - private String snapshotFetchSize = DEFAULT_SNAPSHOT_FETCH_SIZE; - @JsonProperty("collections") - private List collections = new ArrayList<>(); - @JsonProperty("ssl") - private Boolean ssl = SSL_ENABLED; - @JsonProperty("ssl_invalid_host_allowed") - private Boolean sslInvalidHostAllowed = SSL_INVALID_HOST_ALLOWED; - - @Override - public List buildConnectors() { - return collections.stream().map(collection -> { - final String connectorName = collection.getTopicPrefix() + "." + collection.getCollectionName(); - final Map config = buildConfig(collection); - return new Connector(connectorName, config, this.forceUpdate); - }).collect(Collectors.toList()); - } - - public IngestionMode getIngestionMode() { - return this.ingestionMode; - } - - public CredentialsConfig getCredentialsConfig() { - return this.credentialsConfig; - } - - public String getHostname() { - return this.hostname; - } - - public String getPort() { - return this.port; - } - - public Boolean getSSLEnabled() { - return this.ssl; - } - - public Boolean getSSLInvalidHostAllowed() { - return this.sslInvalidHostAllowed; - } - - public List getCollections() { - return this.collections; - } - - public ExportConfig getExportConfig() { - return this.exportConfig; - } - - private Map buildConfig(final CollectionConfig collection) { - Map config = new HashMap<>(); - config.put("connector.class", CONNECTOR_CLASS); - config.put("mongodb.connection.string", String.format(MONGODB_CONNECTION_STRING_FORMAT, hostname, port)); - config.put("mongodb.user", credentialsConfig.getUsername()); - config.put("mongodb.password", credentialsConfig.getPassword()); - config.put("snapshot.mode", DEFAULT_SNAPSHOT_MODE); - config.put("snapshot.fetch.size", snapshotFetchSize); - config.put("topic.prefix", collection.getTopicPrefix()); - config.put("collection.include.list", collection.getCollectionName()); - config.put("mongodb.ssl.enabled", ssl.toString()); - config.put("mongodb.ssl.invalid.hostname.allowed", sslInvalidHostAllowed.toString()); - // Non-configurable properties used to transform CDC data before sending to Kafka. - config.put("transforms", "unwrap"); - config.put("transforms.unwrap.type", "io.debezium.connector.mongodb.transforms.ExtractNewDocumentState"); - config.put("transforms.unwrap.drop.tombstones", "true"); - config.put("transforms.unwrap.delete.handling.mode", "rewrite"); - config.put("transforms.unwrap.add.fields", "op,rs,collection,source.ts_ms,source.db,source.snapshot,ts_ms"); - return config; - } - - public enum IngestionMode { - EXPORT_STREAM("export_stream"), - EXPORT("export"), - STREAM("stream"); - - private static final Map OPTIONS_MAP = Arrays.stream(IngestionMode.values()) - .collect(Collectors.toMap( - value -> value.type, - value -> value - )); - - private final String type; - - IngestionMode(final String type) { - this.type = type; - } - - @JsonCreator - public static IngestionMode fromTypeValue(final String type) { - return OPTIONS_MAP.get(type.toLowerCase()); - } - } - - public static class CollectionConfig { - @JsonProperty("topic_prefix") - @NotNull - private String topicPrefix; - - @JsonProperty("collection") - @NotNull - private String collectionName; - - public String getCollectionName() { - return collectionName; - } - - public String getTopicPrefix() { - return topicPrefix; - } - } - - public static class ExportConfig { - private static int DEFAULT_ITEMS_PER_PARTITION = 4000; - private static String DEFAULT_READ_PREFERENCE = "secondaryPreferred"; - @JsonProperty("acknowledgments") - private Boolean acknowledgments = false; - @JsonProperty("items_per_partition") - private Integer itemsPerPartition = DEFAULT_ITEMS_PER_PARTITION; - @JsonProperty("read_preference") - private String readPreference = DEFAULT_READ_PREFERENCE; - - public boolean getAcknowledgements() { - return this.acknowledgments; - } - - public Integer getItemsPerPartition() { - return this.itemsPerPartition; - } - - public String getReadPreference() { - return this.readPreference; - } - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/MySQLConfig.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/MySQLConfig.java deleted file mode 100644 index ea9d0f9f95..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/MySQLConfig.java +++ /dev/null @@ -1,109 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.configuration; - -import com.fasterxml.jackson.annotation.JsonProperty; -import jakarta.validation.constraints.NotNull; -import org.apache.kafka.connect.runtime.WorkerConfig; -import org.opensearch.dataprepper.plugins.kafkaconnect.util.Connector; - -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Properties; -import java.util.stream.Collectors; - -public class MySQLConfig extends ConnectorConfig { - public static final String CONNECTOR_CLASS = "io.debezium.connector.mysql.MySqlConnector"; - private static final String SCHEMA_HISTORY_PRODUCER_PREFIX = "schema.history.internal.producer."; - private static final String SCHEMA_HISTORY_CONSUMER_PREFIX = "schema.history.internal.consumer."; - private static final String TOPIC_DEFAULT_PARTITIONS = "10"; - private static final String TOPIC_DEFAULT_REPLICATION_FACTOR = "-1"; - private static final String SCHEMA_HISTORY = "schemahistory"; - private static final String DEFAULT_SNAPSHOT_MODE = "initial"; - private static final String DEFAULT_PORT = "3306"; - - @JsonProperty("hostname") - @NotNull - private String hostname; - @JsonProperty("port") - private String port = DEFAULT_PORT; - @JsonProperty("credentials") - private CredentialsConfig credentialsConfig; - @JsonProperty("snapshot_mode") - private String snapshotMode = DEFAULT_SNAPSHOT_MODE; - @JsonProperty("tables") - private List tables = new ArrayList<>(); - - @Override - public List buildConnectors() { - return tables.stream().map(table -> { - final String connectorName = table.getTopicPrefix() + "." + table.getTableName(); - final Map config = buildConfig(table, connectorName); - return new Connector(connectorName, config, this.forceUpdate); - }).collect(Collectors.toList()); - } - - private Map buildConfig(final TableConfig table, final String connectorName) { - int databaseServerId = Math.abs(connectorName.hashCode()); - final Map config = new HashMap<>(); - final Properties authProperties = this.getAuthProperties(); - if (authProperties != null) { - authProperties.forEach((k, v) -> { - if (k == WorkerConfig.BOOTSTRAP_SERVERS_CONFIG) { - this.setBootstrapServers(v.toString()); - return; - } - if (v instanceof Class) { - config.put(SCHEMA_HISTORY_PRODUCER_PREFIX + k, ((Class) v).getName()); - config.put(SCHEMA_HISTORY_CONSUMER_PREFIX + k, ((Class) v).getName()); - return; - } - config.put(SCHEMA_HISTORY_PRODUCER_PREFIX + k, v.toString()); - config.put(SCHEMA_HISTORY_CONSUMER_PREFIX + k, v.toString()); - }); - } - config.put("topic.creation.default.partitions", TOPIC_DEFAULT_PARTITIONS); - config.put("topic.creation.default.replication.factor", TOPIC_DEFAULT_REPLICATION_FACTOR); - config.put("connector.class", CONNECTOR_CLASS); - config.put("database.hostname", hostname); - config.put("database.port", port); - config.put("database.user", credentialsConfig.getUsername()); - config.put("database.password", credentialsConfig.getPassword()); - config.put("snapshot.mode", snapshotMode); - config.put("topic.prefix", table.getTopicPrefix()); - config.put("table.include.list", table.getTableName()); - config.put("schema.history.internal.kafka.bootstrap.servers", this.getBootstrapServers()); - config.put("schema.history.internal.kafka.topic", String.join(".", List.of(table.getTopicPrefix(), table.getTableName(), SCHEMA_HISTORY))); - config.put("database.server.id", Integer.toString(databaseServerId)); - // Non-configurable properties used to transform CDC data before sending to Kafka. - config.put("transforms", "unwrap"); - config.put("transforms.unwrap.type", "io.debezium.transforms.ExtractNewRecordState"); - config.put("transforms.unwrap.drop.tombstones", "true"); - config.put("transforms.unwrap.delete.handling.mode", "rewrite"); - config.put("transforms.unwrap.add.fields", "op,table,source.ts_ms,source.db,source.snapshot,ts_ms"); - return config; - } - - private static class TableConfig { - @JsonProperty("topic_prefix") - @NotNull - private String topicPrefix; - - @JsonProperty("table") - @NotNull - private String tableName; - - public String getTableName() { - return tableName; - } - - public String getTopicPrefix() { - return topicPrefix; - } - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/PostgreSQLConfig.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/PostgreSQLConfig.java deleted file mode 100644 index 205cfcd7a4..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/PostgreSQLConfig.java +++ /dev/null @@ -1,124 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.configuration; - -import com.fasterxml.jackson.annotation.JsonCreator; -import com.fasterxml.jackson.annotation.JsonProperty; -import jakarta.validation.constraints.NotNull; -import org.opensearch.dataprepper.plugins.kafkaconnect.util.Connector; - -import java.util.ArrayList; -import java.util.Arrays; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.stream.Collectors; - -public class PostgreSQLConfig extends ConnectorConfig { - public static final String CONNECTOR_CLASS = "io.debezium.connector.postgresql.PostgresConnector"; - private static final String TOPIC_DEFAULT_PARTITIONS = "10"; - private static final String TOPIC_DEFAULT_REPLICATION_FACTOR = "-1"; - private static final String DEFAULT_PORT = "5432"; - private static final String DEFAULT_SNAPSHOT_MODE = "initial"; - private static final PluginName DEFAULT_DECODING_PLUGIN = PluginName.PGOUTPUT; // default plugin for Aurora PostgreSQL - @JsonProperty("hostname") - @NotNull - private String hostname; - @JsonProperty("port") - private String port = DEFAULT_PORT; - /** - * The name of the PostgreSQL logical decoding plug-in installed on the PostgreSQL server. - * Supported values are decoderbufs, and pgoutput. - */ - @JsonProperty("plugin_name") - private PluginName pluginName = DEFAULT_DECODING_PLUGIN; - @JsonProperty("credentials") - private CredentialsConfig credentialsConfig; - @JsonProperty("snapshot_mode") - private String snapshotMode = DEFAULT_SNAPSHOT_MODE; - @JsonProperty("tables") - private List tables = new ArrayList<>(); - - @Override - public List buildConnectors() { - return tables.stream().map(table -> { - final String connectorName = table.getTopicPrefix() + "." + table.getTableName(); - final Map config = buildConfig(table); - return new Connector(connectorName, config, this.forceUpdate); - }).collect(Collectors.toList()); - } - - private Map buildConfig(final TableConfig tableName) { - Map config = new HashMap<>(); - config.put("topic.creation.default.partitions", TOPIC_DEFAULT_PARTITIONS); - config.put("topic.creation.default.replication.factor", TOPIC_DEFAULT_REPLICATION_FACTOR); - config.put("connector.class", CONNECTOR_CLASS); - config.put("plugin.name", pluginName.type); - config.put("database.hostname", hostname); - config.put("database.port", port); - config.put("database.user", credentialsConfig.getUsername()); - config.put("database.password", credentialsConfig.getPassword()); - config.put("snapshot.mode", snapshotMode); - config.put("topic.prefix", tableName.getTopicPrefix()); - config.put("database.dbname", tableName.getDatabaseName()); - config.put("table.include.list", tableName.getTableName()); - // Non-configurable properties used to transform CDC data before sending to Kafka. - config.put("transforms", "unwrap"); - config.put("transforms.unwrap.type", "io.debezium.transforms.ExtractNewRecordState"); - config.put("transforms.unwrap.drop.tombstones", "true"); - config.put("transforms.unwrap.delete.handling.mode", "rewrite"); - config.put("transforms.unwrap.add.fields", "op,table,source.ts_ms,source.db,source.snapshot,ts_ms"); - return config; - } - - public enum PluginName { - DECODERBUFS("decoderbufs"), - PGOUTPUT("pgoutput"); - - private static final Map OPTIONS_MAP = Arrays.stream(PostgreSQLConfig.PluginName.values()) - .collect(Collectors.toMap( - value -> value.type, - value -> value - )); - - private final String type; - - PluginName(final String type) { - this.type = type; - } - - @JsonCreator - public static PostgreSQLConfig.PluginName fromTypeValue(final String type) { - return OPTIONS_MAP.get(type.toLowerCase()); - } - } - - private static class TableConfig { - @JsonProperty("database") - @NotNull - private String databaseName; - - @JsonProperty("topic_prefix") - @NotNull - private String topicPrefix; - - @JsonProperty("table") - @NotNull - private String tableName; - - public String getDatabaseName() { - return databaseName; - } - - public String getTableName() { - return tableName; - } - - public String getTopicPrefix() { - return topicPrefix; - } - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/DefaultKafkaConnectConfigSupplier.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/DefaultKafkaConnectConfigSupplier.java deleted file mode 100644 index 39e7f6609a..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/DefaultKafkaConnectConfigSupplier.java +++ /dev/null @@ -1,18 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.extension; - -public class DefaultKafkaConnectConfigSupplier implements KafkaConnectConfigSupplier { - private final KafkaConnectConfig kafkaConnectConfig; - public DefaultKafkaConnectConfigSupplier(KafkaConnectConfig kafkaConnectConfig) { - this.kafkaConnectConfig = kafkaConnectConfig; - } - - @Override - public KafkaConnectConfig getConfig() { - return this.kafkaConnectConfig; - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfig.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfig.java deleted file mode 100644 index fbed48d949..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfig.java +++ /dev/null @@ -1,96 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.extension; - -import com.fasterxml.jackson.annotation.JsonProperty; -import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.AwsConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; -import org.opensearch.dataprepper.plugins.kafka.util.KafkaClusterAuthConfig; - -import java.time.Duration; -import java.util.List; -import java.util.Objects; -import java.util.Properties; - -public class KafkaConnectConfig implements KafkaClusterAuthConfig { - private static final long CONNECTOR_TIMEOUT_MS = 360000L; // 360 seconds - private static final long CONNECT_TIMEOUT_MS = 60000L; // 60 seconds - - @JsonProperty("worker_properties") - private WorkerProperties workerProperties = new WorkerProperties(); - - @JsonProperty("connect_start_timeout") - private Duration connectStartTimeout = Duration.ofMillis(CONNECT_TIMEOUT_MS); - - @JsonProperty("connector_start_timeout") - private Duration connectorStartTimeout = Duration.ofMillis(CONNECTOR_TIMEOUT_MS); - - @JsonProperty("bootstrap_servers") - private List bootstrapServers; - - private AuthConfig authConfig; - private EncryptionConfig encryptionConfig; - private AwsConfig awsConfig; - - public Duration getConnectStartTimeout() { - return connectStartTimeout; - } - - public Duration getConnectorStartTimeout() { - return connectorStartTimeout; - } - - public void setBootstrapServers(final List bootstrapServers) { - this.bootstrapServers = bootstrapServers; - if (Objects.nonNull(bootstrapServers)) { - this.workerProperties.setBootstrapServers(String.join(",", bootstrapServers));; - } - } - - public void setAuthProperties(final Properties authProperties) { - this.workerProperties.setAuthProperties(authProperties); - } - - public void setAuthConfig(AuthConfig authConfig) { - this.authConfig = authConfig; - } - - public void setAwsConfig(AwsConfig awsConfig) { - this.awsConfig = awsConfig; - } - - public void setEncryptionConfig(EncryptionConfig encryptionConfig) { - this.encryptionConfig = encryptionConfig; - } - - public WorkerProperties getWorkerProperties() { - return workerProperties; - } - - @Override - public AwsConfig getAwsConfig() { - return awsConfig; - } - - @Override - public AuthConfig getAuthConfig() { - return authConfig; - } - - @Override - public EncryptionConfig getEncryptionConfig() { - return encryptionConfig; - } - - @Override - public List getBootstrapServers() { - if (Objects.nonNull(bootstrapServers)) { - return bootstrapServers; - } - return null; - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigExtension.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigExtension.java deleted file mode 100644 index 18c6aee682..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigExtension.java +++ /dev/null @@ -1,30 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.extension; - -import org.opensearch.dataprepper.model.annotations.DataPrepperExtensionPlugin; -import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; -import org.opensearch.dataprepper.model.plugin.ExtensionPlugin; -import org.opensearch.dataprepper.model.plugin.ExtensionPoints; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -@DataPrepperExtensionPlugin(modelType = KafkaConnectConfig.class, rootKeyJsonPath = "/kafka_connect_config") -public class KafkaConnectConfigExtension implements ExtensionPlugin { - private static final Logger LOG = LoggerFactory.getLogger(KafkaConnectConfigExtension.class); - private DefaultKafkaConnectConfigSupplier defaultKafkaConnectConfigSupplier; - - @DataPrepperPluginConstructor - public KafkaConnectConfigExtension(final KafkaConnectConfig kafkaConnectConfig) { - this.defaultKafkaConnectConfigSupplier = new DefaultKafkaConnectConfigSupplier(kafkaConnectConfig); - } - - @Override - public void apply(ExtensionPoints extensionPoints) { - LOG.info("Applying Kafka Connect Config Extension."); - extensionPoints.addExtensionProvider(new KafkaConnectConfigProvider(this.defaultKafkaConnectConfigSupplier)); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigProvider.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigProvider.java deleted file mode 100644 index 2e2f5e01f6..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigProvider.java +++ /dev/null @@ -1,27 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.extension; - -import org.opensearch.dataprepper.model.plugin.ExtensionProvider; - -import java.util.Optional; - -public class KafkaConnectConfigProvider implements ExtensionProvider { - private final KafkaConnectConfigSupplier kafkaConnectConfigSupplier; - public KafkaConnectConfigProvider(KafkaConnectConfigSupplier kafkaConnectConfigSupplier) { - this.kafkaConnectConfigSupplier = kafkaConnectConfigSupplier; - } - - @Override - public Optional provideInstance(Context context) { - return Optional.of(this.kafkaConnectConfigSupplier); - } - - @Override - public Class supportedClass() { - return KafkaConnectConfigSupplier.class; - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigSupplier.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigSupplier.java deleted file mode 100644 index c5805378d9..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigSupplier.java +++ /dev/null @@ -1,10 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.extension; - -public interface KafkaConnectConfigSupplier { - KafkaConnectConfig getConfig(); -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/WorkerProperties.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/WorkerProperties.java deleted file mode 100644 index 1de8fff29c..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/WorkerProperties.java +++ /dev/null @@ -1,213 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.extension; - -import com.fasterxml.jackson.annotation.JsonProperty; -import org.apache.kafka.connect.runtime.WorkerConfig; - -import java.time.Duration; -import java.util.HashMap; -import java.util.Map; -import java.util.Properties; - -public class WorkerProperties { - private static final String KEY_CONVERTER = "org.apache.kafka.connect.json.JsonConverter"; - private static final String KEY_CONVERTER_SCHEMAS_ENABLE = "false"; - private static final String VALUE_CONVERTER_SCHEMAS_ENABLE = "false"; - private static final String VALUE_CONVERTER = "org.apache.kafka.connect.json.JsonConverter"; - private static final Integer OFFSET_STORAGE_PARTITIONS = 25; - private static final Long OFFSET_FLUSH_INTERVAL_MS = 60000L; - private static final Long OFFSET_FLUSH_TIMEOUT_MS = 5000L; - private static final Integer STATUS_STORAGE_PARTITIONS = 5; - private static final Long HEARTBEAT_INTERVAL_MS = 3000L; - private static final Long SESSION_TIMEOUT_MS = 30000L; - private static final long CONNECTOR_REBALANCE_DELAY_MS = 300000L; // 300 seconds - private static final String DEFAULT_GROUP_ID = "localGroup"; - private static final String DEFAULT_CLIENT_ID = "localClient"; - private static final String DEFAULT_CONFIG_STORAGE_TOPIC = "config-storage-topic"; - private static final String DEFAULT_OFFSET_STORAGE_TOPIC = "offset-storage-topic"; - private static final String DEFAULT_STATUS_STORAGE_TOPIC = "status-storage-topic"; - private final Integer offsetStorageReplicationFactor = -1; - private final Integer configStorageReplicationFactor = -1; - private final Integer statusStorageReplicationFactor = -1; - private String keyConverter = KEY_CONVERTER; - private String keyConverterSchemasEnable = KEY_CONVERTER_SCHEMAS_ENABLE; - private String valueConverter = VALUE_CONVERTER; - private String valueConverterSchemasEnable = VALUE_CONVERTER_SCHEMAS_ENABLE; - - @JsonProperty("group_id") - private String groupId = DEFAULT_GROUP_ID; - @JsonProperty("config_storage_topic") - private String configStorageTopic = DEFAULT_CONFIG_STORAGE_TOPIC; - @JsonProperty("offset_storage_topic") - private String offsetStorageTopic = DEFAULT_OFFSET_STORAGE_TOPIC; - @JsonProperty("status_storage_topic") - private String statusStorageTopic = DEFAULT_STATUS_STORAGE_TOPIC; - @JsonProperty("client_id") - private String clientId = DEFAULT_CLIENT_ID; - @JsonProperty("offset_storage_partitions") - private Integer offsetStoragePartitions = OFFSET_STORAGE_PARTITIONS; - @JsonProperty("offset_flush_interval") - private Duration offsetFlushInterval = Duration.ofMillis(OFFSET_FLUSH_INTERVAL_MS); - @JsonProperty("offset_flush_timeout") - private Duration offsetFlushTimeout = Duration.ofMillis(OFFSET_FLUSH_TIMEOUT_MS); - @JsonProperty("status_storage_partitions") - private Integer statusStoragePartitions = STATUS_STORAGE_PARTITIONS; - @JsonProperty("heartbeat_interval") - private Duration heartBeatInterval = Duration.ofMillis(HEARTBEAT_INTERVAL_MS); - @JsonProperty("session_timeout") - private Duration sessionTimeout = Duration.ofMillis(SESSION_TIMEOUT_MS); - @JsonProperty("connector_rebalance_max_delay") - private Duration connectorRebalanceDelay = Duration.ofMillis(CONNECTOR_REBALANCE_DELAY_MS); - private String keyConverterSchemaRegistryUrl; - private String valueConverterSchemaRegistryUrl; - private String bootstrapServers; - private Properties authProperties; - - public String getKeyConverter() { - return keyConverter; - } - - public String getKeyConverterSchemasEnable() { - return keyConverterSchemasEnable; - } - - public String getKeyConverterSchemaRegistryUrl() { - return keyConverterSchemaRegistryUrl; - } - - public String getValueConverter() { - return valueConverter; - } - - public String getValueConverterSchemasEnable() { - return valueConverterSchemasEnable; - } - - public String getValueConverterSchemaRegistryUrl() { - return valueConverterSchemaRegistryUrl; - } - - public Integer getOffsetStoragePartitions() { - return offsetStoragePartitions; - } - - public Long getOffsetFlushInterval() { - return offsetFlushInterval.toMillis(); - } - - public Long getOffsetFlushTimeout() { - return offsetFlushTimeout.toMillis(); - } - - public Long getRebalanceMaxDelay() { - return connectorRebalanceDelay.toMillis(); - } - - public Integer getStatusStoragePartitions() { - return statusStoragePartitions; - } - - public Long getHeartBeatInterval() { - return heartBeatInterval.toMillis(); - } - - public Long getSessionTimeout() { - return sessionTimeout.toMillis(); - } - - public String getBootstrapServers() { - return bootstrapServers; - } - - public void setBootstrapServers(final String bootstrapServers) { - this.bootstrapServers = bootstrapServers; - } - - public String getGroupId() { - return groupId; - } - - public String getClientId() { - return clientId; - } - - public String getConfigStorageTopic() { - return configStorageTopic; - } - - public Integer getConfigStorageReplicationFactor() { - return configStorageReplicationFactor; - } - - public String getOffsetStorageTopic() { - return offsetStorageTopic; - } - - public Integer getOffsetStorageReplicationFactor() { - return offsetStorageReplicationFactor; - } - - public String getStatusStorageTopic() { - return statusStorageTopic; - } - - public Integer getStatusStorageReplicationFactor() { - return statusStorageReplicationFactor; - } - - public void setAuthProperties(Properties authProperties) { - this.authProperties = authProperties; - } - - public Map buildKafkaConnectPropertyMap() { - final String producerPrefix = "producer."; - Map workerProps = new HashMap<>(); - if (authProperties != null) { - authProperties.forEach((k, v) -> { - if (k == WorkerConfig.BOOTSTRAP_SERVERS_CONFIG) { - this.setBootstrapServers(v.toString()); - return; - } - if (v instanceof Class) { - workerProps.put(k.toString(), ((Class) v).getName()); - workerProps.put(producerPrefix + k, ((Class) v).getName()); - return; - } - workerProps.put(k.toString(), v.toString()); - workerProps.put(producerPrefix + k, v.toString()); - }); - } - workerProps.put("bootstrap.servers", this.getBootstrapServers()); - workerProps.put("group.id", this.getGroupId()); - workerProps.put("client.id", this.getClientId()); - workerProps.put("offset.storage.topic", this.getOffsetStorageTopic()); - workerProps.put("offset.storage.replication.factor", this.getOffsetStorageReplicationFactor().toString()); - workerProps.put("config.storage.topic", this.getConfigStorageTopic()); - workerProps.put("config.storage.replication.factor", this.getConfigStorageReplicationFactor().toString()); - workerProps.put("status.storage.topic", this.getStatusStorageTopic()); - workerProps.put("status.storage.replication.factor", this.getStatusStorageReplicationFactor().toString()); - workerProps.put("key.converter", this.getKeyConverter()); - workerProps.put("key.converter.schemas.enable", this.getKeyConverterSchemasEnable()); - if (this.getKeyConverterSchemaRegistryUrl() != null) { - workerProps.put("key.converter.schema.registry.url", this.getKeyConverterSchemaRegistryUrl()); - } - workerProps.put("value.converter", this.getValueConverter()); - workerProps.put("value.converter.schemas.enable", this.getValueConverterSchemasEnable()); - if (this.getValueConverterSchemaRegistryUrl() != null) { - workerProps.put("value.converter.schema.registry.url", this.getValueConverterSchemaRegistryUrl()); - } - workerProps.put("offset.storage.partitions", this.getOffsetStoragePartitions().toString()); - workerProps.put("offset.flush.interval.ms", this.getOffsetFlushInterval().toString()); - workerProps.put("offset.flush.timeout.ms", this.getOffsetFlushTimeout().toString()); - workerProps.put("status.storage.partitions", this.getStatusStoragePartitions().toString()); - workerProps.put("heartbeat.interval.ms", this.getHeartBeatInterval().toString()); - workerProps.put("session.timeout.ms", this.getSessionTimeout().toString()); - workerProps.put("scheduled.rebalance.max.delay.ms", this.getRebalanceMaxDelay().toString()); - return workerProps; - } -} - diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/meter/KafkaConnectMetrics.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/meter/KafkaConnectMetrics.java deleted file mode 100644 index db7250c259..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/meter/KafkaConnectMetrics.java +++ /dev/null @@ -1,168 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.meter; - -import io.micrometer.core.instrument.Tag; -import io.micrometer.core.instrument.Tags; -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import javax.management.MBeanServer; -import javax.management.MBeanServerFactory; -import javax.management.MalformedObjectNameException; -import javax.management.ObjectName; -import java.lang.management.ManagementFactory; -import java.util.List; -import java.util.Set; -import java.util.concurrent.Callable; -import java.util.function.BiFunction; - -import static java.util.Collections.emptyList; - -public class KafkaConnectMetrics { - private static final Logger LOG = LoggerFactory.getLogger(KafkaConnectMetrics.class); - private static final String JMX_DOMAIN = "kafka.connect"; - private static final String CONNECT_WORKER_METRICS_NAME = "connect-worker-metrics"; - private static final List CONNECT_WORKER_METRICS_LIST = List.of( - "task-count", - "connector-count", - "connector-startup-attempts-total", - "connector-startup-success-total", - "connector-startup-failure-total", - "task-startup-attempts-total", - "task-startup-success-total", - "task-startup-failure-total" - ); - private static final String SOURCE_TASK_METRICS_NAME = "source-task-metrics"; - private static final List SOURCE_TASK_METRICS_LIST = List.of( - "source-record-write-total", - "source-record-write-rate", - "source-record-poll-total", - "source-record-poll-rate", - "source-record-active-count-max", - "source-record-active-count-avg", - "source-record-active-count" - ); - private static final String CLIENT_ID_KEY = "client-id"; - private static final String CLIENT_ID = "client.id"; - private static final String NODE_ID_KEY = "node-id"; - private static final String NODE_ID = "node.id"; - private static final String CONNECTOR = "connector"; - private static final String TASK = "task"; - - private final PluginMetrics pluginMetrics; - - private final MBeanServer mBeanServer; - - private final Iterable tags; - - - public KafkaConnectMetrics(final PluginMetrics pluginMetrics) { - this(pluginMetrics, emptyList()); - } - - public KafkaConnectMetrics(final PluginMetrics pluginMetrics, - final Iterable tags) { - this(pluginMetrics, getMBeanServer(), tags); - } - - public KafkaConnectMetrics(final PluginMetrics pluginMetrics, - final MBeanServer mBeanServer, - final Iterable tags) { - this.pluginMetrics = pluginMetrics; - this.mBeanServer = mBeanServer; - this.tags = tags; - } - - private static MBeanServer getMBeanServer() { - List mBeanServers = MBeanServerFactory.findMBeanServer(null); - if (!mBeanServers.isEmpty()) { - return mBeanServers.get(0); - } - return ManagementFactory.getPlatformMBeanServer(); - } - - private static String sanitize(String value) { - return value.replaceAll("-", "."); - } - - public void bindConnectMetrics() { - registerMetricsEventually(CONNECT_WORKER_METRICS_NAME, (o, tags) -> { - CONNECT_WORKER_METRICS_LIST.forEach( - (metricName) -> registerFunctionGaugeForObject(o, metricName, tags) - ); - return null; - }); - } - - public void bindConnectorMetrics() { - registerMetricsEventually(SOURCE_TASK_METRICS_NAME, (o, tags) -> { - SOURCE_TASK_METRICS_LIST.forEach( - (metricName) -> registerFunctionGaugeForObject(o, metricName, tags) - ); - return null; - }); - } - - private void registerMetricsEventually(String type, - BiFunction perObject) { - try { - Set objs = mBeanServer.queryNames(new ObjectName(JMX_DOMAIN + ":type=" + type + ",*"), null); - if (!objs.isEmpty()) { - for (ObjectName o : objs) { - perObject.apply(o, Tags.concat(tags, nameTag(o))); - } - } - } catch (MalformedObjectNameException e) { - throw new RuntimeException("Error registering Kafka Connect JMX based metrics", e); - } - } - - private Iterable nameTag(ObjectName name) { - Tags tags = Tags.empty(); - - String clientId = name.getKeyProperty(CLIENT_ID_KEY); - if (clientId != null) { - tags = Tags.concat(tags, CLIENT_ID, clientId); - } - - String nodeId = name.getKeyProperty(NODE_ID_KEY); - if (nodeId != null) { - tags = Tags.concat(tags, NODE_ID, nodeId); - } - - String connectorName = name.getKeyProperty(CONNECTOR); - if (connectorName != null) { - tags = Tags.concat(tags, CONNECTOR, connectorName); - } - - String taskName = name.getKeyProperty(TASK); - if (taskName != null) { - tags = Tags.concat(tags, TASK, taskName); - } - - return tags; - } - - private void registerFunctionGaugeForObject(ObjectName o, String jmxMetricName, Tags allTags) { - pluginMetrics.gaugeWithTags( - sanitize(jmxMetricName), - allTags, - mBeanServer, - s -> safeDouble(() -> s.getAttribute(o, jmxMetricName)) - ); - } - - private double safeDouble(Callable callable) { - try { - if (callable.call() == null) return Double.NaN; - return Double.parseDouble(callable.call().toString()); - } catch (Exception e) { - return Double.NaN; - } - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/KafkaConnectSource.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/KafkaConnectSource.java deleted file mode 100644 index 03370a359a..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/KafkaConnectSource.java +++ /dev/null @@ -1,114 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.source; - -import org.apache.kafka.connect.runtime.WorkerConfig; -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.buffer.Buffer; -import org.opensearch.dataprepper.model.configuration.PipelineDescription; -import org.opensearch.dataprepper.model.record.Record; -import org.opensearch.dataprepper.model.source.Source; -import org.opensearch.dataprepper.plugins.kafka.extension.KafkaClusterConfigSupplier; -import org.opensearch.dataprepper.plugins.kafka.util.KafkaSecurityConfigurer; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.ConnectorConfig; -import org.opensearch.dataprepper.plugins.kafkaconnect.extension.KafkaConnectConfig; -import org.opensearch.dataprepper.plugins.kafkaconnect.extension.KafkaConnectConfigSupplier; -import org.opensearch.dataprepper.plugins.kafkaconnect.extension.WorkerProperties; -import org.opensearch.dataprepper.plugins.kafkaconnect.util.Connector; -import org.opensearch.dataprepper.plugins.kafkaconnect.util.KafkaConnect; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.util.List; -import java.util.Map; -import java.util.Objects; -import java.util.Properties; - -/** - * The abstraction of the kafka connect source. - * The kafka connect and connectors are configured and runs async here. - */ -@SuppressWarnings("deprecation") -public abstract class KafkaConnectSource implements Source> { - private static final Logger LOG = LoggerFactory.getLogger(KafkaConnectSource.class); - public final ConnectorConfig connectorConfig; - private final String pipelineName; - private KafkaConnectConfig kafkaConnectConfig; - private KafkaConnect kafkaConnect; - - public KafkaConnectSource(final ConnectorConfig connectorConfig, - final PluginMetrics pluginMetrics, - final PipelineDescription pipelineDescription, - final KafkaClusterConfigSupplier kafkaClusterConfigSupplier, - final KafkaConnectConfigSupplier kafkaConnectConfigSupplier) { - this.connectorConfig = connectorConfig; - this.pipelineName = pipelineDescription.getPipelineName(); - if (shouldStartKafkaConnect()) { - if (kafkaClusterConfigSupplier == null || kafkaConnectConfigSupplier == null) { - throw new IllegalArgumentException("Extensions: KafkaClusterConfig and KafkaConnectConfig cannot be null"); - } - this.kafkaConnectConfig = kafkaConnectConfigSupplier.getConfig(); - this.updateConfig(kafkaClusterConfigSupplier); - this.kafkaConnect = KafkaConnect.getPipelineInstance( - pipelineName, - pluginMetrics, - kafkaConnectConfig.getConnectStartTimeout(), - kafkaConnectConfig.getConnectorStartTimeout()); - } - } - - @Override - public void start(Buffer> buffer) { - if (shouldStartKafkaConnect()) { - LOG.info("Starting Kafka Connect Source for pipeline: {}", pipelineName); - // Please make sure buildWokerProperties is always first to execute. - final WorkerProperties workerProperties = this.kafkaConnectConfig.getWorkerProperties(); - Map workerProps = workerProperties.buildKafkaConnectPropertyMap(); - if (workerProps.get(WorkerConfig.BOOTSTRAP_SERVERS_CONFIG) == null || workerProps.get(WorkerConfig.BOOTSTRAP_SERVERS_CONFIG).isEmpty()) { - throw new IllegalArgumentException("Bootstrap Servers cannot be null or empty"); - } - final List connectors = this.connectorConfig.buildConnectors(); - kafkaConnect.addConnectors(connectors); - kafkaConnect.initialize(workerProps); - kafkaConnect.start(); - } - } - - @Override - public void stop() { - if (shouldStartKafkaConnect()) { - LOG.info("Stopping Kafka Connect Source for pipeline: {}", pipelineName); - kafkaConnect.stop(); - } - } - - public boolean shouldStartKafkaConnect() { - return true; - } - - private void updateConfig(final KafkaClusterConfigSupplier kafkaClusterConfigSupplier) { - if (kafkaConnectConfig.getBootstrapServers() == null) { - this.kafkaConnectConfig.setBootstrapServers(kafkaClusterConfigSupplier.getBootStrapServers()); - } - if (kafkaConnectConfig.getAuthConfig() == null) { - kafkaConnectConfig.setAuthConfig(kafkaClusterConfigSupplier.getAuthConfig()); - } - if (kafkaConnectConfig.getAwsConfig() == null) { - kafkaConnectConfig.setAwsConfig(kafkaClusterConfigSupplier.getAwsConfig()); - } - if (kafkaConnectConfig.getEncryptionConfig() == null) { - kafkaConnectConfig.setEncryptionConfig(kafkaClusterConfigSupplier.getEncryptionConfig()); - } - Properties authProperties = new Properties(); - KafkaSecurityConfigurer.setAuthProperties(authProperties, kafkaConnectConfig, LOG); - this.kafkaConnectConfig.setAuthProperties(authProperties); - // Update Connector Config - if (Objects.nonNull(kafkaConnectConfig.getBootstrapServers())) { - this.connectorConfig.setBootstrapServers(String.join(",", kafkaConnectConfig.getBootstrapServers())); - } - this.connectorConfig.setAuthProperties(authProperties); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/MongoDBSource.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/MongoDBSource.java deleted file mode 100644 index 2099bd1e2c..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/MongoDBSource.java +++ /dev/null @@ -1,127 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.source; - -import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; -import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; -import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; -import org.opensearch.dataprepper.model.buffer.Buffer; -import org.opensearch.dataprepper.model.codec.ByteDecoder; -import org.opensearch.dataprepper.model.codec.JsonDecoder; -import org.opensearch.dataprepper.model.configuration.PipelineDescription; -import org.opensearch.dataprepper.model.record.Record; -import org.opensearch.dataprepper.model.source.Source; -import org.opensearch.dataprepper.model.source.coordinator.SourceCoordinator; -import org.opensearch.dataprepper.model.source.coordinator.UsesSourceCoordination; -import org.opensearch.dataprepper.plugins.kafka.extension.KafkaClusterConfigSupplier; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.MongoDBConfig; -import org.opensearch.dataprepper.plugins.kafkaconnect.extension.KafkaConnectConfigSupplier; -import org.opensearch.dataprepper.plugins.kafkaconnect.source.mongoDB.MongoDBService; -import org.opensearch.dataprepper.plugins.kafkaconnect.source.mongoDB.MongoDBSnapshotProgressState; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.util.List; -import java.util.Objects; - -/** - * The starting point of the mysql source which ingest CDC data using Kafka Connect and Debezium Connector. - */ -@SuppressWarnings("deprecation") -@DataPrepperPlugin(name = "mongodb", pluginType = Source.class, pluginConfigurationType = MongoDBConfig.class) -public class MongoDBSource extends KafkaConnectSource implements UsesSourceCoordination { - private static final Logger LOG = LoggerFactory.getLogger(MongoDBSource.class); - private static final String COLLECTION_SPLITTER = "\\."; - - private final AwsCredentialsSupplier awsCredentialsSupplier; - - private final PluginMetrics pluginMetrics; - - private final AcknowledgementSetManager acknowledgementSetManager; - - private MongoDBService mongoDBService; - - private SourceCoordinator sourceCoordinator; - - private ByteDecoder byteDecoder; - - @DataPrepperPluginConstructor - public MongoDBSource(final MongoDBConfig mongoDBConfig, - final PluginMetrics pluginMetrics, - final PipelineDescription pipelineDescription, - final AcknowledgementSetManager acknowledgementSetManager, - final AwsCredentialsSupplier awsCredentialsSupplier, - final KafkaClusterConfigSupplier kafkaClusterConfigSupplier, - final KafkaConnectConfigSupplier kafkaConnectConfigSupplier) { - super(mongoDBConfig, pluginMetrics, pipelineDescription, kafkaClusterConfigSupplier, kafkaConnectConfigSupplier); - this.pluginMetrics = pluginMetrics; - this.acknowledgementSetManager = acknowledgementSetManager; - this.awsCredentialsSupplier = awsCredentialsSupplier; - this.byteDecoder = new JsonDecoder(); - this.validateCollections(); - } - - @Override - public void start(Buffer> buffer) { - super.start(buffer); - if (shouldStartInitialLoad()) { - LOG.info("Starting initial load"); - this.mongoDBService = MongoDBService.create((MongoDBConfig) this.connectorConfig, sourceCoordinator, buffer, acknowledgementSetManager, pluginMetrics); - this.mongoDBService.start(); - } - } - - @Override - public void stop(){ - super.stop(); - if (shouldStartInitialLoad() && Objects.nonNull(mongoDBService) && Objects.nonNull(sourceCoordinator)) { - LOG.info("Stopping initial load"); - mongoDBService.stop(); - sourceCoordinator.giveUpPartitions(); - } - } - - @Override - public void setSourceCoordinator(final SourceCoordinator sourceCoordinator) { - this.sourceCoordinator = (SourceCoordinator) sourceCoordinator; - } - - @Override - public Class getPartitionProgressStateClass() { - return MongoDBSnapshotProgressState.class; - } - - @Override - public ByteDecoder getDecoder() { - return byteDecoder; - } - - @Override - public boolean shouldStartKafkaConnect() { - final MongoDBConfig mongoDBConfig = (MongoDBConfig) this.connectorConfig; - return mongoDBConfig.getIngestionMode() == MongoDBConfig.IngestionMode.EXPORT_STREAM - || mongoDBConfig.getIngestionMode() == MongoDBConfig.IngestionMode.STREAM; - } - - private boolean shouldStartInitialLoad() { - final MongoDBConfig mongoDBConfig = (MongoDBConfig) this.connectorConfig; - return mongoDBConfig.getIngestionMode() == MongoDBConfig.IngestionMode.EXPORT_STREAM - || mongoDBConfig.getIngestionMode() == MongoDBConfig.IngestionMode.EXPORT; - } - - private void validateCollections() { - MongoDBConfig config = (MongoDBConfig) this.connectorConfig; - List collectionConfigs = config.getCollections(); - collectionConfigs.forEach(collectionConfig -> { - List collection = List.of(collectionConfig.getCollectionName().split(COLLECTION_SPLITTER)); - if (collection.size() < 2) { - throw new IllegalArgumentException("Invalid Collection Name. Must be in db.collection format"); - } - }); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/MySQLSource.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/MySQLSource.java deleted file mode 100644 index 5dbcd16f4e..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/MySQLSource.java +++ /dev/null @@ -1,32 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.source; - -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; -import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; -import org.opensearch.dataprepper.model.configuration.PipelineDescription; -import org.opensearch.dataprepper.model.source.Source; -import org.opensearch.dataprepper.plugins.kafka.extension.KafkaClusterConfigSupplier; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.MySQLConfig; -import org.opensearch.dataprepper.plugins.kafkaconnect.extension.KafkaConnectConfigSupplier; - -/** - * The starting point of the mysql source which ingest CDC data using Kafka Connect and Debezium Connector. - */ -@SuppressWarnings("deprecation") -@DataPrepperPlugin(name = "mysql", pluginType = Source.class, pluginConfigurationType = MySQLConfig.class) -public class MySQLSource extends KafkaConnectSource { - - @DataPrepperPluginConstructor - public MySQLSource(final MySQLConfig mySQLConfig, - final PluginMetrics pluginMetrics, - final PipelineDescription pipelineDescription, - final KafkaClusterConfigSupplier kafkaClusterConfigSupplier, - final KafkaConnectConfigSupplier kafkaConnectConfigSupplier) { - super(mySQLConfig, pluginMetrics, pipelineDescription, kafkaClusterConfigSupplier, kafkaConnectConfigSupplier); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/PostgreSQLSource.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/PostgreSQLSource.java deleted file mode 100644 index 7e7c24d08c..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/PostgreSQLSource.java +++ /dev/null @@ -1,32 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.source; - -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; -import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; -import org.opensearch.dataprepper.model.configuration.PipelineDescription; -import org.opensearch.dataprepper.model.source.Source; -import org.opensearch.dataprepper.plugins.kafka.extension.KafkaClusterConfigSupplier; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.PostgreSQLConfig; -import org.opensearch.dataprepper.plugins.kafkaconnect.extension.KafkaConnectConfigSupplier; - -/** - * The starting point of the mysql source which ingest CDC data using Kafka Connect and Debezium Connector. - */ -@SuppressWarnings("deprecation") -@DataPrepperPlugin(name = "postgresql", pluginType = Source.class, pluginConfigurationType = PostgreSQLConfig.class) -public class PostgreSQLSource extends KafkaConnectSource { - - @DataPrepperPluginConstructor - public PostgreSQLSource(final PostgreSQLConfig postgreSQLConfig, - final PluginMetrics pluginMetrics, - final PipelineDescription pipelineDescription, - final KafkaClusterConfigSupplier kafkaClusterConfigSupplier, - final KafkaConnectConfigSupplier kafkaConnectConfigSupplier) { - super(postgreSQLConfig, pluginMetrics, pipelineDescription, kafkaClusterConfigSupplier, kafkaConnectConfigSupplier); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBHelper.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBHelper.java deleted file mode 100644 index 96af2c7663..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBHelper.java +++ /dev/null @@ -1,130 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.source.mongoDB; - -import com.mongodb.client.MongoClient; -import com.mongodb.client.MongoClients; -import org.bson.Document; -import org.bson.conversions.Bson; -import org.bson.types.BSONTimestamp; -import org.bson.types.Binary; -import org.bson.types.Code; -import org.bson.types.Decimal128; -import org.bson.types.ObjectId; -import org.bson.types.Symbol; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.MongoDBConfig; - -import static com.mongodb.client.model.Filters.and; -import static com.mongodb.client.model.Filters.gte; -import static com.mongodb.client.model.Filters.lte; - -public class MongoDBHelper { - private static final String MONGO_CONNECTION_STRING_TEMPLATE = "mongodb://%s:%s@%s:%s/?replicaSet=rs0&directConnection=true&readpreference=%s&ssl=%s&tlsAllowInvalidHostnames=%s"; - private static final String BINARY_PARTITION_FORMAT = "%s-%s"; - private static final String BINARY_PARTITION_SPLITTER = "-"; - private static final String TIMESTAMP_PARTITION_FORMAT = "%s-%s"; - private static final String TIMESTAMP_PARTITION_SPLITTER = "-"; - - public static MongoClient getMongoClient(final MongoDBConfig mongoDBConfig) { - String username = mongoDBConfig.getCredentialsConfig().getUsername(); - String password = mongoDBConfig.getCredentialsConfig().getPassword(); - String hostname = mongoDBConfig.getHostname(); - String port = mongoDBConfig.getPort(); - String ssl = mongoDBConfig.getSSLEnabled().toString(); - String invalidHostAllowed = mongoDBConfig.getSSLInvalidHostAllowed().toString(); - String readPreference = mongoDBConfig.getExportConfig().getReadPreference(); - String connectionString = String.format(MONGO_CONNECTION_STRING_TEMPLATE, username, password, hostname, port, readPreference, ssl, invalidHostAllowed); - - return MongoClients.create(connectionString); - } - - public static String getPartitionStringFromMongoDBId(Object id, String className) { - switch (className) { - case "org.bson.Document": - return ((Document) id).toJson(); - case "org.bson.types.Binary": - final byte type = ((Binary) id).getType(); - final byte[] data = ((Binary) id).getData(); - String typeString = String.valueOf((int) type); - String dataString = new String(data); - return String.format(BINARY_PARTITION_FORMAT, typeString, dataString); - case "org.bson.types.BSONTimestamp": - final int inc = ((BSONTimestamp) id).getInc(); - final int time = ((BSONTimestamp) id).getTime(); - return String.format(TIMESTAMP_PARTITION_FORMAT, inc, time); - case "org.bson.types.code": - return ((Code) id).getCode(); - default: - return id.toString(); - } - } - - public static Bson buildAndQuery(String gte, String lte, String className) { - switch (className) { - case "java.lang.Integer": - return and( - gte("_id", Integer.parseInt(gte)), - lte("_id", Integer.parseInt(lte)) - ); - case "java.lang.Double": - return and( - gte("_id", Double.parseDouble(gte)), - lte("_id", Double.parseDouble(lte)) - ); - case "java.lang.String": - return and( - gte("_id", gte), - lte("_id", lte) - ); - case "java.lang.Long": - return and( - gte("_id", Long.parseLong(gte)), - lte("_id", Long.parseLong(lte)) - ); - case "org.bson.types.ObjectId": - return and( - gte("_id", new ObjectId(gte)), - lte("_id", new ObjectId(lte)) - ); - case "org.bson.types.Decimal128": - return and( - gte("_id", Decimal128.parse(gte)), - lte("_id", Decimal128.parse(lte)) - ); - case "org.bson.types.Binary": - String[] gteString = gte.split(BINARY_PARTITION_SPLITTER, 2); - String[] lteString = lte.split(BINARY_PARTITION_SPLITTER, 2); - return and( - gte("_id", new Binary(Byte.parseByte(gteString[0]), gteString[1].getBytes())), - lte("_id", new Binary(Byte.parseByte(lteString[0]), lteString[1].getBytes())) - ); - case "org.bson.types.BSONTimestamp": - String[] gteTimestampString = gte.split(TIMESTAMP_PARTITION_SPLITTER, 2); - String[] lteTimestampString = lte.split(TIMESTAMP_PARTITION_SPLITTER, 2); - return and( - gte("_id", new BSONTimestamp(Integer.parseInt(gteTimestampString[0]), Integer.parseInt(gteTimestampString[1]))), - lte("_id", new BSONTimestamp(Integer.parseInt(lteTimestampString[0]), Integer.parseInt(lteTimestampString[1]))) - ); - case "org.bson.types.code": - return and( - gte("_id", new Code(gte)), - lte("_id", new Code(lte)) - ); - case "org.bson.types.Symbol": - return and( - gte("_id", new Symbol(gte)), - lte("_id", new Symbol(lte)) - ); - case "org.bson.Document": - return and( - gte("_id", Document.parse(gte)), - lte("_id", Document.parse(lte)) - ); - default: - throw new RuntimeException("Unexpected _id class supported: " + className); - } - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBPartitionCreationSupplier.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBPartitionCreationSupplier.java deleted file mode 100644 index e585bb8348..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBPartitionCreationSupplier.java +++ /dev/null @@ -1,131 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.source.mongoDB; - -import com.mongodb.client.FindIterable; -import com.mongodb.client.MongoClient; -import com.mongodb.client.MongoCollection; -import com.mongodb.client.MongoCursor; -import com.mongodb.client.MongoDatabase; -import com.mongodb.client.model.Filters; -import org.bson.Document; -import org.opensearch.dataprepper.model.source.coordinator.PartitionIdentifier; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.MongoDBConfig; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.time.Instant; -import java.util.ArrayList; -import java.util.Collections; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.function.Function; -import java.util.stream.Collectors; - -public class MongoDBPartitionCreationSupplier implements Function, List> { - public static final String GLOBAL_STATE_PARTITIONED_COLLECTION_KEY = "partitionedCollections"; - private static final Logger LOG = LoggerFactory.getLogger(MongoDBPartitionCreationSupplier.class); - private static final String MONGODB_PARTITION_KEY_FORMAT = "%s|%s|%s|%s"; // partition format: ||| - private static final String COLLECTION_SPLITTER = "\\."; - - private final MongoDBConfig mongoDBConfig; - - public MongoDBPartitionCreationSupplier(final MongoDBConfig mongoDBConfig) { - this.mongoDBConfig = mongoDBConfig; - } - - @Override - public List apply(final Map globalStateMap) { - Map partitionedCollections = (Map) globalStateMap.getOrDefault(GLOBAL_STATE_PARTITIONED_COLLECTION_KEY, new HashMap<>()); - List collectionsToInitPartitions = this.getCollectionsToInitPartitions(mongoDBConfig, partitionedCollections); - - if (collectionsToInitPartitions.isEmpty()) { - return Collections.emptyList(); - } - - final List allPartitionIdentifiers = collectionsToInitPartitions - .parallelStream() - .flatMap(collectionName -> { - List partitions = this.buildPartitions(collectionName); - partitionedCollections.put(collectionName, Instant.now().toEpochMilli()); - return partitions.stream(); - }) - .collect(Collectors.toList()); - - globalStateMap.put(GLOBAL_STATE_PARTITIONED_COLLECTION_KEY, partitionedCollections); - return allPartitionIdentifiers; - } - - private List getCollectionsToInitPartitions(final MongoDBConfig mongoDBConfig, - final Map partitionedCollections) { - return mongoDBConfig.getCollections() - .stream() - .map(MongoDBConfig.CollectionConfig::getCollectionName) - .filter(collectionName -> !partitionedCollections.containsKey(collectionName)) - .collect(Collectors.toList()); - } - - private List buildPartitions(final String collectionName) { - List collectionPartitions = new ArrayList<>(); - List collection = List.of(collectionName.split(COLLECTION_SPLITTER)); - if (collection.size() < 2) { - throw new IllegalArgumentException("Invalid Collection Name. Must be in db.collection format"); - } - try (MongoClient mongoClient = MongoDBHelper.getMongoClient(mongoDBConfig)) { - MongoDatabase db = mongoClient.getDatabase(collection.get(0)); - MongoCollection col = db.getCollection(collection.get(1)); - int chunkSize = this.mongoDBConfig.getExportConfig().getItemsPerPartition(); - FindIterable startIterable = col.find() - .projection(new Document("_id", 1)) - .sort(new Document("_id", 1)) - .limit(1); - while (true) { - try (MongoCursor startCursor = startIterable.iterator()) { - if (!startCursor.hasNext()) { - break; - } - Document startDoc = startCursor.next(); - Object gteValue = startDoc.get("_id"); - String className = gteValue.getClass().getName(); - - // Get end doc - Document endDoc = startIterable.skip(chunkSize - 1).limit(1).first(); - if (endDoc == null) { - // this means we have reached the end of the doc - endDoc = col.find() - .projection(new Document("_id", 1)) - .sort(new Document("_id", -1)) - .limit(1) - .first(); - } - if (endDoc == null) { - break; - } - - Object lteValue = endDoc.get("_id"); - String gteValueString = MongoDBHelper.getPartitionStringFromMongoDBId(gteValue, className); - String lteValueString = MongoDBHelper.getPartitionStringFromMongoDBId(lteValue, className); - LOG.info("Chunk of " + collectionName + ": {gte: " + gteValueString + ", lte: " + lteValueString + "}"); - collectionPartitions.add( - PartitionIdentifier - .builder() - .withPartitionKey(String.format(MONGODB_PARTITION_KEY_FORMAT, collectionName, gteValueString, lteValueString, className)) - .build()); - - startIterable = col.find(Filters.gt("_id", lteValue)) - .projection(new Document("_id", 1)) - .sort(new Document("_id", 1)) - .limit(1); - } catch (Exception e) { - LOG.error("Failed to read start cursor when build partitions", e); - throw new RuntimeException(e); - } - } - } - return collectionPartitions; - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBService.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBService.java deleted file mode 100644 index 07dc439d53..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBService.java +++ /dev/null @@ -1,96 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.source.mongoDB; - -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; -import org.opensearch.dataprepper.model.buffer.Buffer; -import org.opensearch.dataprepper.model.record.Record; -import org.opensearch.dataprepper.model.source.coordinator.SourceCoordinator; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.MongoDBConfig; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.time.Duration; -import java.util.concurrent.Executors; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.ScheduledFuture; -import java.util.concurrent.TimeUnit; - -public class MongoDBService { - private static final Logger LOG = LoggerFactory.getLogger(MongoDBService.class); - private static final Duration EXECUTOR_SERVICE_SHUTDOWN_TIMEOUT = Duration.ofSeconds(30); - private final PluginMetrics pluginMetrics; - private final MongoDBConfig mongoDBConfig; - private final Buffer> buffer; - private final MongoDBPartitionCreationSupplier mongoDBPartitionCreationSupplier; - private final ScheduledExecutorService scheduledExecutorService; - private final SourceCoordinator sourceCoordinator; - private final AcknowledgementSetManager acknowledgementSetManager; - private MongoDBSnapshotWorker snapshotWorker; - private ScheduledFuture snapshotWorkerFuture; - - - private MongoDBService( - final MongoDBConfig mongoDBConfig, - final SourceCoordinator sourceCoordinator, - final Buffer> buffer, - final ScheduledExecutorService scheduledExecutorService, - final AcknowledgementSetManager acknowledgementSetManager, - final PluginMetrics pluginMetrics) { - this.pluginMetrics = pluginMetrics; - this.mongoDBConfig = mongoDBConfig; - this.buffer = buffer; - this.scheduledExecutorService = scheduledExecutorService; - this.acknowledgementSetManager = acknowledgementSetManager; - this.sourceCoordinator = sourceCoordinator; - this.sourceCoordinator.initialize(); - this.mongoDBPartitionCreationSupplier = new MongoDBPartitionCreationSupplier(mongoDBConfig); - } - - public static MongoDBService create( - final MongoDBConfig mongoDBConfig, - final SourceCoordinator sourceCoordinator, - final Buffer> buffer, - final AcknowledgementSetManager acknowledgementSetManager, - final PluginMetrics pluginMetrics) { - return new MongoDBService( - mongoDBConfig, - sourceCoordinator, - buffer, - Executors.newSingleThreadScheduledExecutor(), - acknowledgementSetManager, - pluginMetrics); - } - - public void start() { - snapshotWorker = new MongoDBSnapshotWorker( - sourceCoordinator, - buffer, - mongoDBPartitionCreationSupplier, - pluginMetrics, - acknowledgementSetManager, - mongoDBConfig); - snapshotWorkerFuture = scheduledExecutorService.schedule(() -> snapshotWorker.run(), 0L, TimeUnit.MILLISECONDS); - } - - public void stop() { - scheduledExecutorService.shutdown(); - try { - snapshotWorkerFuture.cancel(true); - if (scheduledExecutorService.awaitTermination(EXECUTOR_SERVICE_SHUTDOWN_TIMEOUT.getSeconds(), TimeUnit.SECONDS)) { - LOG.info("Successfully waited for the snapshot worker to terminate"); - } else { - LOG.warn("snapshot worker did not terminate in time, forcing termination"); - scheduledExecutorService.shutdownNow(); - } - } catch (InterruptedException e) { - LOG.error("Interrupted while waiting for the snapshot worker to terminate", e); - scheduledExecutorService.shutdownNow(); - } - - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBSnapshotProgressState.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBSnapshotProgressState.java deleted file mode 100644 index 95ea8ba67e..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBSnapshotProgressState.java +++ /dev/null @@ -1,36 +0,0 @@ -package org.opensearch.dataprepper.plugins.kafkaconnect.source.mongoDB; - -import com.fasterxml.jackson.annotation.JsonProperty; - -public class MongoDBSnapshotProgressState { - @JsonProperty("totalRecords") - private long total; - @JsonProperty("successRecords") - private long success; - @JsonProperty("failedRecords") - private long failed; - - public long getTotal() { - return total; - } - - public long getSuccess() { - return success; - } - - public long getFailed() { - return failed; - } - - public void setTotal(long total) { - this.total = total; - } - - public void setSuccess(long successRecords) { - this.success = successRecords; - } - - public void setFailure(long failedRecords) { - this.failed = failedRecords; - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBSnapshotWorker.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBSnapshotWorker.java deleted file mode 100644 index 32819857eb..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBSnapshotWorker.java +++ /dev/null @@ -1,209 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.source.mongoDB; - -import com.fasterxml.jackson.core.JsonProcessingException; -import com.fasterxml.jackson.core.type.TypeReference; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.mongodb.client.MongoClient; -import com.mongodb.client.MongoCollection; -import com.mongodb.client.MongoCursor; -import com.mongodb.client.MongoDatabase; -import io.micrometer.core.instrument.Counter; -import org.bson.Document; -import org.bson.conversions.Bson; -import org.bson.json.JsonMode; -import org.bson.json.JsonWriterSettings; -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; -import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; -import org.opensearch.dataprepper.model.buffer.Buffer; -import org.opensearch.dataprepper.model.event.Event; -import org.opensearch.dataprepper.model.event.JacksonEvent; -import org.opensearch.dataprepper.model.opensearch.OpenSearchBulkActions; -import org.opensearch.dataprepper.model.record.Record; -import org.opensearch.dataprepper.model.source.coordinator.SourceCoordinator; -import org.opensearch.dataprepper.model.source.coordinator.SourcePartition; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.MongoDBConfig; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.time.Duration; -import java.util.List; -import java.util.Map; -import java.util.Optional; - -public class MongoDBSnapshotWorker implements Runnable { - private static final Logger LOG = LoggerFactory.getLogger(MongoDBSnapshotWorker.class); - private static final Duration BACKOFF_ON_EXCEPTION = Duration.ofSeconds(60); - private static final Duration BACKOFF_ON_EMPTY_PARTITION = Duration.ofSeconds(60); - private static final Duration ACKNOWLEDGEMENT_SET_TIMEOUT = Duration.ofHours(2); - private static final String SUCCESS_ITEM_COUNTER_NAME = "exportRecordsSuccessTotal"; - private static final String FAILURE_ITEM_COUNTER_NAME = "exportRecordsFailedTotal"; - private static final String SUCCESS_PARTITION_COUNTER_NAME = "exportPartitionSuccessTotal"; - private static final String FAILURE_PARTITION_COUNTER_NAME = "exportPartitionFailureTotal"; - private static final String EVENT_SOURCE_COLLECTION_ATTRIBUTE = "__collection"; - private static final String EVENT_SOURCE_DB_ATTRIBUTE = "__source_db"; - private static final String EVENT_SOURCE_OPERATION = "__op"; - private static final String EVENT_SOURCE_TS_MS = "__source_ts_ms"; - private static final String EVENT_TYPE = "EXPORT"; - private static final String PARTITION_KEY_SPLITTER = "\\|"; - private static final String COLLECTION_SPLITTER = "\\."; - private final SourceCoordinator sourceCoordinator; - private static int DEFAULT_BUFFER_WRITE_TIMEOUT_MS = 5000; - private final Buffer> buffer; - private final MongoDBPartitionCreationSupplier mongoDBPartitionCreationSupplier; - private final AcknowledgementSetManager acknowledgementSetManager; - private final MongoDBConfig mongoDBConfig; - private final Counter successItemsCounter; - private final Counter failureItemsCounter; - private final Counter successPartitionCounter; - private final Counter failureParitionCounter; - private final ObjectMapper objectMapper = new ObjectMapper(); - private final TypeReference> mapTypeReference = new TypeReference>() { - }; - - - public MongoDBSnapshotWorker(final SourceCoordinator sourceCoordinator, - final Buffer> buffer, - final MongoDBPartitionCreationSupplier mongoDBPartitionCreationSupplier, - final PluginMetrics pluginMetrics, - final AcknowledgementSetManager acknowledgementSetManager, - final MongoDBConfig mongoDBConfig) { - this.sourceCoordinator = sourceCoordinator; - this.buffer = buffer; - this.mongoDBPartitionCreationSupplier = mongoDBPartitionCreationSupplier; - this.acknowledgementSetManager = acknowledgementSetManager; - this.mongoDBConfig = mongoDBConfig; - this.successItemsCounter = pluginMetrics.counter(SUCCESS_ITEM_COUNTER_NAME); - this.failureItemsCounter = pluginMetrics.counter(FAILURE_ITEM_COUNTER_NAME); - this.successPartitionCounter = pluginMetrics.counter(SUCCESS_PARTITION_COUNTER_NAME); - this.failureParitionCounter = pluginMetrics.counter(FAILURE_PARTITION_COUNTER_NAME); - } - - @Override - public void run() { - while (!Thread.currentThread().isInterrupted()) { - try { - final Optional> snapshotPartition = sourceCoordinator.getNextPartition(mongoDBPartitionCreationSupplier); - if (snapshotPartition.isEmpty()) { - try { - LOG.info("get empty partition"); - Thread.sleep(BACKOFF_ON_EMPTY_PARTITION.toMillis()); - continue; - } catch (final InterruptedException e) { - LOG.info("The worker was interrupted while sleeping after acquiring no indices to process, stopping processing"); - return; - } - } - LOG.info("get partition success {}", snapshotPartition.get().getPartitionKey()); - try { - final Optional acknowledgementSet = createAcknowledgementSet(snapshotPartition.get()); - this.startProcessPartition(snapshotPartition.get()); - if (acknowledgementSet.isEmpty()) { - sourceCoordinator.completePartition(snapshotPartition.get().getPartitionKey(), false); - } else { - sourceCoordinator.updatePartitionForAcknowledgmentWait(snapshotPartition.get().getPartitionKey(), ACKNOWLEDGEMENT_SET_TIMEOUT); - acknowledgementSet.get().complete(); - } - successPartitionCounter.increment(); - } catch (final Exception e) { - LOG.error("Received an exception while processing the partition.", e); - sourceCoordinator.giveUpPartitions(); - failureParitionCounter.increment(); - } - } catch (final Exception e) { - LOG.error("Received an exception while trying to snapshot documentDB, backing off and retrying", e); - try { - Thread.sleep(BACKOFF_ON_EXCEPTION.toMillis()); - } catch (final InterruptedException ex) { - LOG.info("The DocumentDBSnapshotWorker was interrupted before backing off and retrying, stopping processing"); - return; - } - } - } - } - - private void startProcessPartition(SourcePartition partition) { - List partitionKeys = List.of(partition.getPartitionKey().split(PARTITION_KEY_SPLITTER)); - if (partitionKeys.size() < 4) { - throw new RuntimeException("Invalid Partition Key. Must as db.collection|gte|lte format."); - } - List collection = List.of(partitionKeys.get(0).split(COLLECTION_SPLITTER)); - final String gte = partitionKeys.get(1); - final String lte = partitionKeys.get(2); - final String className = partitionKeys.get(3); - if (collection.size() < 2) { - throw new RuntimeException("Invalid Collection Name. Must as db.collection format"); - } - try (MongoClient mongoClient = MongoDBHelper.getMongoClient(mongoDBConfig)) { - MongoDatabase db = mongoClient.getDatabase(collection.get(0)); - MongoCollection col = db.getCollection(collection.get(1)); - Bson query = MongoDBHelper.buildAndQuery(gte, lte, className); - long totalRecords = 0L; - long successRecords = 0L; - long failedRecords = 0L; - try (MongoCursor cursor = col.find(query).iterator()) { - while (cursor.hasNext()) { - totalRecords += 1; - try { - JsonWriterSettings writerSettings = JsonWriterSettings.builder() - .outputMode(JsonMode.RELAXED) - .objectIdConverter((value, writer) -> writer.writeString(value.toHexString())) - .build(); - String record = cursor.next().toJson(writerSettings); - Map data = convertToMap(record); - data.putIfAbsent(EVENT_SOURCE_DB_ATTRIBUTE, collection.get(0)); - data.putIfAbsent(EVENT_SOURCE_COLLECTION_ATTRIBUTE, collection.get(1)); - data.putIfAbsent(EVENT_SOURCE_OPERATION, OpenSearchBulkActions.CREATE.toString()); - data.putIfAbsent(EVENT_SOURCE_TS_MS, 0); - if (buffer.isByteBuffer()) { - buffer.writeBytes(objectMapper.writeValueAsBytes(data), null, DEFAULT_BUFFER_WRITE_TIMEOUT_MS); - } else { - buffer.write(getEventFromData(data), DEFAULT_BUFFER_WRITE_TIMEOUT_MS); - } - successItemsCounter.increment(); - successRecords += 1; - } catch (Exception e) { - LOG.error("failed to add record to buffer with error {}", e.getMessage()); - failureItemsCounter.increment(); - failedRecords += 1; - } - } - final MongoDBSnapshotProgressState progressState = new MongoDBSnapshotProgressState(); - progressState.setTotal(totalRecords); - progressState.setSuccess(successRecords); - progressState.setFailure(failedRecords); - sourceCoordinator.saveProgressStateForPartition(partition.getPartitionKey(), progressState); - } catch (Exception e) { - throw new RuntimeException(e); - } - } - } - - private Optional createAcknowledgementSet(SourcePartition partition) { - if (mongoDBConfig.getExportConfig().getAcknowledgements()) { - return Optional.of(this.acknowledgementSetManager.create((result) -> { - if (result) { - this.sourceCoordinator.completePartition(partition.getPartitionKey(), true); - } - }, ACKNOWLEDGEMENT_SET_TIMEOUT)); - } - return Optional.empty(); - } - - private Map convertToMap(String jsonData) throws JsonProcessingException { - return objectMapper.readValue(jsonData, mapTypeReference); - } - - private Record getEventFromData(Map data) { - Event event = JacksonEvent.builder() - .withEventType(EVENT_TYPE) - .withData(data) - .build(); - return new Record<>(event); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/Connector.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/Connector.java deleted file mode 100644 index a413a89bbd..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/Connector.java +++ /dev/null @@ -1,33 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.util; - -import java.util.Map; - -public class Connector { - private final String name; - private final Map config; - private final Boolean allowReplace; - - public Connector(final String name, final Map config, final Boolean allowReplace) { - this.name = name; - this.config = config; - this.allowReplace = allowReplace; - } - - public String getName() { - return this.name; - } - - public Map getConfig() { - config.putIfAbsent("name", name); - return config; - } - - public Boolean getAllowReplace() { - return allowReplace; - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/KafkaConnect.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/KafkaConnect.java deleted file mode 100644 index 3a1e82bace..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/KafkaConnect.java +++ /dev/null @@ -1,375 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.util; - -import org.apache.kafka.common.utils.Time; -import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy; -import org.apache.kafka.connect.errors.AlreadyExistsException; -import org.apache.kafka.connect.errors.NotFoundException; -import org.apache.kafka.connect.json.JsonConverter; -import org.apache.kafka.connect.json.JsonConverterConfig; -import org.apache.kafka.connect.runtime.Connect; -import org.apache.kafka.connect.runtime.Worker; -import org.apache.kafka.connect.runtime.WorkerConfig; -import org.apache.kafka.connect.runtime.WorkerConfigTransformer; -import org.apache.kafka.connect.runtime.distributed.DistributedConfig; -import org.apache.kafka.connect.runtime.distributed.DistributedHerder; -import org.apache.kafka.connect.runtime.distributed.NotLeaderException; -import org.apache.kafka.connect.runtime.isolation.Plugins; -import org.apache.kafka.connect.runtime.rest.ConnectRestServer; -import org.apache.kafka.connect.runtime.rest.RestClient; -import org.apache.kafka.connect.runtime.rest.RestServer; -import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo; -import org.apache.kafka.connect.storage.ConfigBackingStore; -import org.apache.kafka.connect.storage.Converter; -import org.apache.kafka.connect.storage.KafkaConfigBackingStore; -import org.apache.kafka.connect.storage.KafkaOffsetBackingStore; -import org.apache.kafka.connect.storage.KafkaStatusBackingStore; -import org.apache.kafka.connect.storage.StatusBackingStore; -import org.apache.kafka.connect.util.ConnectUtils; -import org.apache.kafka.connect.util.SharedTopicAdmin; -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.plugins.kafkaconnect.meter.KafkaConnectMetrics; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.net.URI; -import java.time.Clock; -import java.time.Duration; -import java.util.ArrayList; -import java.util.Collection; -import java.util.Collections; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.TimeUnit; -import java.util.stream.Collectors; - -import static org.apache.kafka.clients.CommonClientConfigs.CLIENT_ID_CONFIG; - -/** - * The KafkaConnect infra. - * Unique with single instance for each pipeline. - */ -public class KafkaConnect { - private static final Logger LOG = LoggerFactory.getLogger(KafkaConnect.class); - private static volatile Map instanceMap = new HashMap<>(); - private static final long RETRY_INTERVAL_MS = 3000L; // 3 seconds - private static final int LATCH_WAIT_TIME = 1; // 1 minute - private static final String RUNNING = "RUNNING"; - private final Map connectorMap; - private final KafkaConnectMetrics kafkaConnectMetrics; - private final Time time = Time.SYSTEM; - private final Clock clock = Clock.systemUTC(); - private DistributedHerder herder; - private RestServer rest; - private Connect connect; - private final long connectTimeoutMs; // 60 seconds - private final long connectorTimeoutMs; // 30 seconds - - private KafkaConnect(final PluginMetrics pluginMetrics, - final Duration connectTimeout, - final Duration connectorTimeout) { - this.connectorMap = new HashMap<>(); - this.kafkaConnectMetrics = new KafkaConnectMetrics(pluginMetrics); - this.connectTimeoutMs = connectTimeout.toMillis(); - this.connectorTimeoutMs = connectorTimeout.toMillis(); - } - - /** - * For Testing - */ - public KafkaConnect(final DistributedHerder herder, - final RestServer rest, - final Connect connect, - final KafkaConnectMetrics kafkaConnectMetrics) { - this.connectorMap = new HashMap<>(); - this.herder = herder; - this.rest = rest; - this.connect = connect; - this.kafkaConnectMetrics = kafkaConnectMetrics; - this.connectTimeoutMs = 60000L; - this.connectorTimeoutMs = 30000L; - } - - public static KafkaConnect getPipelineInstance(final String pipelineName, - final PluginMetrics pluginMetrics, - final Duration connectTimeout, - final Duration connectorTimeout) { - KafkaConnect instance = instanceMap.get(pipelineName); - if (instance == null) { - synchronized (KafkaConnect.class) { - instance = new KafkaConnect(pluginMetrics, connectTimeout, connectorTimeout); - instanceMap.put(pipelineName, instance); - } - } - return instance; - } - - public synchronized void initialize(Map workerProps) { - DistributedConfig config = new DistributedConfig(workerProps); - RestClient restClient = new RestClient(config); - this.rest = new ConnectRestServer(config.rebalanceTimeout(), restClient, workerProps); - this.herder = initHerder(workerProps, config, restClient); - this.connect = new Connect(herder, (ConnectRestServer) rest); - } - - /** - * Add connectors to the Kafka Connect. - * This must be done before the start() is called. - * - * @param connectors connectors to be added. - */ - public void addConnectors(List connectors) { - connectors.forEach(connector -> { - this.connectorMap.put(connector.getName(), connector); - }); - } - - /** - * Start the kafka connect. - * Will add all connectors, and cleanup unused connectors at once. - */ - public synchronized void start() { - if (this.connect == null) { - throw new RuntimeException("Please initialize Kafka Connect first."); - } - if (this.connect.isRunning()) { - LOG.info("Kafka Connect is running, will not start again"); - return; - } - LOG.info("Starting Kafka Connect"); - try { - this.rest.initializeServer(); - this.connect.start(); - waitForConnectRunning(); - this.kafkaConnectMetrics.bindConnectMetrics(); - this.initConnectors(); - } catch (Exception e) { - LOG.error("Failed to start Connect", e); - this.connect.stop(); - throw new RuntimeException(e); - } - } - - /** - * Stop the Kafka Connect. - */ - public void stop() { - if (this.connect == null) { - LOG.info("Kafka Connect is running, will not start again"); - return; - } - LOG.info("Stopping Kafka Connect"); - this.connect.stop(); - } - - private DistributedHerder initHerder(Map workerProps, DistributedConfig config, RestClient restClient) { - LOG.info("Scanning for plugin classes. This might take a moment ..."); - Plugins plugins = new Plugins(workerProps); - plugins.compareAndSwapWithDelegatingLoader(); - String kafkaClusterId = config.kafkaClusterId(); - LOG.info("Kafka cluster ID: {}", kafkaClusterId); - - URI advertisedUrl = rest.advertisedUrl(); - String workerId = advertisedUrl.getHost() + ":" + advertisedUrl.getPort(); - - String clientIdBase = ConnectUtils.clientIdBase(config); - - // Create the admin client to be shared by all backing stores. - Map adminProps = new HashMap<>(config.originals()); - ConnectUtils.addMetricsContextProperties(adminProps, config, kafkaClusterId); - adminProps.put(CLIENT_ID_CONFIG, clientIdBase + "shared-admin"); - SharedTopicAdmin sharedAdmin = new SharedTopicAdmin(adminProps); - - KafkaOffsetBackingStore offsetBackingStore = new KafkaOffsetBackingStore(sharedAdmin, () -> clientIdBase, - plugins.newInternalConverter(true, JsonConverter.class.getName(), - Collections.singletonMap(JsonConverterConfig.SCHEMAS_ENABLE_CONFIG, "false"))); - offsetBackingStore.configure(config); - - ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy = plugins.newPlugin( - config.getString(WorkerConfig.CONNECTOR_CLIENT_POLICY_CLASS_CONFIG), - config, ConnectorClientConfigOverridePolicy.class); - - Worker worker = new Worker(workerId, time, plugins, config, offsetBackingStore, connectorClientConfigOverridePolicy); - WorkerConfigTransformer configTransformer = worker.configTransformer(); - - Converter internalValueConverter = worker.getInternalValueConverter(); - StatusBackingStore statusBackingStore = new KafkaStatusBackingStore(time, internalValueConverter, sharedAdmin, clientIdBase); - statusBackingStore.configure(config); - - ConfigBackingStore configBackingStore = new KafkaConfigBackingStore( - internalValueConverter, - config, - configTransformer, - sharedAdmin, - clientIdBase); - - // Pass the shared admin to the distributed herder as an additional AutoCloseable object that should be closed when the - // herder is stopped. This is easier than having to track and own the lifecycle ourselves. - return new DistributedHerder(config, time, worker, - kafkaClusterId, statusBackingStore, configBackingStore, - advertisedUrl.toString(), restClient, connectorClientConfigOverridePolicy, - Collections.emptyList(), sharedAdmin); - } - - /** - * - * @throws InterruptedException - */ - private void waitForConnectRunning() throws InterruptedException { - long startTime = clock.millis(); - boolean isRunning = false; - while (clock.millis() - startTime < connectTimeoutMs) { - LOG.info("Waiting Kafka Connect running"); - isRunning = this.connect.isRunning(); - if (isRunning) break; - TimeUnit.MILLISECONDS.sleep(RETRY_INTERVAL_MS); - } - if (!isRunning) { - throw new RuntimeException("Timed out waiting for Kafka Connect running"); - } - LOG.info("Kafka Connect is running"); - } - - /** - * Initialize connectors. - * The Kafka Connectors are managed in orders: - * 1. Delete Connectors not in pipeline configurations - * 2. Register Connectors - * 3. Wait for all connectors in running state. - * 4. Bind connectors' metrics - */ - private void initConnectors() throws InterruptedException { - this.deleteConnectors(); - this.registerConnectors(); - this.waitForConnectorsRunning(); - this.kafkaConnectMetrics.bindConnectorMetrics(); - } - - /** - * Register Connector to Kafka Connect. - * Designed as private method to prevent register the connector after connect is started. - */ - private void registerConnectors() throws InterruptedException { - CountDownLatch connectorLatch = new CountDownLatch(connectorMap.size()); - List exceptionMessages = new ArrayList<>(); - connectorMap.forEach((connectorName, connector) -> { - herder.connectorConfig(connectorName, (e, config) -> { - boolean shouldUpdate; - if (config == null) { - shouldUpdate = true; - } else { - shouldUpdate = connector.getAllowReplace() || (!config.equals(connector.getConfig())); - } - herder.putConnectorConfig(connectorName, connector.getConfig(), shouldUpdate, (error, result) -> { - if (error != null) { - if (error instanceof NotLeaderException || error instanceof AlreadyExistsException) { - LOG.info(error.getMessage()); - } else { - LOG.error("Failed to put connector config: {}", connectorName); - exceptionMessages.add(error.getMessage()); - } - } else { - // Handle the successful registration - LOG.info("Success put connector config: {}", connectorName); - } - connectorLatch.countDown(); - }); - }); - }); - // Block and wait for all tasks to complete - if (!connectorLatch.await(LATCH_WAIT_TIME, TimeUnit.MINUTES)) { - throw new RuntimeException("Timed out waiting for initConnectors"); - } else { - if (!exceptionMessages.isEmpty()) { - throw new RuntimeException(String.join(", ", exceptionMessages)); - } - LOG.info("InitConnectors completed"); - } - } - - /** - * Delete Connectors from Kafka Connect. - * Designed as private method to prevent delete the connector after connect is started. - */ - private void deleteConnectors() throws InterruptedException { - Collection connectorsToDelete = this.herder.connectors() - .stream() - .filter(connectorName -> !connectorMap.containsKey(connectorName)) - .collect(Collectors.toList()); - List exceptionMessages = new ArrayList<>(); - CountDownLatch deleteLatch = new CountDownLatch(connectorsToDelete.size()); - connectorsToDelete.forEach(connectorName -> { - herder.deleteConnectorConfig(connectorName, (error, result) -> { - if (error != null) { - if (error instanceof NotLeaderException || error instanceof NotFoundException) { - LOG.info(error.getMessage()); - } else { - LOG.error("Failed to delete connector config: {}", connectorName); - exceptionMessages.add(error.getMessage()); - } - } else { - // Handle the successful registration - LOG.info("Success delete connector config: {}", connectorName); - } - deleteLatch.countDown(); - }); - }); - // Block and wait for all tasks to complete - if (!deleteLatch.await(LATCH_WAIT_TIME, TimeUnit.MINUTES)) { - throw new RuntimeException("Timed out waiting for deleteConnectors"); - } else { - if (!exceptionMessages.isEmpty()) { - throw new RuntimeException(String.join(", ", exceptionMessages)); - } - LOG.info("deleteConnectors completed"); - } - } - - private void waitForConnectorsRunning() throws InterruptedException { - LOG.info("Waiting for connectors to be running"); - Set connectorNames = this.connectorMap.keySet(); - List exceptionMessages = new ArrayList<>(); - CountDownLatch countDownLatch = new CountDownLatch(connectorNames.size()); - connectorNames.parallelStream().forEach(connectorName -> { - long startTime = clock.millis(); - boolean isRunning = false; - while (clock.millis() - startTime < connectorTimeoutMs) { - try { - ConnectorStateInfo info = herder.connectorStatus(connectorName); - if (RUNNING.equals(info.connector().state())) { - // Connector is running, decrement the latch count - isRunning = true; - break; - } - } catch (Exception e) { - LOG.info(e.getMessage()); - } - try { - TimeUnit.MILLISECONDS.sleep(RETRY_INTERVAL_MS); - } catch (InterruptedException e) { - break; - } - } - countDownLatch.countDown(); - if (!isRunning) { - exceptionMessages.add(String.format("Connector %s is not running in desired period of time", connectorName)); - } - }); - // Block and wait for all tasks to complete - if (!countDownLatch.await(LATCH_WAIT_TIME, TimeUnit.MINUTES)) { - throw new RuntimeException("Timed out waiting for running state check"); - } else { - if (!exceptionMessages.isEmpty()) { - throw new RuntimeException(String.join(", ", exceptionMessages)); - } - LOG.info("All connectors are running"); - } - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/SecretManagerHelper.java b/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/SecretManagerHelper.java deleted file mode 100644 index a4a31c52bb..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/SecretManagerHelper.java +++ /dev/null @@ -1,58 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.util; - -import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; -import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; -import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration; -import software.amazon.awssdk.core.retry.RetryPolicy; -import software.amazon.awssdk.regions.Region; -import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient; -import software.amazon.awssdk.services.secretsmanager.model.GetSecretValueRequest; -import software.amazon.awssdk.services.secretsmanager.model.GetSecretValueResponse; -import software.amazon.awssdk.services.sts.StsClient; -import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider; -import software.amazon.awssdk.services.sts.model.AssumeRoleRequest; - -import java.util.UUID; - -public class SecretManagerHelper { - private static final String SESSION_PREFIX = "data-prepper-secretmanager-session"; - public static String getSecretValue(final String stsRoleArn, final String region, final String secretId) { - AwsCredentialsProvider credentialsProvider = DefaultCredentialsProvider.create(); - ClientOverrideConfiguration clientOverrideConfiguration = ClientOverrideConfiguration - .builder() - .retryPolicy(RetryPolicy.defaultRetryPolicy()) - .build(); - - if (stsRoleArn != null && !stsRoleArn.isEmpty()) { - String sessionName = SESSION_PREFIX + UUID.randomUUID(); - StsClient stsClient = StsClient.builder() - .overrideConfiguration(clientOverrideConfiguration) - .region(Region.of(region)) - .credentialsProvider(credentialsProvider) - .build(); - AssumeRoleRequest assumeRoleRequest = AssumeRoleRequest - .builder() - .roleArn(stsRoleArn) - .roleSessionName(sessionName) - .build(); - credentialsProvider = StsAssumeRoleCredentialsProvider - .builder() - .stsClient(stsClient) - .refreshRequest(assumeRoleRequest) - .build(); - } - SecretsManagerClient secretsManagerClient = SecretsManagerClient.builder() - .overrideConfiguration(clientOverrideConfiguration) - .credentialsProvider(credentialsProvider) - .region(Region.of(region)) - .build(); - final GetSecretValueRequest request = GetSecretValueRequest.builder().secretId(secretId).build(); - final GetSecretValueResponse response = secretsManagerClient.getSecretValue(request); - return response.secretString(); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/CredentialConfigTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/CredentialConfigTest.java deleted file mode 100644 index e36af1e63a..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/CredentialConfigTest.java +++ /dev/null @@ -1,87 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.configuration; - -import org.junit.jupiter.api.Test; -import org.mockito.MockedStatic; -import org.opensearch.dataprepper.plugins.kafkaconnect.util.SecretManagerHelper; - -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.is; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.Mockito.mockStatic; - -public class CredentialConfigTest { - private final String testUserName = "testUser"; - private final String testPassword = "testPassword"; - private final String testStsRole = "testRole"; - private final String testRegion = "testRegion"; - private final String testSecretId = "testSecritId"; - - @Test - void test_credential_config_plaintext() { - CredentialsConfig credentialsConfig = new CredentialsConfig( - new CredentialsConfig.PlainText(testUserName, testPassword), null); - assertThat(credentialsConfig.getUsername(), is(testUserName)); - assertThat(credentialsConfig.getPassword(), is(testPassword)); - } - - @Test - void test_credential_config_plaintext_invalid() { - assertThrows(IllegalArgumentException.class, () -> new CredentialsConfig( - new CredentialsConfig.PlainText(null, null), null)); - assertThrows(IllegalArgumentException.class, () -> new CredentialsConfig( - new CredentialsConfig.PlainText(testUserName, null), null)); - assertThrows(IllegalArgumentException.class, () -> new CredentialsConfig( - new CredentialsConfig.PlainText(null, testPassword), null)); - } - - @Test - void test_credential_config_secret_manager() { - final String expectedSecret = "{\"username\":\"expectedUsername\",\"password\":\"expectedPassword\"}"; - try (MockedStatic mockedStatic = mockStatic(SecretManagerHelper.class)) { - mockedStatic.when(() -> SecretManagerHelper.getSecretValue(testStsRole, testRegion, testSecretId)).thenReturn(expectedSecret); - CredentialsConfig credentialsConfig = new CredentialsConfig( - null, new CredentialsConfig.SecretManager(testStsRole, testRegion, testSecretId)); - assertThat(credentialsConfig.getUsername(), is("expectedUsername")); - assertThat(credentialsConfig.getPassword(), is("expectedPassword")); - } - } - - @Test - void test_credential_config_failure_on_secret_manager() { - try (MockedStatic mockedStatic = mockStatic(SecretManagerHelper.class)) { - mockedStatic.when(() -> SecretManagerHelper.getSecretValue(testStsRole, testRegion, testSecretId)).thenThrow(new RuntimeException()); - assertThrows(RuntimeException.class, () -> new CredentialsConfig( - null, new CredentialsConfig.SecretManager(testStsRole, testRegion, testSecretId))); - final String invalidSecret = "{}"; - mockedStatic.when(() -> SecretManagerHelper.getSecretValue(testStsRole, testRegion, testSecretId)).thenReturn(invalidSecret); - assertThrows(RuntimeException.class, () -> new CredentialsConfig( - null, new CredentialsConfig.SecretManager(testStsRole, testRegion, testSecretId))); - } - } - - @Test - void test_credential_config_secret_manager_invalid_input() { - assertThrows(IllegalArgumentException.class, () -> new CredentialsConfig( - null, new CredentialsConfig.SecretManager(null, null, null))); - assertThrows(IllegalArgumentException.class, () -> new CredentialsConfig( - null, new CredentialsConfig.SecretManager(null, null, testSecretId))); - assertThrows(IllegalArgumentException.class, () -> new CredentialsConfig( - null, new CredentialsConfig.SecretManager(null, testRegion, null))); - } - - @Test - void test_invalid_credential_config() { - // Must be set - assertThrows(IllegalArgumentException.class, () -> new CredentialsConfig(null, null)); - // Cannot both set - assertThrows(IllegalArgumentException.class, () -> new CredentialsConfig( - new CredentialsConfig.PlainText(testUserName, testPassword), - new CredentialsConfig.SecretManager(testStsRole, testRegion, testSecretId) - )); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/MongoDBConfigTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/MongoDBConfigTest.java deleted file mode 100644 index ef3bf843b8..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/MongoDBConfigTest.java +++ /dev/null @@ -1,81 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.configuration; - -import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; -import org.junit.jupiter.api.Test; -import org.opensearch.dataprepper.plugins.kafkaconnect.util.Connector; -import org.yaml.snakeyaml.Yaml; - -import java.io.FileReader; -import java.io.IOException; -import java.io.Reader; -import java.io.StringReader; -import java.util.Map; - -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.instanceOf; -import static org.hamcrest.Matchers.is; -import static org.hamcrest.Matchers.notNullValue; - -public class MongoDBConfigTest { - - @Test - public void test_get_mongodb_connectors() throws IOException { - MongoDBConfig testConfig = buildTestConfig("sample-mongodb-pipeline.yaml"); - assertThat(testConfig, notNullValue()); - assertThat(testConfig.buildConnectors(), notNullValue()); - assertThat(testConfig.buildConnectors().size(), is(1)); - // verify Connector - Connector mongodbConnector = testConfig.buildConnectors().get(0); - assertThat(mongodbConnector, instanceOf(Connector.class)); - final Map actualConfig = mongodbConnector.getConfig(); - assertThat(actualConfig.get("connector.class"), is(MongoDBConfig.CONNECTOR_CLASS)); - assertThat(actualConfig.get("mongodb.connection.string"), is("mongodb://localhost:27017/?replicaSet=rs0&directConnection=true")); - assertThat(actualConfig.get("mongodb.user"), is("debezium")); - assertThat(actualConfig.get("mongodb.password"), is("dbz")); - assertThat(actualConfig.get("snapshot.mode"), is("never")); - assertThat(actualConfig.get("topic.prefix"), is("prefix1")); - assertThat(actualConfig.get("collection.include.list"), is("test.customers")); - assertThat(actualConfig.get("mongodb.ssl.enabled"), is("false")); - } - - @Test - public void test_get_mongodb_config_props() throws IOException { - MongoDBConfig testConfig = buildTestConfig("sample-mongodb-pipeline.yaml"); - assertThat(testConfig, notNullValue()); - assertThat(testConfig.getIngestionMode(), is(MongoDBConfig.IngestionMode.EXPORT_STREAM)); - assertThat(testConfig.getCredentialsConfig().getUsername(), is("debezium")); - assertThat(testConfig.getHostname(), is("localhost")); - assertThat(testConfig.getPort(), is("27017")); - assertThat(testConfig.getSSLEnabled(), is(false)); - assertThat(testConfig.getSSLInvalidHostAllowed(), is(false)); - assertThat(testConfig.getCollections().size(), is(1)); - assertThat(testConfig.getExportConfig().getAcknowledgements(), is(false)); - assertThat(testConfig.getExportConfig().getItemsPerPartition(), is(4000)); - assertThat(testConfig.getExportConfig().getReadPreference(), is("secondaryPreferred")); - } - - private MongoDBConfig buildTestConfig(final String resourceFileName) throws IOException { - //Added to load Yaml file - Start - Yaml yaml = new Yaml(); - FileReader fileReader = new FileReader(getClass().getClassLoader().getResource(resourceFileName).getFile()); - Object data = yaml.load(fileReader); - if (data instanceof Map) { - Map propertyMap = (Map) data; - Map logPipelineMap = (Map) propertyMap.get("log-pipeline"); - Map sourceMap = (Map) logPipelineMap.get("source"); - Map kafkaConnectConfigMap = (Map) sourceMap.get("mongodb"); - ObjectMapper mapper = new ObjectMapper(); - mapper.registerModule(new JavaTimeModule()); - String json = mapper.writeValueAsString(kafkaConnectConfigMap); - Reader reader = new StringReader(json); - return mapper.readValue(reader, MongoDBConfig.class); - } - return null; - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/MySQLConfigTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/MySQLConfigTest.java deleted file mode 100644 index fa4d3526da..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/MySQLConfigTest.java +++ /dev/null @@ -1,76 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.configuration; - -import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; -import org.junit.jupiter.api.Test; -import org.opensearch.dataprepper.plugins.kafkaconnect.util.Connector; -import org.yaml.snakeyaml.Yaml; - -import java.io.FileReader; -import java.io.IOException; -import java.io.Reader; -import java.io.StringReader; -import java.util.Map; -import java.util.Properties; - -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.instanceOf; -import static org.hamcrest.Matchers.is; -import static org.hamcrest.Matchers.notNullValue; - -public class MySQLConfigTest { - @Test - public void test_get_mysql_connectors() throws IOException { - final String bootstrapServers = "localhost:9092"; - final Properties authProperties = new Properties(); - authProperties.put("bootstrap.servers", bootstrapServers); - authProperties.put("testClass", this.getClass()); - authProperties.put("testKey", "testValue"); - MySQLConfig testConfig = buildTestConfig("sample-mysql-pipeline.yaml"); - assertThat(testConfig, notNullValue()); - assertThat(testConfig.buildConnectors(), notNullValue()); - assertThat(testConfig.buildConnectors().size(), is(1)); - // verify Connector - testConfig.setAuthProperties(authProperties); - Connector mysqlConnector = testConfig.buildConnectors().get(0); - assertThat(mysqlConnector, instanceOf(Connector.class)); - final Map actualConfig = mysqlConnector.getConfig(); - assertThat(actualConfig.get("connector.class"), is(MySQLConfig.CONNECTOR_CLASS)); - assertThat(actualConfig.get("database.hostname"), is("localhost")); - assertThat(actualConfig.get("database.port"), is("3306")); - assertThat(actualConfig.get("database.user"), is("debezium")); - assertThat(actualConfig.get("database.password"), is("dbz")); - assertThat(actualConfig.get("snapshot.mode"), is("initial")); - assertThat(actualConfig.get("topic.prefix"), is("prefix1")); - assertThat(actualConfig.get("table.include.list"), is("inventory.customers")); - assertThat(actualConfig.get("schema.history.internal.kafka.bootstrap.servers"), is(bootstrapServers)); - assertThat(actualConfig.get("schema.history.internal.producer.testKey"), is(authProperties.getProperty("testKey"))); - assertThat(actualConfig.get("schema.history.internal.consumer.testKey"), is(authProperties.getProperty("testKey"))); - assertThat(actualConfig.get("schema.history.internal.producer.testClass"), is(this.getClass().getName())); - assertThat(actualConfig.get("schema.history.internal.consumer.testClass"), is(this.getClass().getName())); - } - - private MySQLConfig buildTestConfig(final String resourceFileName) throws IOException { - //Added to load Yaml file - Start - Yaml yaml = new Yaml(); - FileReader fileReader = new FileReader(getClass().getClassLoader().getResource(resourceFileName).getFile()); - Object data = yaml.load(fileReader); - if (data instanceof Map) { - Map propertyMap = (Map) data; - Map logPipelineMap = (Map) propertyMap.get("log-pipeline"); - Map sourceMap = (Map) logPipelineMap.get("source"); - Map kafkaConnectConfigMap = (Map) sourceMap.get("mysql"); - ObjectMapper mapper = new ObjectMapper(); - mapper.registerModule(new JavaTimeModule()); - String json = mapper.writeValueAsString(kafkaConnectConfigMap); - Reader reader = new StringReader(json); - return mapper.readValue(reader, MySQLConfig.class); - } - return null; - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/PostgreSQLConfigTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/PostgreSQLConfigTest.java deleted file mode 100644 index 036cda30aa..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/configuration/PostgreSQLConfigTest.java +++ /dev/null @@ -1,67 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.configuration; - -import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; -import org.junit.jupiter.api.Test; -import org.opensearch.dataprepper.plugins.kafkaconnect.util.Connector; -import org.yaml.snakeyaml.Yaml; - -import java.io.FileReader; -import java.io.IOException; -import java.io.Reader; -import java.io.StringReader; -import java.util.Map; - -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.instanceOf; -import static org.hamcrest.Matchers.is; -import static org.hamcrest.Matchers.notNullValue; - -public class PostgreSQLConfigTest { - @Test - public void test_get_postgresql_connectors() throws IOException { - PostgreSQLConfig testConfig = buildTestConfig("sample-postgres-pipeline.yaml"); - assertThat(testConfig, notNullValue()); - assertThat(testConfig.buildConnectors(), notNullValue()); - assertThat(testConfig.buildConnectors().size(), is(1)); - // verify Connector - Connector postgresqlConnector = testConfig.buildConnectors().get(0); - assertThat(postgresqlConnector, instanceOf(Connector.class)); - assertThat(postgresqlConnector.getName(), is("psql.public.customers")); - final Map actualConfig = postgresqlConnector.getConfig(); - assertThat(actualConfig.get("connector.class"), is(PostgreSQLConfig.CONNECTOR_CLASS)); - assertThat(actualConfig.get("plugin.name"), is("pgoutput")); - assertThat(actualConfig.get("database.hostname"), is("localhost")); - assertThat(actualConfig.get("database.port"), is("5432")); - assertThat(actualConfig.get("database.user"), is("debezium")); - assertThat(actualConfig.get("database.password"), is("dbz")); - assertThat(actualConfig.get("snapshot.mode"), is("initial")); - assertThat(actualConfig.get("topic.prefix"), is("psql")); - assertThat(actualConfig.get("database.dbname"), is("postgres")); - assertThat(actualConfig.get("table.include.list"), is("public.customers")); - } - - private PostgreSQLConfig buildTestConfig(final String resourceFileName) throws IOException { - //Added to load Yaml file - Start - Yaml yaml = new Yaml(); - FileReader fileReader = new FileReader(getClass().getClassLoader().getResource(resourceFileName).getFile()); - Object data = yaml.load(fileReader); - if (data instanceof Map) { - Map propertyMap = (Map) data; - Map logPipelineMap = (Map) propertyMap.get("log-pipeline"); - Map sourceMap = (Map) logPipelineMap.get("source"); - Map kafkaConnectConfigMap = (Map) sourceMap.get("postgresql"); - ObjectMapper mapper = new ObjectMapper(); - mapper.registerModule(new JavaTimeModule()); - String json = mapper.writeValueAsString(kafkaConnectConfigMap); - Reader reader = new StringReader(json); - return mapper.readValue(reader, PostgreSQLConfig.class); - } - return null; - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/DefaultKafkaConnectConfigSupplierTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/DefaultKafkaConnectConfigSupplierTest.java deleted file mode 100644 index d4bfb32504..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/DefaultKafkaConnectConfigSupplierTest.java +++ /dev/null @@ -1,29 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.extension; - -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; - -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.MatcherAssert.assertThat; - -@ExtendWith(MockitoExtension.class) -public class DefaultKafkaConnectConfigSupplierTest { - @Mock - private KafkaConnectConfig kafkaConnectConfig; - - private DefaultKafkaConnectConfigSupplier createObjectUnderTest() { - return new DefaultKafkaConnectConfigSupplier(kafkaConnectConfig); - } - - @Test - void test_get_config() { - assertThat(createObjectUnderTest().getConfig(), equalTo(kafkaConnectConfig)); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigExtensionTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigExtensionTest.java deleted file mode 100644 index e2a705978b..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigExtensionTest.java +++ /dev/null @@ -1,44 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.extension; - -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.ArgumentCaptor; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.model.plugin.ExtensionPoints; -import org.opensearch.dataprepper.model.plugin.ExtensionProvider; - -import static org.hamcrest.CoreMatchers.instanceOf; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.mockito.Mockito.verify; - -@ExtendWith(MockitoExtension.class) -public class KafkaConnectConfigExtensionTest { - @Mock - private ExtensionPoints extensionPoints; - - @Mock - private KafkaConnectConfig kafkaConnectConfig; - - private KafkaConnectConfigExtension createObjectUnderTest() { - return new KafkaConnectConfigExtension(kafkaConnectConfig); - } - - @Test - void apply_should_addExtensionProvider() { - createObjectUnderTest().apply(extensionPoints); - final ArgumentCaptor extensionProviderArgumentCaptor = - ArgumentCaptor.forClass(ExtensionProvider.class); - - verify(extensionPoints).addExtensionProvider(extensionProviderArgumentCaptor.capture()); - - final ExtensionProvider actualExtensionProvider = extensionProviderArgumentCaptor.getValue(); - - assertThat(actualExtensionProvider, instanceOf(KafkaConnectConfigProvider.class)); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigProviderTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigProviderTest.java deleted file mode 100644 index 3162f0d193..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigProviderTest.java +++ /dev/null @@ -1,52 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.extension; - -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.model.plugin.ExtensionProvider; - -import java.util.Optional; - -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.CoreMatchers.notNullValue; -import static org.hamcrest.CoreMatchers.sameInstance; -import static org.hamcrest.MatcherAssert.assertThat; - -@ExtendWith(MockitoExtension.class) -public class KafkaConnectConfigProviderTest { - @Mock - private KafkaConnectConfigSupplier kafkaConnectConfigSupplier; - - @Mock - private ExtensionProvider.Context context; - - private KafkaConnectConfigProvider createObjectUnderTest() { - return new KafkaConnectConfigProvider(kafkaConnectConfigSupplier); - } - - @Test - void supportedClass_returns_kafkaConnectConfigSupplier() { - assertThat(createObjectUnderTest().supportedClass(), equalTo(KafkaConnectConfigSupplier.class)); - } - - @Test - void provideInstance_returns_the_kafkaConnectConfigSupplier_from_the_constructor() { - final KafkaConnectConfigProvider objectUnderTest = createObjectUnderTest(); - - final Optional optionalKafkaConnectConfigSupplier = objectUnderTest.provideInstance(context); - assertThat(optionalKafkaConnectConfigSupplier, notNullValue()); - assertThat(optionalKafkaConnectConfigSupplier.isPresent(), equalTo(true)); - assertThat(optionalKafkaConnectConfigSupplier.get(), equalTo(kafkaConnectConfigSupplier)); - - final Optional anotherOptionalKafkaConnectConfigSupplier = objectUnderTest.provideInstance(context); - assertThat(anotherOptionalKafkaConnectConfigSupplier, notNullValue()); - assertThat(anotherOptionalKafkaConnectConfigSupplier.isPresent(), equalTo(true)); - assertThat(anotherOptionalKafkaConnectConfigSupplier.get(), sameInstance(optionalKafkaConnectConfigSupplier.get())); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigTest.java deleted file mode 100644 index 672cb43903..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/extension/KafkaConnectConfigTest.java +++ /dev/null @@ -1,124 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.extension; - -import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.databind.module.SimpleModule; -import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; -import org.hamcrest.CoreMatchers; -import org.junit.jupiter.api.Test; -import org.opensearch.dataprepper.model.types.ByteCount; -import org.opensearch.dataprepper.parser.ByteCountDeserializer; -import org.opensearch.dataprepper.parser.DataPrepperDurationDeserializer; -import org.opensearch.dataprepper.parser.model.DataPrepperConfiguration; -import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.AwsConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; - -import java.io.File; -import java.io.IOException; -import java.io.Reader; -import java.io.StringReader; -import java.time.Duration; -import java.util.List; -import java.util.Map; -import java.util.Properties; - -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.is; -import static org.hamcrest.Matchers.notNullValue; -import static org.hamcrest.Matchers.nullValue; - -public class KafkaConnectConfigTest { - - private static SimpleModule simpleModule = new SimpleModule() - .addDeserializer(Duration.class, new DataPrepperDurationDeserializer()) - .addDeserializer(ByteCount.class, new ByteCountDeserializer()); - private static ObjectMapper OBJECT_MAPPER = new ObjectMapper(new YAMLFactory()).registerModule(simpleModule); - - private KafkaConnectConfig makeConfig(String filePath) throws IOException { - final File configurationFile = new File(filePath); - final DataPrepperConfiguration dataPrepperConfiguration = OBJECT_MAPPER.readValue(configurationFile, DataPrepperConfiguration.class); - assertThat(dataPrepperConfiguration, CoreMatchers.notNullValue()); - assertThat(dataPrepperConfiguration.getPipelineExtensions(), CoreMatchers.notNullValue()); - final Map kafkaConnectConfigMap = (Map) dataPrepperConfiguration.getPipelineExtensions().getExtensionMap().get("kafka_connect_config"); - String json = OBJECT_MAPPER.writeValueAsString(kafkaConnectConfigMap); - Reader reader = new StringReader(json); - return OBJECT_MAPPER.readValue(reader, KafkaConnectConfig.class); - } - - @Test - public void test_config_setter_getter() throws IOException { - KafkaConnectConfig testConfig = makeConfig("src/test/resources/sample-data-prepper-config-with-kafka-connect-config-extension.yaml"); - AuthConfig authConfig = new AuthConfig(); - AwsConfig awsConfig = new AwsConfig(); - EncryptionConfig encryptionConfig = new EncryptionConfig(); - List bootstrapServer = List.of("testhost:123"); - testConfig.setAuthConfig(authConfig); - testConfig.setAwsConfig(awsConfig); - testConfig.setEncryptionConfig(encryptionConfig); - testConfig.setBootstrapServers(bootstrapServer); - assertThat(testConfig.getAuthConfig(), is(authConfig)); - assertThat(testConfig.getAwsConfig(), is(awsConfig)); - assertThat(testConfig.getEncryptionConfig(), is(encryptionConfig)); - assertThat(testConfig.getBootstrapServers(), is(bootstrapServer)); - assertThat(testConfig.getConnectorStartTimeout().getSeconds(), is(3L)); - assertThat(testConfig.getConnectStartTimeout().getSeconds(), is(3L)); - } - - @Test - public void test_config_get_worker_properties() throws IOException { - final String bootstrapServers = "localhost:9092"; - final Properties authProperties = new Properties(); - authProperties.put("bootstrap.servers", bootstrapServers); - authProperties.put("testClass", KafkaConnectConfigTest.class); - authProperties.put("testKey", "testValue"); - KafkaConnectConfig testConfig = makeConfig("src/test/resources/sample-data-prepper-config-with-kafka-connect-config-extension.yaml"); - testConfig.setAuthProperties(authProperties); - // verify WorkerProperties - assertThat(testConfig.getWorkerProperties(), notNullValue()); - Map workerProperties = testConfig.getWorkerProperties().buildKafkaConnectPropertyMap(); - assertThat(workerProperties.get("bootstrap.servers"), is(bootstrapServers)); - assertThat(workerProperties.get("group.id"), is("test-group")); - assertThat(workerProperties.get("client.id"), is("test-client")); - assertThat(workerProperties.get("offset.storage.topic"), is("test-offsets")); - assertThat(workerProperties.get("config.storage.topic"), is("test-configs")); - assertThat(workerProperties.get("status.storage.topic"), is("test-status")); - assertThat(workerProperties.get("key.converter"), is("org.apache.kafka.connect.json.JsonConverter")); - assertThat(workerProperties.get("value.converter"), is("org.apache.kafka.connect.json.JsonConverter")); - assertThat(workerProperties.get("offset.storage.partitions"), is("2")); - assertThat(workerProperties.get("offset.flush.interval.ms"), is("6000")); - assertThat(workerProperties.get("offset.flush.timeout.ms"), is("500")); - assertThat(workerProperties.get("status.storage.partitions"), is("1")); - assertThat(workerProperties.get("heartbeat.interval.ms"), is("300")); - assertThat(workerProperties.get("session.timeout.ms"), is("3000")); - assertThat(workerProperties.get("scheduled.rebalance.max.delay.ms"), is("60000")); - assertThat(workerProperties.get("testClass"), is(this.getClass().getName())); - assertThat(workerProperties.get("producer.testClass"), is(this.getClass().getName())); - assertThat(workerProperties.get("testKey"), is(authProperties.getProperty("testKey"))); - assertThat(workerProperties.get("producer.testKey"), is(authProperties.getProperty("testKey"))); - } - - @Test - public void test_config_default_worker_properties() throws IOException { - KafkaConnectConfig testConfig = makeConfig("src/test/resources/sample-data-prepper-config-with-default-kafka-connect-config-extension.yaml"); - assertThat(testConfig, notNullValue()); - assertThat(testConfig.getConnectStartTimeout().getSeconds(), is(60L)); - assertThat(testConfig.getConnectorStartTimeout().getSeconds(), is(360L)); - assertThat(testConfig.getBootstrapServers(), nullValue()); - WorkerProperties testWorkerProperties = testConfig.getWorkerProperties(); - assertThat(testWorkerProperties, notNullValue()); - Map workerProperties = testWorkerProperties.buildKafkaConnectPropertyMap(); - assertThat(workerProperties.get("bootstrap.servers"), nullValue()); - assertThat(workerProperties.get("offset.storage.partitions"), is("25")); - assertThat(workerProperties.get("offset.flush.interval.ms"), is("60000")); - assertThat(workerProperties.get("offset.flush.timeout.ms"), is("5000")); - assertThat(workerProperties.get("status.storage.partitions"), is("5")); - assertThat(workerProperties.get("heartbeat.interval.ms"), is("3000")); - assertThat(workerProperties.get("session.timeout.ms"), is("30000")); - assertThat(workerProperties.get("scheduled.rebalance.max.delay.ms"), is("300000")); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/meter/KafkaConnectMetricsTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/meter/KafkaConnectMetricsTest.java deleted file mode 100644 index 55ef90db5d..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/meter/KafkaConnectMetricsTest.java +++ /dev/null @@ -1,69 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.meter; - -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.Mock; -import org.opensearch.dataprepper.metrics.PluginMetrics; - -import javax.management.MBeanServer; -import javax.management.MalformedObjectNameException; -import javax.management.ObjectName; -import java.util.Set; - -import static java.util.Collections.emptyList; -import static java.util.Collections.emptySet; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.notNullValue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.atLeastOnce; -import static org.mockito.Mockito.lenient; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -public class KafkaConnectMetricsTest { - @Mock - private PluginMetrics pluginMetrics; - - @Mock - private MBeanServer mBeanServer; - - private Iterable tags = emptyList(); - - @BeforeEach - void setUp() throws Exception { - pluginMetrics = mock(PluginMetrics.class); - mBeanServer = mock(MBeanServer.class); - lenient().when(mBeanServer.getAttribute(any(), any())).thenReturn(1); - } - - @Test - void testConstructor() { - assertThat(new KafkaConnectMetrics(pluginMetrics), notNullValue()); - when(mBeanServer.queryNames(any(), any())).thenReturn(emptySet()); - assertThat(new KafkaConnectMetrics(pluginMetrics, tags), notNullValue()); - } - - @Test - void testBindConnectMetrics() throws MalformedObjectNameException { - final KafkaConnectMetrics kafkaConnectMetrics = new KafkaConnectMetrics(pluginMetrics, mBeanServer, tags); - when(mBeanServer.queryNames(any(), any())).thenReturn(Set.of(new ObjectName("test:*"))); - kafkaConnectMetrics.bindConnectMetrics(); - verify(mBeanServer).queryNames(any(), any()); - verify(pluginMetrics, atLeastOnce()).gaugeWithTags(any(), any(), any(), any()); - } - - @Test - void testBindConnectorMetrics() throws MalformedObjectNameException { - final KafkaConnectMetrics kafkaConnectMetrics = new KafkaConnectMetrics(pluginMetrics, mBeanServer, tags); - when(mBeanServer.queryNames(any(), any())).thenReturn(Set.of(new ObjectName("test:type=test,connector=test,client-id=test1,node-id=test1,task=task1"))); - kafkaConnectMetrics.bindConnectorMetrics(); - verify(mBeanServer).queryNames(any(), any()); - verify(pluginMetrics, atLeastOnce()).gaugeWithTags(any(), any(), any(), any()); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/KafkaConnectSourceTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/KafkaConnectSourceTest.java deleted file mode 100644 index 2e52176fab..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/KafkaConnectSourceTest.java +++ /dev/null @@ -1,180 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.source; - -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.MockedStatic; -import org.mockito.junit.jupiter.MockitoExtension; -import org.mockito.stubbing.Answer; -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.buffer.Buffer; -import org.opensearch.dataprepper.model.configuration.PipelineDescription; -import org.opensearch.dataprepper.model.record.Record; -import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.AwsConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; -import org.opensearch.dataprepper.plugins.kafka.extension.KafkaClusterConfigSupplier; -import org.opensearch.dataprepper.plugins.kafka.util.KafkaSecurityConfigurer; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.MySQLConfig; -import org.opensearch.dataprepper.plugins.kafkaconnect.extension.KafkaConnectConfig; -import org.opensearch.dataprepper.plugins.kafkaconnect.extension.KafkaConnectConfigSupplier; -import org.opensearch.dataprepper.plugins.kafkaconnect.extension.WorkerProperties; -import org.opensearch.dataprepper.plugins.kafkaconnect.util.KafkaConnect; - -import java.util.Collections; -import java.util.List; - -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.Mockito.any; -import static org.mockito.Mockito.doNothing; -import static org.mockito.Mockito.lenient; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.mockStatic; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -@ExtendWith(MockitoExtension.class) -public class KafkaConnectSourceTest { - private final String TEST_PIPELINE_NAME = "test_pipeline"; - private KafkaConnectSource kafkaConnectSource; - - @Mock - private MySQLConfig mySQLConfig; - - @Mock - private KafkaConnectConfig kafkaConnectConfig; - - @Mock - private PluginMetrics pluginMetrics; - - @Mock - private PipelineDescription pipelineDescription; - - @Mock - private Buffer> buffer; - - @Mock - private KafkaConnect kafkaConnect; - - @Mock - private KafkaClusterConfigSupplier kafkaClusterConfigSupplier; - - @Mock - private KafkaConnectConfigSupplier kafkaConnectConfigSupplier; - - private String bootstrapServers = "localhost:9092"; - - public KafkaConnectSource createSourceUnderTest() { - return new MySQLSource(mySQLConfig, pluginMetrics, pipelineDescription, kafkaClusterConfigSupplier, kafkaConnectConfigSupplier); - } - - @BeforeEach - void setUp() { - WorkerProperties workerProperties = new WorkerProperties(); - workerProperties.setBootstrapServers(bootstrapServers); - kafkaConnectConfigSupplier = mock(KafkaConnectConfigSupplier.class); - lenient().when(kafkaConnectConfigSupplier.getConfig()).thenReturn(kafkaConnectConfig); - lenient().when(kafkaConnectConfig.getWorkerProperties()).thenReturn(workerProperties); - lenient().when(mySQLConfig.buildConnectors()).thenReturn(Collections.emptyList()); - - pipelineDescription = mock(PipelineDescription.class); - lenient().when(pipelineDescription.getPipelineName()).thenReturn(TEST_PIPELINE_NAME); - pluginMetrics = mock(PluginMetrics.class); - } - - @Test - void testStartKafkaConnectSource() throws InterruptedException { - try (MockedStatic mockedStatic = mockStatic(KafkaConnect.class); - MockedStatic mockedSecurityConfigurer = mockStatic(KafkaSecurityConfigurer.class)) { - mockedSecurityConfigurer.when(() -> KafkaSecurityConfigurer.setAuthProperties(any(), any(), any())).thenAnswer((Answer) invocation -> null); - kafkaConnect = mock(KafkaConnect.class); - doNothing().when(kafkaConnect).addConnectors(any()); - doNothing().when(kafkaConnect).start(); - doNothing().when(kafkaConnect).stop(); - // Set up the mock behavior for the static method getInstance() - mockedStatic.when(() -> KafkaConnect.getPipelineInstance(any(), any(), any(), any())).thenReturn(kafkaConnect); - kafkaConnectSource = createSourceUnderTest(); - kafkaConnectSource.start(buffer); - verify(kafkaConnect).addConnectors(any()); - verify(kafkaConnect).start(); - Thread.sleep(10); - kafkaConnectSource.stop(); - verify(kafkaConnect).stop(); - } - } - - @Test - void testStartKafkaConnectSourceError() { - WorkerProperties workerProperties = new WorkerProperties(); - workerProperties.setBootstrapServers(null); - lenient().when(kafkaConnectConfig.getWorkerProperties()).thenReturn(workerProperties); - try (MockedStatic mockedStatic = mockStatic(KafkaConnect.class); - MockedStatic mockedSecurityConfigurer = mockStatic(KafkaSecurityConfigurer.class)) { - mockedSecurityConfigurer.when(() -> KafkaSecurityConfigurer.setAuthProperties(any(), any(), any())).thenAnswer((Answer) invocation -> null); - kafkaConnect = mock(KafkaConnect.class); - // Set up the mock behavior for the static method getInstance() - mockedStatic.when(() -> KafkaConnect.getPipelineInstance(any(), any(), any(), any())).thenReturn(kafkaConnect); - kafkaConnectSource = createSourceUnderTest(); - assertThrows(IllegalArgumentException.class, () -> kafkaConnectSource.start(buffer)); - } - } - - @Test - void test_updateConfig_using_kafkaClusterConfigExtension() { - final List bootstrapServers = List.of("localhost:9092"); - final AuthConfig authConfig = mock(AuthConfig.class); - final AwsConfig awsConfig = mock(AwsConfig.class); - final EncryptionConfig encryptionConfig = mock(EncryptionConfig.class); - doNothing().when(kafkaConnectConfig).setBootstrapServers(any()); - doNothing().when(kafkaConnectConfig).setAuthConfig(any()); - doNothing().when(kafkaConnectConfig).setAwsConfig(any()); - doNothing().when(kafkaConnectConfig).setEncryptionConfig(any()); - when(kafkaConnectConfig.getAuthConfig()).thenReturn(null); - when(kafkaConnectConfig.getAwsConfig()).thenReturn(null); - when(kafkaConnectConfig.getEncryptionConfig()).thenReturn(null); - when(kafkaConnectConfig.getBootstrapServers()).thenReturn(null); - when(kafkaClusterConfigSupplier.getBootStrapServers()).thenReturn(bootstrapServers); - when(kafkaClusterConfigSupplier.getAuthConfig()).thenReturn(authConfig); - when(kafkaClusterConfigSupplier.getAwsConfig()).thenReturn(awsConfig); - when(kafkaClusterConfigSupplier.getEncryptionConfig()).thenReturn(encryptionConfig); - try (MockedStatic mockedStatic = mockStatic(KafkaSecurityConfigurer.class)) { - mockedStatic.when(() -> KafkaSecurityConfigurer.setAuthProperties(any(), any(), any())).thenAnswer((Answer) invocation -> null); - kafkaConnectSource = createSourceUnderTest(); - verify(kafkaConnectConfig).setBootstrapServers(bootstrapServers); - verify(kafkaConnectConfig).setAuthConfig(authConfig); - verify(kafkaConnectConfig).setAwsConfig(awsConfig); - verify(kafkaConnectConfig).setEncryptionConfig(encryptionConfig); - } - } - - @Test - void test_updateConfig_not_using_kafkaClusterConfigExtension() { - final List bootstrapServers = List.of("localhost:9092"); - final AuthConfig authConfig = mock(AuthConfig.class); - final AwsConfig awsConfig = mock(AwsConfig.class); - final EncryptionConfig encryptionConfig = mock(EncryptionConfig.class); - lenient().doNothing().when(kafkaConnectConfig).setBootstrapServers(any()); - lenient().doNothing().when(kafkaConnectConfig).setAuthConfig(any()); - lenient().doNothing().when(kafkaConnectConfig).setAwsConfig(any()); - lenient().doNothing().when(kafkaConnectConfig).setEncryptionConfig(any()); - lenient().when(kafkaConnectConfig.getAuthConfig()).thenReturn(authConfig); - lenient().when(kafkaConnectConfig.getAwsConfig()).thenReturn(awsConfig); - lenient().when(kafkaConnectConfig.getEncryptionConfig()).thenReturn(encryptionConfig); - lenient().when(kafkaConnectConfig.getBootstrapServers()).thenReturn(bootstrapServers); - try (MockedStatic mockedStatic = mockStatic(KafkaSecurityConfigurer.class)) { - mockedStatic.when(() -> KafkaSecurityConfigurer.setAuthProperties(any(), any(), any())).thenAnswer((Answer) invocation -> null); - kafkaConnectSource = createSourceUnderTest(); - verify(kafkaConnectConfig, never()).setBootstrapServers(any()); - verify(kafkaConnectConfig, never()).setAuthConfig(any()); - verify(kafkaConnectConfig, never()).setAwsConfig(any()); - verify(kafkaConnectConfig, never()).setEncryptionConfig(any()); - } - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/MongoDBSourceTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/MongoDBSourceTest.java deleted file mode 100644 index fde37870e2..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/MongoDBSourceTest.java +++ /dev/null @@ -1,127 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.source; - -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.MockedStatic; -import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; -import org.opensearch.dataprepper.model.buffer.Buffer; -import org.opensearch.dataprepper.model.codec.ByteDecoder; -import org.opensearch.dataprepper.model.configuration.PipelineDescription; -import org.opensearch.dataprepper.model.record.Record; -import org.opensearch.dataprepper.model.source.coordinator.SourceCoordinator; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.MongoDBConfig; -import org.opensearch.dataprepper.plugins.kafkaconnect.source.mongoDB.MongoDBService; -import org.opensearch.dataprepper.plugins.kafkaconnect.source.mongoDB.MongoDBSnapshotProgressState; - -import java.util.List; - -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.CoreMatchers.instanceOf; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.doNothing; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.mockStatic; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -@ExtendWith(MockitoExtension.class) -public class MongoDBSourceTest { - @Mock - private MongoDBConfig mongoDBConfig; - - @Mock - private PluginMetrics pluginMetrics; - - @Mock - private PipelineDescription pipelineDescription; - - @Mock - private AwsCredentialsSupplier awsCredentialsSupplier; - - @Mock - private AcknowledgementSetManager acknowledgementSetManager; - - @Mock - private SourceCoordinator sourceCoordinator; - - @Mock - private MongoDBService mongoDBService; - - @Mock - private Buffer> buffer; - - @BeforeEach - void setup() { - mongoDBConfig = mock(MongoDBConfig.class); - sourceCoordinator = mock(SourceCoordinator.class); - } - - @Test - void testConstructorValidations() { - when(mongoDBConfig.getIngestionMode()).thenReturn(MongoDBConfig.IngestionMode.EXPORT_STREAM); - assertThrows(IllegalArgumentException.class, () -> new MongoDBSource( - mongoDBConfig, - pluginMetrics, - pipelineDescription, - acknowledgementSetManager, - awsCredentialsSupplier, - null, - null)); - } - - @Test - void testConstructorValidations_invalidCollectionName() { - MongoDBConfig.CollectionConfig collectionConfig = mock(MongoDBConfig.CollectionConfig.class); - when(collectionConfig.getCollectionName()).thenReturn("invalidName"); - when(mongoDBConfig.getIngestionMode()).thenReturn(MongoDBConfig.IngestionMode.EXPORT); - when(mongoDBConfig.getCollections()).thenReturn(List.of(collectionConfig)); - assertThrows(IllegalArgumentException.class, () -> new MongoDBSource( - mongoDBConfig, - pluginMetrics, - pipelineDescription, - acknowledgementSetManager, - awsCredentialsSupplier, - null, - null)); - } - - @Test - void testExportConstructor() { - when(mongoDBConfig.getIngestionMode()).thenReturn(MongoDBConfig.IngestionMode.EXPORT); - doNothing().when(sourceCoordinator).giveUpPartitions(); - MongoDBSource mongoDBSource = new MongoDBSource( - mongoDBConfig, - pluginMetrics, - pipelineDescription, - acknowledgementSetManager, - awsCredentialsSupplier, - null, - null); - mongoDBSource.setSourceCoordinator(sourceCoordinator); - assertThat(mongoDBSource.getPartitionProgressStateClass(), equalTo(MongoDBSnapshotProgressState.class)); - assertThat(mongoDBSource.getDecoder(), instanceOf(ByteDecoder.class)); - try (MockedStatic mockedStatic = mockStatic((MongoDBService.class))) { - mongoDBService = mock(MongoDBService.class); - doNothing().when(mongoDBService).start(); - doNothing().when(mongoDBService).stop(); - mockedStatic.when(() -> MongoDBService.create(any(), any(), any(), any(), any())).thenReturn(mongoDBService); - mongoDBSource.start(buffer); - verify(mongoDBService).start(); - mongoDBSource.stop(); - verify(mongoDBService).stop(); - verify(sourceCoordinator).giveUpPartitions(); - } - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/MySQLSourceTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/MySQLSourceTest.java deleted file mode 100644 index 3c0fbb0046..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/MySQLSourceTest.java +++ /dev/null @@ -1,33 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.source; - -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.configuration.PipelineDescription; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.MySQLConfig; - -import static org.junit.jupiter.api.Assertions.assertThrows; - -@ExtendWith(MockitoExtension.class) -public class MySQLSourceTest { - @Mock - private MySQLConfig mySQLConfig; - - @Mock - private PluginMetrics pluginMetrics; - - @Mock - private PipelineDescription pipelineDescription; - - @Test - void testConstructorValidations() { - assertThrows(IllegalArgumentException.class, () -> new MySQLSource(mySQLConfig, pluginMetrics, pipelineDescription, null, null)); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/PostgreSQLSourceTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/PostgreSQLSourceTest.java deleted file mode 100644 index 2cdf8973cf..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/PostgreSQLSourceTest.java +++ /dev/null @@ -1,33 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.source; - -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.configuration.PipelineDescription; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.PostgreSQLConfig; - -import static org.junit.jupiter.api.Assertions.assertThrows; - -@ExtendWith(MockitoExtension.class) -public class PostgreSQLSourceTest { - @Mock - private PostgreSQLConfig postgreSQLConfig; - - @Mock - private PluginMetrics pluginMetrics; - - @Mock - private PipelineDescription pipelineDescription; - - @Test - void testConstructorValidations() { - assertThrows(IllegalArgumentException.class, () -> new PostgreSQLSource(postgreSQLConfig, pluginMetrics, pipelineDescription, null, null)); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBPartitionCreationSupplierTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBPartitionCreationSupplierTest.java deleted file mode 100644 index c5133acf6a..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBPartitionCreationSupplierTest.java +++ /dev/null @@ -1,126 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.source.mongoDB; - -import com.mongodb.client.FindIterable; -import com.mongodb.client.MongoClient; -import com.mongodb.client.MongoClients; -import com.mongodb.client.MongoCollection; -import com.mongodb.client.MongoCursor; -import com.mongodb.client.MongoDatabase; -import org.bson.Document; -import org.bson.conversions.Bson; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.MockedStatic; -import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.model.source.coordinator.PartitionIdentifier; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.CredentialsConfig; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.MongoDBConfig; - -import java.time.Instant; -import java.util.Collections; -import java.util.HashMap; -import java.util.List; -import java.util.Map; - -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.is; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyInt; -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.lenient; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.mockStatic; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -@ExtendWith(MockitoExtension.class) -public class MongoDBPartitionCreationSupplierTest { - private static String TEST_COLLECTION_NAME = "test.collection"; - @Mock - private MongoDBConfig mongoDBConfig; - - @Mock - private MongoDBConfig.CollectionConfig collectionConfig; - - private MongoDBPartitionCreationSupplier testSupplier; - - @BeforeEach - public void setup() { - mongoDBConfig = mock(MongoDBConfig.class); - collectionConfig = mock(MongoDBConfig.CollectionConfig.class); - lenient().when(collectionConfig.getCollectionName()).thenReturn(TEST_COLLECTION_NAME); - lenient().when(mongoDBConfig.getCollections()).thenReturn(Collections.singletonList(collectionConfig)); - lenient().when(mongoDBConfig.getCredentialsConfig()).thenReturn(new CredentialsConfig(new CredentialsConfig.PlainText("user", "user"), null)); - lenient().when(mongoDBConfig.getExportConfig()).thenReturn(new MongoDBConfig.ExportConfig()); - testSupplier = new MongoDBPartitionCreationSupplier(mongoDBConfig); - } - - @Test - public void test_returnEmptyPartitionListIfAlreadyPartitioned() { - final Map globalStateMap = new HashMap<>(); - final Map partitionedCollections = new HashMap<>(); - partitionedCollections.put(TEST_COLLECTION_NAME, Instant.now().toEpochMilli()); - globalStateMap.put(MongoDBPartitionCreationSupplier.GLOBAL_STATE_PARTITIONED_COLLECTION_KEY, partitionedCollections); - List partitions = testSupplier.apply(globalStateMap); - assert (partitions.isEmpty()); - } - - @Test - public void test_returnPartitionsForCollection() { - try (MockedStatic mockedMongoClientsStatic = mockStatic(MongoClients.class)) { - // Given a collection with 5000 items which should be split to two partitions: 0-3999 and 4000-4999 - MongoClient mongoClient = mock(MongoClient.class); - MongoDatabase mongoDatabase = mock(MongoDatabase.class); - MongoCollection col = mock(MongoCollection.class); - FindIterable findIterable = mock(FindIterable.class); - MongoCursor cursor = mock(MongoCursor.class); - mockedMongoClientsStatic.when(() -> MongoClients.create(anyString())).thenReturn(mongoClient); - when(mongoClient.getDatabase(anyString())).thenReturn(mongoDatabase); - when(mongoDatabase.getCollection(anyString())).thenReturn(col); - when(col.find()).thenReturn(findIterable); - when(col.find(any(Bson.class))).thenReturn(findIterable); - when(findIterable.projection(any())).thenReturn(findIterable); - when(findIterable.sort(any())).thenReturn(findIterable); - when(findIterable.skip(anyInt())).thenReturn(findIterable); - when(findIterable.limit(anyInt())).thenReturn(findIterable); - when(findIterable.iterator()).thenReturn(cursor); - when(cursor.hasNext()).thenReturn(true, true, false); - // mock startDoc and endDoc returns, 0-3999, and 4000-4999 - when(cursor.next()) - .thenReturn(new Document("_id", "0")) - .thenReturn(new Document("_id", "4000")); - when(findIterable.first()) - .thenReturn(new Document("_id", "3999")) - .thenReturn(null) - .thenReturn(new Document("_id", "4999")); - // When Apply Partition create logics - final Map globalStateMap = new HashMap<>(); - List partitions = testSupplier.apply(globalStateMap); - // Then dependencies are called - verify(mongoClient).getDatabase(eq("test")); - verify(mongoClient, times(1)).close(); - verify(mongoDatabase).getCollection(eq("collection")); - // And partitions are created - assertThat(partitions.size(), is(2)); - assertThat(partitions.get(0).getPartitionKey(), is("test.collection|0|3999|java.lang.String")); - assertThat(partitions.get(1).getPartitionKey(), is("test.collection|4000|4999|java.lang.String")); - } - } - - @Test - public void test_returnPartitionsForCollection_error() { - when(collectionConfig.getCollectionName()).thenReturn("invalidDBName"); - final Map globalStateMap = new HashMap<>(); - assertThrows(IllegalArgumentException.class, () -> testSupplier.apply(globalStateMap)); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBServiceTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBServiceTest.java deleted file mode 100644 index a460af7d4c..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBServiceTest.java +++ /dev/null @@ -1,68 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.source.mongoDB; - -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.MockedConstruction; -import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; -import org.opensearch.dataprepper.model.buffer.Buffer; -import org.opensearch.dataprepper.model.record.Record; -import org.opensearch.dataprepper.model.source.coordinator.SourceCoordinator; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.MongoDBConfig; - -import java.util.Optional; - -import static org.mockito.Mockito.mockConstruction; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -@ExtendWith(MockitoExtension.class) -public class MongoDBServiceTest { - @Mock - private MongoDBConfig mongoDBConfig; - - @Mock - private Buffer> buffer; - - @Mock - private AcknowledgementSetManager acknowledgementSetManager; - - @Mock - private SourceCoordinator sourceCoordinator; - - @Mock - private MongoDBPartitionCreationSupplier mongoDBPartitionCreationSupplier; - - @Mock - private PluginMetrics pluginMetrics; - - @Test - public void testConstructor() { - createObjectUnderTest(); - verify(sourceCoordinator).initialize(); - } - - @Test - public void testStartAndStop() throws InterruptedException { - when(sourceCoordinator.getNextPartition(mongoDBPartitionCreationSupplier)).thenReturn(Optional.empty()); - MongoDBService testObject = createObjectUnderTest(); - testObject.start(); - Thread.sleep(100); - testObject.stop(); - } - - private MongoDBService createObjectUnderTest() { - try (final MockedConstruction mockedConstruction = mockConstruction(MongoDBPartitionCreationSupplier.class, (mock, context) -> { - mongoDBPartitionCreationSupplier = mock; - })) { - return MongoDBService.create(mongoDBConfig, sourceCoordinator, buffer, acknowledgementSetManager, pluginMetrics); - } - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBSnapshotWorkerTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBSnapshotWorkerTest.java deleted file mode 100644 index ef4179d526..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/mongoDB/MongoDBSnapshotWorkerTest.java +++ /dev/null @@ -1,245 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.source.mongoDB; - -import com.mongodb.client.FindIterable; -import com.mongodb.client.MongoClient; -import com.mongodb.client.MongoClients; -import com.mongodb.client.MongoCollection; -import com.mongodb.client.MongoCursor; -import com.mongodb.client.MongoDatabase; -import io.micrometer.core.instrument.Counter; -import org.bson.Document; -import org.bson.conversions.Bson; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.CsvSource; -import org.mockito.ArgumentCaptor; -import org.mockito.Mock; -import org.mockito.MockedStatic; -import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; -import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; -import org.opensearch.dataprepper.model.buffer.Buffer; -import org.opensearch.dataprepper.model.event.Event; -import org.opensearch.dataprepper.model.record.Record; -import org.opensearch.dataprepper.model.source.coordinator.SourceCoordinator; -import org.opensearch.dataprepper.model.source.coordinator.SourcePartition; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.CredentialsConfig; -import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.MongoDBConfig; - -import java.util.List; -import java.util.Optional; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.Future; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; -import java.util.function.Consumer; - -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.equalTo; -import static org.hamcrest.Matchers.is; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyBoolean; -import static org.mockito.ArgumentMatchers.anyInt; -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.doAnswer; -import static org.mockito.Mockito.doNothing; -import static org.mockito.Mockito.doThrow; -import static org.mockito.Mockito.lenient; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.mockStatic; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -@ExtendWith(MockitoExtension.class) -public class MongoDBSnapshotWorkerTest { - @Mock - private SourceCoordinator sourceCoordinator; - @Mock - private Buffer> buffer; - @Mock - private MongoDBPartitionCreationSupplier mongoDBPartitionCreationSupplier; - @Mock - private PluginMetrics pluginMetrics; - @Mock - private AcknowledgementSetManager acknowledgementSetManager; - @Mock - private MongoDBConfig mongoDBConfig; - @Mock - private SourcePartition sourcePartition; - @Mock - private Counter counter; - private MongoDBSnapshotWorker testWorker; - private ExecutorService executorService; - - - @BeforeEach - public void setup() throws TimeoutException { - lenient().when(mongoDBConfig.getExportConfig()).thenReturn(new MongoDBConfig.ExportConfig()); - lenient().when(mongoDBConfig.getCredentialsConfig()).thenReturn(new CredentialsConfig(new CredentialsConfig.PlainText("user", "user"), null)); - lenient().when(buffer.isByteBuffer()).thenReturn(false); - lenient().doNothing().when(buffer).write(any(), anyInt()); - lenient().doNothing().when(sourceCoordinator).saveProgressStateForPartition(anyString(), any()); - lenient().when(pluginMetrics.counter(anyString())).thenReturn(counter); - executorService = Executors.newSingleThreadExecutor(); - testWorker = new MongoDBSnapshotWorker(sourceCoordinator, buffer, mongoDBPartitionCreationSupplier, pluginMetrics, acknowledgementSetManager, mongoDBConfig); - } - - @Test - public void test_shouldSleepIfNoPartitionRetrieved() throws InterruptedException { - when(sourceCoordinator.getNextPartition(mongoDBPartitionCreationSupplier)).thenReturn(Optional.empty()); - final Future future = executorService.submit(() -> testWorker.run()); - Thread.sleep(100); - executorService.shutdown(); - future.cancel(true); - assertThat(future.isCancelled(), equalTo(true)); - assertThat(executorService.awaitTermination(100, TimeUnit.MILLISECONDS), equalTo(true)); - } - - @ParameterizedTest - @CsvSource({ - "test.collection|0|1|java.lang.Integer", - "test.collection|0|1|java.lang.Double", - "test.collection|0|1|java.lang.String", - "test.collection|0|1|java.lang.Long", - "test.collection|000000000000000000000000|000000000000000000000001|org.bson.types.ObjectId" - }) - public void test_shouldProcessPartitionSuccess(final String partitionKey) throws InterruptedException, TimeoutException { - this.mockDependencyAndProcessPartition(partitionKey, true); - - final ArgumentCaptor> ingestDataCapture = ArgumentCaptor.forClass(Record.class); - verify(buffer, times(2)).write(ingestDataCapture.capture(), anyInt()); - List> capturedData = ingestDataCapture.getAllValues(); - String data1 = ((Event) capturedData.get(0).getData()).jsonBuilder().includeTags(null).toJsonString(); - String data2 = ((Event) capturedData.get(1).getData()).jsonBuilder().includeTags(null).toJsonString(); - assertThat(data1, is("{\"_id\":0,\"__source_db\":\"test\",\"__collection\":\"collection\",\"__op\":\"create\",\"__source_ts_ms\":0}")); - assertThat(data2, is("{\"_id\":1,\"__source_db\":\"test\",\"__collection\":\"collection\",\"__op\":\"create\",\"__source_ts_ms\":0}")); - } - - @Test - public void test_shouldProcessPartitionSuccess_byteBuffer() throws Exception { - when(buffer.isByteBuffer()).thenReturn(true); - doNothing().when(buffer).writeBytes(any(byte[].class), any(), anyInt()); - this.mockDependencyAndProcessPartition("test.collection|0|1|java.lang.Integer", true); - - final ArgumentCaptor ingestDataCapture = ArgumentCaptor.forClass(byte[].class); - verify(buffer, times(2)).writeBytes(ingestDataCapture.capture(), any(), anyInt()); - List capturedData = ingestDataCapture.getAllValues(); - String data1 = new String(capturedData.get(0)); - String data2 = new String(capturedData.get(1)); - assertThat(data1, is("{\"_id\":0,\"__source_db\":\"test\",\"__collection\":\"collection\",\"__op\":\"create\",\"__source_ts_ms\":0}")); - assertThat(data2, is("{\"_id\":1,\"__source_db\":\"test\",\"__collection\":\"collection\",\"__op\":\"create\",\"__source_ts_ms\":0}")); - } - - @Test - public void test_shouldProcessPartitionSuccess_ackEnabled() throws InterruptedException, TimeoutException { - MongoDBConfig.ExportConfig exportConfig = mock(MongoDBConfig.ExportConfig.class); - when(exportConfig.getAcknowledgements()).thenReturn(true); - when(mongoDBConfig.getExportConfig()).thenReturn(exportConfig); - final AcknowledgementSet acknowledgementSet = mock(AcknowledgementSet.class); - doAnswer(invocation -> { - Consumer consumer = invocation.getArgument(0); - consumer.accept(true); - return acknowledgementSet; - }).when(acknowledgementSetManager).create(any(Consumer.class), any()); - doNothing().when(sourceCoordinator).updatePartitionForAcknowledgmentWait(anyString(), any()); - - this.mockDependencyAndProcessPartition("test.collection|0|1|java.lang.Integer", true); - - final ArgumentCaptor> ingestDataCapture = ArgumentCaptor.forClass(Record.class); - verify(buffer, times(2)).write(ingestDataCapture.capture(), anyInt()); - List> capturedData = ingestDataCapture.getAllValues(); - String data1 = ((Event) capturedData.get(0).getData()).jsonBuilder().includeTags(null).toJsonString(); - String data2 = ((Event) capturedData.get(1).getData()).jsonBuilder().includeTags(null).toJsonString(); - assertThat(data1, is("{\"_id\":0,\"__source_db\":\"test\",\"__collection\":\"collection\",\"__op\":\"create\",\"__source_ts_ms\":0}")); - assertThat(data2, is("{\"_id\":1,\"__source_db\":\"test\",\"__collection\":\"collection\",\"__op\":\"create\",\"__source_ts_ms\":0}")); - } - - @Test - public void test_shouldGiveUpPartitionIfExceptionOccurred() throws InterruptedException { - doNothing().when(sourceCoordinator).giveUpPartitions(); - this.mockDependencyAndProcessPartition("invalidPartition", false); - verify(sourceCoordinator, times(1)).giveUpPartitions(); - } - - @Test - public void test_shouldCountFailureIfBufferFailed() throws Exception { - doThrow(new RuntimeException("")).when(buffer).write(any(), anyInt()); - this.mockDependencyAndProcessPartition("test.collection|0|1|java.lang.Integer", false); - final ArgumentCaptor progressStateCapture = ArgumentCaptor.forClass(MongoDBSnapshotProgressState.class); - verify(sourceCoordinator, times(1)).saveProgressStateForPartition(anyString(), progressStateCapture.capture()); - List progressStates = progressStateCapture.getAllValues(); - assertThat(progressStates.get(0).getTotal(), is(2L)); - assertThat(progressStates.get(0).getSuccess(), is(0L)); - assertThat(progressStates.get(0).getFailed(), is(2L)); - } - - @Test - public void test_shouldThreadSleepIfExceptionOccurred() throws InterruptedException { - doThrow(new RuntimeException("")).when(sourceCoordinator).getNextPartition(mongoDBPartitionCreationSupplier); - final Future future = executorService.submit(() -> testWorker.run()); - Thread.sleep(100); - executorService.shutdown(); - future.cancel(true); - assertThat(future.isCancelled(), equalTo(true)); - assertThat(executorService.awaitTermination(100, TimeUnit.MILLISECONDS), equalTo(true)); - } - - private void mockDependencyAndProcessPartition(String partitionKey, boolean shouldProcessSucceed) throws InterruptedException { - lenient().when(sourcePartition.getPartitionKey()).thenReturn(partitionKey); - lenient().doNothing().when(sourceCoordinator).completePartition(anyString(), anyBoolean()); - lenient().when(sourceCoordinator.getNextPartition(mongoDBPartitionCreationSupplier)) - .thenReturn(Optional.of(sourcePartition)) - .thenReturn(Optional.empty()); - - MongoClient mongoClient = mock(MongoClient.class); - MongoDatabase mongoDatabase = mock(MongoDatabase.class); - MongoCollection col = mock(MongoCollection.class); - FindIterable findIterable = mock(FindIterable.class); - MongoCursor cursor = mock(MongoCursor.class); - lenient().when(mongoClient.getDatabase(anyString())).thenReturn(mongoDatabase); - lenient().when(mongoDatabase.getCollection(anyString())).thenReturn(col); - lenient().when(col.find(any(Bson.class))).thenReturn(findIterable); - lenient().when(findIterable.iterator()).thenReturn(cursor); - lenient().when(cursor.hasNext()).thenReturn(true, true, false); - lenient().when(cursor.next()) - .thenReturn(new Document("_id", 0)) - .thenReturn(new Document("_id", 1)); - - final Future future = executorService.submit(() -> { - try (MockedStatic mockedMongoClientsStatic = mockStatic(MongoClients.class)) { - mockedMongoClientsStatic.when(() -> MongoClients.create(anyString())).thenReturn(mongoClient); - testWorker.run(); - } - }); - Thread.sleep(1000); - executorService.shutdown(); - future.cancel(true); - assertThat(future.isCancelled(), equalTo(true)); - assertThat(executorService.awaitTermination(1000, TimeUnit.MILLISECONDS), equalTo(true)); - if (shouldProcessSucceed) { - // Verify Results - verify(cursor, times(2)).next(); - - final ArgumentCaptor progressStateCapture = ArgumentCaptor.forClass(MongoDBSnapshotProgressState.class); - verify(sourceCoordinator, times(1)).saveProgressStateForPartition(eq(partitionKey), progressStateCapture.capture()); - List progressStates = progressStateCapture.getAllValues(); - assertThat(progressStates.get(0).getTotal(), is(2L)); - assertThat(progressStates.get(0).getSuccess(), is(2L)); - assertThat(progressStates.get(0).getFailed(), is(0L)); - - verify(mongoClient, times(1)).close(); - } - } -} - diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/ConnectorTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/ConnectorTest.java deleted file mode 100644 index f8d377ae37..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/ConnectorTest.java +++ /dev/null @@ -1,28 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.util; - -import org.junit.jupiter.api.Test; - -import java.util.HashMap; -import java.util.Map; - -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.is; - -public class ConnectorTest { - @Test - void testGettersOfConnector() { - final String name = "connectorName"; - final Boolean allowReplace = false; - final Map config = new HashMap<>(); - final Connector connector = new Connector(name, config, allowReplace); - assertThat(connector.getName(), is(name)); - assertThat(connector.getConfig(), is(config)); - assertThat(connector.getConfig().get("name"), is(name)); - assertThat(connector.getAllowReplace(), is(allowReplace)); - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/KafkaConnectTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/KafkaConnectTest.java deleted file mode 100644 index eb13027378..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/KafkaConnectTest.java +++ /dev/null @@ -1,342 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.util; - -import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy; -import org.apache.kafka.connect.errors.AlreadyExistsException; -import org.apache.kafka.connect.errors.NotFoundException; -import org.apache.kafka.connect.runtime.Connect; -import org.apache.kafka.connect.runtime.Herder; -import org.apache.kafka.connect.runtime.Worker; -import org.apache.kafka.connect.runtime.WorkerConfigTransformer; -import org.apache.kafka.connect.runtime.distributed.DistributedConfig; -import org.apache.kafka.connect.runtime.distributed.DistributedHerder; -import org.apache.kafka.connect.runtime.distributed.NotLeaderException; -import org.apache.kafka.connect.runtime.isolation.Plugins; -import org.apache.kafka.connect.runtime.rest.ConnectRestServer; -import org.apache.kafka.connect.runtime.rest.RestClient; -import org.apache.kafka.connect.runtime.rest.RestServer; -import org.apache.kafka.connect.runtime.rest.entities.ConnectorInfo; -import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo; -import org.apache.kafka.connect.runtime.rest.entities.ConnectorType; -import org.apache.kafka.connect.storage.Converter; -import org.apache.kafka.connect.storage.KafkaOffsetBackingStore; -import org.apache.kafka.connect.util.Callback; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.Mock; -import org.mockito.MockedConstruction; -import org.mockito.MockedStatic; -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.plugins.kafkaconnect.extension.WorkerProperties; -import org.opensearch.dataprepper.plugins.kafkaconnect.meter.KafkaConnectMetrics; - -import java.net.URI; -import java.time.Clock; -import java.time.Duration; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; - -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.is; -import static org.hamcrest.Matchers.not; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.Mockito.any; -import static org.mockito.Mockito.doAnswer; -import static org.mockito.Mockito.doNothing; -import static org.mockito.Mockito.doThrow; -import static org.mockito.Mockito.eq; -import static org.mockito.Mockito.lenient; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.mockConstruction; -import static org.mockito.Mockito.mockStatic; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -public class KafkaConnectTest { - private static final String TEST_PIPELINE_NAME = "test"; - private static final WorkerProperties DEFAULT_WORDER_PROPERTY = new WorkerProperties(); - private static final long TEST_CONNECTOR_TIMEOUT_MS = 360000L; // 360 seconds - private static final long TEST_CONNECT_TIMEOUT_MS = 60000L; // 60 seconds - private static final Duration TEST_CONNECTOR_TIMEOUT = Duration.ofMillis(TEST_CONNECTOR_TIMEOUT_MS); - private static final Duration TEST_CONNECT_TIMEOUT = Duration.ofMillis(TEST_CONNECT_TIMEOUT_MS); - @Mock - private KafkaConnectMetrics kafkaConnectMetrics; - - @Mock - private PluginMetrics pluginMetrics; - - @Mock - private DistributedHerder distributedHerder; - - @Mock - private RestServer rest; - - @Mock - private Connect connect; - - - @BeforeEach - void setUp() throws Exception { - kafkaConnectMetrics = mock(KafkaConnectMetrics.class); - distributedHerder = mock(DistributedHerder.class); - rest = mock(RestServer.class); - connect = mock(Connect.class); - DEFAULT_WORDER_PROPERTY.setBootstrapServers("localhost:9002"); - - lenient().when(connect.isRunning()).thenReturn(false).thenReturn(true); - lenient().when(distributedHerder.connectors()).thenReturn(new ArrayList<>()); - ConnectorStateInfo runningState = new ConnectorStateInfo("newConnector", new ConnectorStateInfo.ConnectorState("RUNNING", "worker", "msg"), new ArrayList<>(), ConnectorType.SOURCE); - lenient().when(distributedHerder.connectorStatus(any())).thenReturn(runningState); - lenient().doAnswer(invocation -> { - Callback> callback = invocation.getArgument(1); - // Simulate a successful completion - callback.onCompletion(null, null); - return null; - }).when(distributedHerder).connectorConfig(any(), any(Callback.class)); - lenient().doAnswer(invocation -> { - Callback> callback = invocation.getArgument(3); - // Simulate a successful completion - callback.onCompletion(null, null); - return null; - }).when(distributedHerder).putConnectorConfig(any(String.class), any(Map.class), any(Boolean.class), any(Callback.class)); - lenient().doAnswer(invocation -> { - Callback> callback = invocation.getArgument(1); - // Simulate a successful completion - callback.onCompletion(null, null); - return null; - }).when(distributedHerder).deleteConnectorConfig(any(), any(Callback.class)); - } - - @Test - void testInitializeKafkaConnectWithSingletonForSamePipeline() { - final KafkaConnect kafkaConnect = KafkaConnect.getPipelineInstance(TEST_PIPELINE_NAME, pluginMetrics, TEST_CONNECT_TIMEOUT, TEST_CONNECTOR_TIMEOUT); - final KafkaConnect sameConnect = KafkaConnect.getPipelineInstance(TEST_PIPELINE_NAME, pluginMetrics, TEST_CONNECT_TIMEOUT, TEST_CONNECTOR_TIMEOUT); - assertThat(sameConnect, is(kafkaConnect)); - final String anotherPipeline = "anotherPipeline"; - final KafkaConnect anotherKafkaConnect = KafkaConnect.getPipelineInstance(anotherPipeline, pluginMetrics, TEST_CONNECT_TIMEOUT, TEST_CONNECTOR_TIMEOUT); - assertThat(anotherKafkaConnect, not(kafkaConnect)); - } - - @Test - void testInitializeKafkaConnect() { - Map workerProps = DEFAULT_WORDER_PROPERTY.buildKafkaConnectPropertyMap(); - try (MockedConstruction mockedConfig = mockConstruction(DistributedConfig.class, (mock, context) -> { - when(mock.kafkaClusterId()).thenReturn("test-cluster-id"); - when(mock.getString(any())).thenReturn("test-string"); - }); - MockedConstruction mockedRestClient = mockConstruction(RestClient.class); - MockedConstruction mockedHerder = mockConstruction(DistributedHerder.class); - MockedConstruction mockedRestServer = mockConstruction(ConnectRestServer.class, (mock, context) -> { - when(mock.advertisedUrl()).thenReturn(URI.create("localhost:9002")); - }); - MockedConstruction mockedPlugin = mockConstruction(Plugins.class, (mock, context) -> { - ClassLoader classLoader = mock(ClassLoader.class); - ConnectorClientConfigOverridePolicy connectorPolicy = mock(ConnectorClientConfigOverridePolicy.class); - when(mock.compareAndSwapWithDelegatingLoader()).thenReturn(classLoader); - when(mock.newPlugin(any(), any(), any())).thenReturn(connectorPolicy); - }); - MockedConstruction mockedWorker = mockConstruction(Worker.class, (mock, context) -> { - WorkerConfigTransformer configTransformer = mock(WorkerConfigTransformer.class); - Converter converter = mock(Converter.class); - when(mock.configTransformer()).thenReturn(configTransformer); - when(mock.getInternalValueConverter()).thenReturn(converter); - }); - MockedConstruction mockedOffsetStore = mockConstruction(KafkaOffsetBackingStore.class, (mock, context) -> { - doNothing().when(mock).configure(any()); - }) - ) { - final KafkaConnect kafkaConnect = KafkaConnect.getPipelineInstance(TEST_PIPELINE_NAME, pluginMetrics, TEST_CONNECT_TIMEOUT, TEST_CONNECTOR_TIMEOUT); - kafkaConnect.initialize(workerProps); - } - } - - @Test - void testStartKafkaConnectSuccess() { - final KafkaConnect kafkaConnect = new KafkaConnect(distributedHerder, rest, connect, kafkaConnectMetrics); - doNothing().when(rest).initializeServer(); - doNothing().when(connect).start(); - kafkaConnect.start(); - verify(rest).initializeServer(); - verify(connect).start(); - } - - @Test - void testStartKafkaConnectFail() { - final KafkaConnect kafkaConnect = new KafkaConnect(distributedHerder, rest, connect, kafkaConnectMetrics); - doNothing().when(rest).initializeServer(); - doThrow(new RuntimeException()).when(connect).start(); - doNothing().when(connect).stop(); - assertThrows(RuntimeException.class, kafkaConnect::start); - verify(connect, times(1)).stop(); - - // throw exception immediately if connect is null - final KafkaConnect kafkaConnect2 = new KafkaConnect(distributedHerder, rest, null, kafkaConnectMetrics); - assertThrows(RuntimeException.class, kafkaConnect2::start); - } - - @Test - void testStartKafkaConnectFailTimeout() { - doNothing().when(rest).initializeServer(); - doNothing().when(connect).start(); - doNothing().when(connect).stop(); - when(connect.isRunning()).thenReturn(false); - try (MockedStatic mockedStatic = mockStatic(Clock.class)) { - final Clock clock = mock(Clock.class); - mockedStatic.when(() -> Clock.systemUTC()).thenReturn(clock); - when(clock.millis()).thenReturn(0L).thenReturn(TEST_CONNECT_TIMEOUT_MS + 1); - final KafkaConnect kafkaConnect = new KafkaConnect(distributedHerder, rest, connect, kafkaConnectMetrics); - assertThrows(RuntimeException.class, kafkaConnect::start); - verify(rest).initializeServer(); - verify(connect).start(); - verify(connect).stop(); - verify(clock, times(2)).millis(); - } - } - - @Test - void testStartKafkaConnectWithConnectRunningAlready() { - final KafkaConnect kafkaConnect = new KafkaConnect(distributedHerder, rest, connect, kafkaConnectMetrics); - when(connect.isRunning()).thenReturn(true); - kafkaConnect.start(); - verify(rest, never()).initializeServer(); - verify(connect, never()).start(); - } - - @Test - void testStopKafkaConnect() { - final KafkaConnect kafkaConnect = new KafkaConnect(distributedHerder, rest, connect, kafkaConnectMetrics); - kafkaConnect.stop(); - verify(connect).stop(); - // should ignore stop if connect is null - final KafkaConnect kafkaConnect2 = new KafkaConnect(distributedHerder, rest, null, kafkaConnectMetrics); - kafkaConnect2.stop(); - } - - @Test - void testInitConnectorsWhenStartKafkaConnectSuccess() { - final String oldConnectorName = "oldConnector"; - final Connector newConnector = mock(Connector.class); - final String newConnectorName = "newConnector"; - final Map newConnectorConfig = new HashMap<>(); - when(newConnector.getName()).thenReturn(newConnectorName); - when(newConnector.getConfig()).thenReturn(newConnectorConfig); - when(distributedHerder.connectors()).thenReturn(List.of(oldConnectorName)); - - final KafkaConnect kafkaConnect = new KafkaConnect(distributedHerder, rest, connect, kafkaConnectMetrics); - kafkaConnect.addConnectors(List.of(newConnector)); - kafkaConnect.start(); - verify(distributedHerder).putConnectorConfig(eq(newConnectorName), eq(newConnectorConfig), eq(true), any(Callback.class)); - verify(distributedHerder).deleteConnectorConfig(eq(oldConnectorName), any(Callback.class)); - } - - @Test - void testInitConnectorsWithoutConnectorConfigChange() { - final Connector newConnector = mock(Connector.class); - final String newConnectorName = "newConnector"; - final Map newConnectorConfig = new HashMap<>(); - when(newConnector.getName()).thenReturn(newConnectorName); - when(newConnector.getConfig()).thenReturn(newConnectorConfig); - when(newConnector.getAllowReplace()).thenReturn(false); - doAnswer(invocation -> { - Callback> callback = invocation.getArgument(1); - // Simulate a successful completion - callback.onCompletion(null, newConnectorConfig); - return null; - }).when(distributedHerder).connectorConfig(any(), any(Callback.class)); - - final KafkaConnect kafkaConnect = new KafkaConnect(distributedHerder, rest, connect, kafkaConnectMetrics); - kafkaConnect.addConnectors(List.of(newConnector)); - kafkaConnect.start(); - verify(distributedHerder).putConnectorConfig(eq(newConnectorName), eq(newConnectorConfig), eq(false), any(Callback.class)); - } - - @Test - void testInitConnectorsErrorsWhenDeleteConnector() { - final String oldConnectorName = "oldConnector"; - when(distributedHerder.connectors()).thenReturn(List.of(oldConnectorName)); - final KafkaConnect kafkaConnect = new KafkaConnect(distributedHerder, rest, connect, kafkaConnectMetrics); - doAnswer(invocation -> { - Callback> callback = invocation.getArgument(1); - // Simulate a successful completion - callback.onCompletion(new RuntimeException(), null); - return null; - }).when(distributedHerder).deleteConnectorConfig(eq(oldConnectorName), any(Callback.class)); - assertThrows(RuntimeException.class, kafkaConnect::start); - // NotLeaderException or NotFoundException should be ignored. - doAnswer(invocation -> { - Callback> callback = invocation.getArgument(1); - callback.onCompletion(new NotLeaderException("Only Leader can delete.", "leaderUrl"), null); - return null; - }).when(distributedHerder).deleteConnectorConfig(eq(oldConnectorName), any(Callback.class)); - kafkaConnect.start(); - doAnswer(invocation -> { - Callback> callback = invocation.getArgument(1); - // Simulate a successful completion - callback.onCompletion(new NotFoundException("Not Found"), null); - return null; - }).when(distributedHerder).deleteConnectorConfig(eq(oldConnectorName), any(Callback.class)); - kafkaConnect.start(); - } - - @Test - void testInitConnectorsErrorsWhenPutConnector() { - final Connector newConnector = mock(Connector.class); - final String newConnectorName = "newConnector"; - final Map newConnectorConfig = new HashMap<>(); - when(newConnector.getName()).thenReturn(newConnectorName); - when(newConnector.getConfig()).thenReturn(newConnectorConfig); - final KafkaConnect kafkaConnect = new KafkaConnect(distributedHerder, rest, connect, kafkaConnectMetrics); - kafkaConnect.addConnectors(List.of(newConnector)); - // RuntimeException should be thrown - doAnswer(invocation -> { - Callback> callback = invocation.getArgument(3); - callback.onCompletion(new RuntimeException(), null); - return null; - }).when(distributedHerder).putConnectorConfig(eq(newConnectorName), eq(newConnectorConfig), eq(true), any(Callback.class)); - assertThrows(RuntimeException.class, kafkaConnect::start); - // NotLeaderException or NotFoundException should be ignored. - doAnswer(invocation -> { - Callback> callback = invocation.getArgument(3); - callback.onCompletion(new NotLeaderException("not leader", "leaderUrl"), null); - return null; - }).when(distributedHerder).putConnectorConfig(eq(newConnectorName), eq(newConnectorConfig), eq(true), any(Callback.class)); - kafkaConnect.start(); - doAnswer(invocation -> { - Callback> callback = invocation.getArgument(3); - callback.onCompletion(new AlreadyExistsException("Already added"), null); - return null; - }).when(distributedHerder).putConnectorConfig(eq(newConnectorName), eq(newConnectorConfig), eq(true), any(Callback.class)); - kafkaConnect.start(); - } - - @Test - void testInitConnectorsErrorsWhenConnectorsNotRunning() { - // should throw exception if connector failed in Running state for 30 seconds - final Connector newConnector = mock(Connector.class); - final String newConnectorName = "newConnector"; - final Map newConnectorConfig = new HashMap<>(); - when(newConnector.getName()).thenReturn(newConnectorName); - when(newConnector.getConfig()).thenReturn(newConnectorConfig); - when(distributedHerder.connectorStatus(eq(newConnectorName))).thenReturn(null); - - try (MockedStatic mockedStatic = mockStatic(Clock.class)) { - final Clock clock = mock(Clock.class); - mockedStatic.when(() -> Clock.systemUTC()).thenReturn(clock); - when(clock.millis()).thenReturn(0L).thenReturn(0L).thenReturn(0L).thenReturn(0L).thenReturn(TEST_CONNECTOR_TIMEOUT_MS + 1); - final KafkaConnect kafkaConnect = new KafkaConnect(distributedHerder, rest, connect, kafkaConnectMetrics); - kafkaConnect.addConnectors(List.of(newConnector)); - assertThrows(RuntimeException.class, kafkaConnect::start); - verify(distributedHerder, times(1)).connectorStatus(any()); - verify(distributedHerder).putConnectorConfig(eq(newConnectorName), eq(newConnectorConfig), eq(true), any(Callback.class)); - verify(clock, times(5)).millis(); - } - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/SecretManagerHelperTest.java b/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/SecretManagerHelperTest.java deleted file mode 100644 index 207d651108..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/util/SecretManagerHelperTest.java +++ /dev/null @@ -1,85 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.kafkaconnect.util; - -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.Mock; -import org.mockito.MockedStatic; -import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; -import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration; -import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient; -import software.amazon.awssdk.services.secretsmanager.SecretsManagerClientBuilder; -import software.amazon.awssdk.services.secretsmanager.model.GetSecretValueRequest; -import software.amazon.awssdk.services.secretsmanager.model.GetSecretValueResponse; -import software.amazon.awssdk.services.sts.StsClient; -import software.amazon.awssdk.services.sts.StsClientBuilder; -import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider; - -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.is; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.lenient; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.mockStatic; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -public class SecretManagerHelperTest { - private final String expectedSecretString = "expectedSecret"; - private final String testStsRole = "testRole"; - private final String testRegion = "testRegion"; - private final String testSecretId = "testSecritId"; - @Mock - private SecretsManagerClientBuilder secretsManagerClientBuilder; - @Mock - private SecretsManagerClient secretsManagerClient; - @Mock - private GetSecretValueResponse getSecretValueResponse; - - @BeforeEach - void setup() { - secretsManagerClientBuilder = mock(SecretsManagerClientBuilder.class); - secretsManagerClient = mock(SecretsManagerClient.class); - getSecretValueResponse = mock(GetSecretValueResponse.class); - lenient().when(secretsManagerClientBuilder.overrideConfiguration(any(ClientOverrideConfiguration.class))).thenReturn(secretsManagerClientBuilder); - lenient().when(secretsManagerClientBuilder.credentialsProvider(any(AwsCredentialsProvider.class))).thenReturn(secretsManagerClientBuilder); - lenient().when(secretsManagerClientBuilder.region(any())).thenReturn(secretsManagerClientBuilder); - lenient().when(secretsManagerClientBuilder.build()).thenReturn(secretsManagerClient); - lenient().when(secretsManagerClient.getSecretValue(any(GetSecretValueRequest.class))).thenReturn(getSecretValueResponse); - lenient().when(getSecretValueResponse.secretString()).thenReturn(expectedSecretString); - } - - @Test - void test_get_secret_without_sts() { - try (MockedStatic mockedStatic = mockStatic(SecretsManagerClient.class)) { - mockedStatic.when(() -> SecretsManagerClient.builder()).thenReturn(secretsManagerClientBuilder); - String result = SecretManagerHelper.getSecretValue("", testRegion, testSecretId); - assertThat(result, is(expectedSecretString)); - verify(secretsManagerClientBuilder, times(1)).credentialsProvider(any(AwsCredentialsProvider.class)); - } - } - - @Test - void test_get_secret_with_sts() { - try (MockedStatic mockedSts = mockStatic(StsClient.class); - MockedStatic mockedStatic = mockStatic(SecretsManagerClient.class)) { - StsClient stsClient = mock(StsClient.class); - StsClientBuilder stsClientBuilder = mock(StsClientBuilder.class); - when(stsClientBuilder.overrideConfiguration(any(ClientOverrideConfiguration.class))).thenReturn(stsClientBuilder); - when(stsClientBuilder.credentialsProvider(any(AwsCredentialsProvider.class))).thenReturn(stsClientBuilder); - when(stsClientBuilder.region(any())).thenReturn(stsClientBuilder); - when(stsClientBuilder.build()).thenReturn(stsClient); - - mockedSts.when(() -> StsClient.builder()).thenReturn(stsClientBuilder); - mockedStatic.when(() -> SecretsManagerClient.builder()).thenReturn(secretsManagerClientBuilder); - String result = SecretManagerHelper.getSecretValue(testStsRole, testRegion, testSecretId); - assertThat(result, is(expectedSecretString)); - verify(secretsManagerClientBuilder, times(1)).credentialsProvider(any(StsAssumeRoleCredentialsProvider.class)); - } - } -} diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/resources/sample-data-prepper-config-with-default-kafka-connect-config-extension.yaml b/data-prepper-plugins/kafka-connect-plugins/src/test/resources/sample-data-prepper-config-with-default-kafka-connect-config-extension.yaml deleted file mode 100644 index 753d4f9f5d..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/resources/sample-data-prepper-config-with-default-kafka-connect-config-extension.yaml +++ /dev/null @@ -1,8 +0,0 @@ -extensions: - kafka_connect_config: - worker_properties: - group_id: test-group - client_id: test-client - config_storage_topic: test-configs - offset_storage_topic: test-offsets - status_storage_topic: test-status \ No newline at end of file diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/resources/sample-data-prepper-config-with-kafka-connect-config-extension.yaml b/data-prepper-plugins/kafka-connect-plugins/src/test/resources/sample-data-prepper-config-with-kafka-connect-config-extension.yaml deleted file mode 100644 index e41c7d04fe..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/resources/sample-data-prepper-config-with-kafka-connect-config-extension.yaml +++ /dev/null @@ -1,19 +0,0 @@ -extensions: - kafka_connect_config: - bootstrap_servers: - - test:123 - connect_start_timeout: 3000ms - connector_start_timeout: 3s - worker_properties: - group_id: test-group - client_id: test-client - config_storage_topic: test-configs - offset_storage_topic: test-offsets - status_storage_topic: test-status - offset_storage_partitions: 2 #optional and default is 25 - offset_flush_interval: 6s #optional and default is 60000 (60s) - offset_flush_timeout: 500ms #optional and default is 5000 (5s) - status_storage_partitions: 1 #optional and default is 5 - heartbeat_interval: 300ms #optional and default is 3000 (3s) - session_timeout: 3s #optional and default is 30000 (30s) - connector_rebalance_max_delay: 60s \ No newline at end of file diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/resources/sample-mongodb-pipeline.yaml b/data-prepper-plugins/kafka-connect-plugins/src/test/resources/sample-mongodb-pipeline.yaml deleted file mode 100644 index 0b10fe2891..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/resources/sample-mongodb-pipeline.yaml +++ /dev/null @@ -1,14 +0,0 @@ -log-pipeline: - source: - mongodb: - hostname: localhost - ingestion_mode: export_stream - credentials: - plaintext: - username: debezium - password: dbz - collections: - - topic_prefix: prefix1 - collection: test.customers - sink: - - noop: \ No newline at end of file diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/resources/sample-mysql-pipeline.yaml b/data-prepper-plugins/kafka-connect-plugins/src/test/resources/sample-mysql-pipeline.yaml deleted file mode 100644 index bdbaeff015..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/resources/sample-mysql-pipeline.yaml +++ /dev/null @@ -1,13 +0,0 @@ -log-pipeline: - source: - mysql: - hostname: localhost - credentials: - plaintext: - username: debezium - password: dbz - tables: - - topic_prefix: prefix1 - table: inventory.customers - sink: - - noop: \ No newline at end of file diff --git a/data-prepper-plugins/kafka-connect-plugins/src/test/resources/sample-postgres-pipeline.yaml b/data-prepper-plugins/kafka-connect-plugins/src/test/resources/sample-postgres-pipeline.yaml deleted file mode 100644 index c843e3ccff..0000000000 --- a/data-prepper-plugins/kafka-connect-plugins/src/test/resources/sample-postgres-pipeline.yaml +++ /dev/null @@ -1,14 +0,0 @@ -log-pipeline: - source: - postgresql: - hostname: localhost - credentials: - plaintext: - username: debezium - password: dbz - tables: - - topic_prefix: psql - database: postgres - table: public.customers - sink: - - noop: \ No newline at end of file diff --git a/data-prepper-plugins/kafka-plugins/build.gradle b/data-prepper-plugins/kafka-plugins/build.gradle index 96d8cc1c64..0ef1eb572d 100644 --- a/data-prepper-plugins/kafka-plugins/build.gradle +++ b/data-prepper-plugins/kafka-plugins/build.gradle @@ -30,26 +30,22 @@ dependencies { implementation project(':data-prepper-plugins:aws-plugin-api') implementation 'org.apache.kafka:kafka-clients:3.6.1' implementation 'org.apache.kafka:connect-json:3.6.1' + implementation project(':data-prepper-plugins:http-common') implementation libs.avro.core implementation 'com.fasterxml.jackson.core:jackson-databind' implementation 'io.micrometer:micrometer-core' implementation libs.commons.lang3 implementation 'io.confluent:kafka-avro-serializer:7.4.0' + implementation 'io.confluent:kafka-json-schema-serializer:7.4.0' implementation 'io.confluent:kafka-schema-registry-client:7.4.0' - implementation ('io.confluent:kafka-schema-registry:7.4.0:tests') { - exclude group: 'org.glassfish.jersey.containers', module: 'jersey-container-servlet' - exclude group: 'org.glassfish.jersey.inject', module: 'jersey-hk2' - exclude group: 'org.glassfish.jersey.ext', module: 'jersey-bean-validation' - } implementation 'software.amazon.awssdk:sts' implementation 'software.amazon.awssdk:auth' implementation 'software.amazon.awssdk:kafka' implementation 'software.amazon.awssdk:kms' implementation 'software.amazon.msk:aws-msk-iam-auth:2.0.3' implementation 'software.amazon.glue:schema-registry-serde:1.1.15' - implementation 'io.confluent:kafka-json-schema-serializer:7.4.0' implementation project(':data-prepper-plugins:failures-common') - implementation 'com.github.fge:json-schema-validator:2.2.14' + implementation 'com.github.java-json-tools:json-schema-validator:2.2.14' implementation 'commons-collections:commons-collections:3.2.2' implementation 'software.amazon.awssdk:s3' implementation 'software.amazon.awssdk:apache-client' @@ -65,7 +61,6 @@ dependencies { testImplementation 'org.apache.kafka:kafka_2.13:3.6.1' testImplementation 'org.apache.kafka:kafka_2.13:3.6.1:test' testImplementation 'org.apache.curator:curator-test:5.5.0' - testImplementation 'io.confluent:kafka-schema-registry:7.4.0' testImplementation('com.kjetland:mbknor-jackson-jsonschema_2.13:1.0.39') testImplementation group: 'org.powermock', name: 'powermock-api-mockito2', version: '2.0.9' testImplementation project(':data-prepper-plugins:otel-metrics-source') @@ -74,8 +69,15 @@ dependencies { testImplementation libs.protobuf.util testImplementation libs.commons.io testImplementation libs.armeria.grpc + testImplementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml' integrationTestImplementation testLibs.junit.vintage + integrationTestImplementation 'io.confluent:kafka-schema-registry:7.4.0' + integrationTestImplementation ('io.confluent:kafka-schema-registry:7.4.0:tests') { + exclude group: 'org.glassfish.jersey.containers', module: 'jersey-container-servlet' + exclude group: 'org.glassfish.jersey.inject', module: 'jersey-hk2' + exclude group: 'org.glassfish.jersey.ext', module: 'jersey-bean-validation' + } constraints { implementation('org.mozilla:rhino') { @@ -130,4 +132,3 @@ task integrationTest(type: Test) { includeTestsMatching '*IT' } } - diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferIT.java index e250da0eeb..f959149bfa 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferIT.java @@ -127,7 +127,7 @@ void setUp() { } private KafkaBuffer createObjectUnderTest() { - return new KafkaBuffer(pluginSetting, kafkaBufferConfig, pluginFactory, acknowledgementSetManager, null, null, null); + return new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, null, null, null); } @Test @@ -150,6 +150,11 @@ void write_and_read() throws TimeoutException { // TODO: The metadata is not included. It needs to be included in the Buffer, though not in the Sink. This may be something we make configurable in the consumer/producer - whether to serialize the metadata or not. //assertThat(onlyResult.getData().getMetadata(), equalTo(record.getData().getMetadata())); assertThat(onlyResult.getData().toMap(), equalTo(record.getData().toMap())); + assertThat(objectUnderTest.getRecordsInFlight(), equalTo(0)); + assertThat(objectUnderTest.getInnerBufferRecordsInFlight(), equalTo(1)); + objectUnderTest.checkpoint(readResult.getValue()); + assertThat(objectUnderTest.getRecordsInFlight(), equalTo(0)); + assertThat(objectUnderTest.getInnerBufferRecordsInFlight(), equalTo(0)); } @Test @@ -159,7 +164,7 @@ void write_and_read_max_request_test() throws TimeoutException, NoSuchFieldExcep final Map topicConfigMap = Map.of( "name", topicName, "group_id", "buffergroup-" + RandomStringUtils.randomAlphabetic(6), - "create_topic", false + "create_topic", true ); final Map bufferConfigMap = Map.of( "topics", List.of(topicConfigMap), diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferOTelIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferOTelIT.java index 11ce71482e..c8d0dd9e88 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferOTelIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferOTelIT.java @@ -295,7 +295,7 @@ private void validateMetric(Event event) { @Test void test_otel_metrics_with_kafka_buffer() throws Exception { - KafkaBuffer kafkaBuffer = new KafkaBuffer(pluginSetting, kafkaBufferConfig, pluginFactory, acknowledgementSetManager, new OTelMetricDecoder(), null, null); + KafkaBuffer kafkaBuffer = new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, new OTelMetricDecoder(), null, null); buffer = new KafkaDelegatingBuffer(kafkaBuffer); final ExportMetricsServiceRequest request = createExportMetricsServiceRequest(); buffer.writeBytes(request.toByteArray(), null, 10_000); @@ -367,7 +367,7 @@ private void validateLog(OpenTelemetryLog logRecord) throws Exception { @Test void test_otel_logs_with_kafka_buffer() throws Exception { - KafkaBuffer kafkaBuffer = new KafkaBuffer(pluginSetting, kafkaBufferConfig, pluginFactory, acknowledgementSetManager, new OTelLogsDecoder(), null, null); + KafkaBuffer kafkaBuffer = new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, new OTelLogsDecoder(), null, null); buffer = new KafkaDelegatingBuffer(kafkaBuffer); final ExportLogsServiceRequest request = createExportLogsRequest(); buffer.writeBytes(request.toByteArray(), null, 10_000); @@ -453,7 +453,7 @@ private void validateSpan(Span span) throws Exception { @Test void test_otel_traces_with_kafka_buffer() throws Exception { - KafkaBuffer kafkaBuffer = new KafkaBuffer(pluginSetting, kafkaBufferConfig, pluginFactory, acknowledgementSetManager, new OTelTraceDecoder(), null, null); + KafkaBuffer kafkaBuffer = new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, new OTelTraceDecoder(), null, null); buffer = new KafkaDelegatingBuffer(kafkaBuffer); final ExportTraceServiceRequest request = createExportTraceRequest(); buffer.writeBytes(request.toByteArray(), null, 10_000); diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBuffer_KmsIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBuffer_KmsIT.java index 2b07936065..047e896c7f 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBuffer_KmsIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBuffer_KmsIT.java @@ -124,7 +124,7 @@ void setUp() { } private KafkaBuffer createObjectUnderTest() { - return new KafkaBuffer(pluginSetting, kafkaBufferConfig, pluginFactory, acknowledgementSetManager, null, ignored -> DefaultCredentialsProvider.create(), null); + return new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, null, ignored -> DefaultCredentialsProvider.create(), null); } @Nested diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBuffer.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBuffer.java index dd78583b58..83f6090fa9 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBuffer.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBuffer.java @@ -18,21 +18,24 @@ import org.opensearch.dataprepper.model.buffer.SizeOverflowException; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.event.Event; -import org.opensearch.dataprepper.model.plugin.PluginFactory; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.buffer.blockingbuffer.BlockingBuffer; import org.apache.kafka.common.errors.RecordTooLargeException; import org.apache.kafka.common.errors.RecordBatchTooLargeException; import org.opensearch.dataprepper.plugins.kafka.admin.KafkaAdminAccessor; import org.opensearch.dataprepper.plugins.kafka.buffer.serialization.BufferSerializationFactory; +import org.opensearch.dataprepper.plugins.kafka.common.KafkaMdc; import org.opensearch.dataprepper.plugins.kafka.common.serialization.CommonSerializationFactory; import org.opensearch.dataprepper.plugins.kafka.common.serialization.SerializationFactory; +import org.opensearch.dataprepper.plugins.kafka.common.thread.KafkaPluginThreadFactory; import org.opensearch.dataprepper.plugins.kafka.consumer.KafkaCustomConsumer; import org.opensearch.dataprepper.plugins.kafka.consumer.KafkaCustomConsumerFactory; import org.opensearch.dataprepper.plugins.kafka.producer.KafkaCustomProducer; import org.opensearch.dataprepper.plugins.kafka.producer.KafkaCustomProducerFactory; +import org.opensearch.dataprepper.plugins.kafka.service.TopicServiceFactory; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.slf4j.MDC; import java.time.Duration; import java.util.Collection; @@ -53,6 +56,7 @@ public class KafkaBuffer extends AbstractBuffer> { public static final int INNER_BUFFER_BATCH_SIZE = 250000; static final String WRITE = "Write"; static final String READ = "Read"; + static final String MDC_KAFKA_PLUGIN_VALUE = "buffer"; private final KafkaCustomProducer producer; private final KafkaAdminAccessor kafkaAdminAccessor; private final AbstractBuffer> innerBuffer; @@ -62,17 +66,17 @@ public class KafkaBuffer extends AbstractBuffer> { private ByteDecoder byteDecoder; @DataPrepperPluginConstructor - public KafkaBuffer(final PluginSetting pluginSetting, final KafkaBufferConfig kafkaBufferConfig, final PluginFactory pluginFactory, + public KafkaBuffer(final PluginSetting pluginSetting, final KafkaBufferConfig kafkaBufferConfig, final AcknowledgementSetManager acknowledgementSetManager, final ByteDecoder byteDecoder, final AwsCredentialsSupplier awsCredentialsSupplier, final CircuitBreaker circuitBreaker) { super(kafkaBufferConfig.getCustomMetricPrefix().orElse(pluginSetting.getName()), pluginSetting.getPipelineName()); final SerializationFactory serializationFactory = new BufferSerializationFactory(new CommonSerializationFactory()); - final KafkaCustomProducerFactory kafkaCustomProducerFactory = new KafkaCustomProducerFactory(serializationFactory, awsCredentialsSupplier); + final KafkaCustomProducerFactory kafkaCustomProducerFactory = new KafkaCustomProducerFactory(serializationFactory, awsCredentialsSupplier, new TopicServiceFactory()); this.byteDecoder = byteDecoder; final String metricPrefixName = kafkaBufferConfig.getCustomMetricPrefix().orElse(pluginSetting.getName()); final PluginMetrics producerMetrics = PluginMetrics.fromNames(metricPrefixName + WRITE, pluginSetting.getPipelineName()); - producer = kafkaCustomProducerFactory.createProducer(kafkaBufferConfig, pluginFactory, pluginSetting, null, null, producerMetrics, null, false); + producer = kafkaCustomProducerFactory.createProducer(kafkaBufferConfig, null, null, producerMetrics, null, false); final KafkaCustomConsumerFactory kafkaCustomConsumerFactory = new KafkaCustomConsumerFactory(serializationFactory, awsCredentialsSupplier); innerBuffer = new BlockingBuffer<>(INNER_BUFFER_CAPACITY, INNER_BUFFER_BATCH_SIZE, pluginSetting.getPipelineName()); this.shutdownInProgress = new AtomicBoolean(false); @@ -80,7 +84,7 @@ public KafkaBuffer(final PluginSetting pluginSetting, final KafkaBufferConfig ka final List consumers = kafkaCustomConsumerFactory.createConsumersForTopic(kafkaBufferConfig, kafkaBufferConfig.getTopic(), innerBuffer, consumerMetrics, acknowledgementSetManager, byteDecoder, shutdownInProgress, false, circuitBreaker); this.kafkaAdminAccessor = new KafkaAdminAccessor(kafkaBufferConfig, List.of(kafkaBufferConfig.getTopic().getGroupId())); - this.executorService = Executors.newFixedThreadPool(consumers.size()); + this.executorService = Executors.newFixedThreadPool(consumers.size(), KafkaPluginThreadFactory.defaultExecutorThreadFactory(MDC_KAFKA_PLUGIN_VALUE)); consumers.forEach(this.executorService::submit); this.drainTimeout = kafkaBufferConfig.getDrainTimeout(); @@ -89,6 +93,7 @@ public KafkaBuffer(final PluginSetting pluginSetting, final KafkaBufferConfig ka @Override public void writeBytes(final byte[] bytes, final String key, int timeoutInMillis) throws Exception { try { + setMdc(); producer.produceRawData(bytes, key); } catch (final Exception e) { LOG.error(e.getMessage(), e); @@ -102,15 +107,21 @@ public void writeBytes(final byte[] bytes, final String key, int timeoutInMillis throw new RuntimeException(e); } } + finally { + resetMdc(); + } } @Override public void doWrite(Record record, int timeoutInMillis) throws TimeoutException { try { + setMdc(); producer.produceRecords(record); } catch (final Exception e) { LOG.error(e.getMessage(), e); throw new RuntimeException(e); + } finally { + resetMdc(); } } @@ -121,29 +132,50 @@ public boolean isByteBuffer() { @Override public void doWriteAll(Collection> records, int timeoutInMillis) throws Exception { - for ( Record record: records ) { + for (Record record : records) { doWrite(record, timeoutInMillis); } } @Override public Map.Entry>, CheckpointState> doRead(int timeoutInMillis) { - return innerBuffer.read(timeoutInMillis); + try { + setMdc(); + return innerBuffer.read(timeoutInMillis); + } finally { + resetMdc(); + } } @Override public void postProcess(final Long recordsInBuffer) { - innerBuffer.postProcess(recordsInBuffer); + try { + setMdc(); + + innerBuffer.postProcess(recordsInBuffer); + } finally { + resetMdc(); + } } @Override public void doCheckpoint(CheckpointState checkpointState) { - innerBuffer.doCheckpoint(checkpointState); + try { + setMdc(); + innerBuffer.checkpoint(checkpointState); + } finally { + resetMdc(); + } } @Override public boolean isEmpty() { - return kafkaAdminAccessor.areTopicsEmpty() && innerBuffer.isEmpty(); + try { + setMdc(); + return kafkaAdminAccessor.areTopicsEmpty() && innerBuffer.isEmpty(); + } finally { + resetMdc(); + } } @Override @@ -156,23 +188,41 @@ public boolean isWrittenOffHeapOnly() { return true; } + int getInnerBufferRecordsInFlight() { + return innerBuffer.getRecordsInFlight(); + } + @Override public void shutdown() { - shutdownInProgress.set(true); - executorService.shutdown(); - try { - if (executorService.awaitTermination(EXECUTOR_SERVICE_SHUTDOWN_TIMEOUT, TimeUnit.SECONDS)) { - LOG.info("Successfully waited for consumer task to terminate"); - } else { - LOG.warn("Consumer task did not terminate in time, forcing termination"); + setMdc(); + + shutdownInProgress.set(true); + executorService.shutdown(); + + try { + if (executorService.awaitTermination(EXECUTOR_SERVICE_SHUTDOWN_TIMEOUT, TimeUnit.SECONDS)) { + LOG.info("Successfully waited for consumer task to terminate"); + } else { + LOG.warn("Consumer task did not terminate in time, forcing termination"); + executorService.shutdownNow(); + } + } catch (final InterruptedException e) { + LOG.error("Interrupted while waiting for consumer task to terminate", e); executorService.shutdownNow(); } - } catch (final InterruptedException e) { - LOG.error("Interrupted while waiting for consumer task to terminate", e); - executorService.shutdownNow(); + + innerBuffer.shutdown(); + } finally { + resetMdc(); } + } + + private static void setMdc() { + MDC.put(KafkaMdc.MDC_KAFKA_PLUGIN_KEY, MDC_KAFKA_PLUGIN_VALUE); + } - innerBuffer.shutdown(); + private static void resetMdc() { + MDC.remove(KafkaMdc.MDC_KAFKA_PLUGIN_KEY); } } diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/common/KafkaMdc.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/common/KafkaMdc.java new file mode 100644 index 0000000000..9ae8985908 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/common/KafkaMdc.java @@ -0,0 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.common;public class KafkaMdc { + public static final String MDC_KAFKA_PLUGIN_KEY = "kafkaPluginType"; +} diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/common/thread/KafkaPluginThreadFactory.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/common/thread/KafkaPluginThreadFactory.java new file mode 100644 index 0000000000..a05540c320 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/common/thread/KafkaPluginThreadFactory.java @@ -0,0 +1,57 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.common.thread; + +import org.opensearch.dataprepper.plugins.kafka.common.KafkaMdc; +import org.slf4j.MDC; + +import java.util.concurrent.Executors; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * An implementation of {@link ThreadFactory} for Kafka plugin threads. + */ +public class KafkaPluginThreadFactory implements ThreadFactory { + private final ThreadFactory delegateThreadFactory; + private final String threadPrefix; + private final String kafkaPluginType; + private final AtomicInteger threadNumber = new AtomicInteger(1); + + KafkaPluginThreadFactory( + final ThreadFactory delegateThreadFactory, + final String kafkaPluginType) { + this.delegateThreadFactory = delegateThreadFactory; + this.threadPrefix = "kafka-" + kafkaPluginType + "-"; + this.kafkaPluginType = kafkaPluginType; + } + + /** + * Creates an instance specifically for use with {@link Executors}. + * + * @param kafkaPluginType The name of the plugin type. e.g. sink, source, buffer + * @return An instance of the {@link KafkaPluginThreadFactory}. + */ + public static KafkaPluginThreadFactory defaultExecutorThreadFactory(final String kafkaPluginType) { + return new KafkaPluginThreadFactory(Executors.defaultThreadFactory(), kafkaPluginType); + } + + @Override + public Thread newThread(final Runnable runnable) { + final Thread thread = delegateThreadFactory.newThread(() -> { + MDC.put(KafkaMdc.MDC_KAFKA_PLUGIN_KEY, kafkaPluginType); + try { + runnable.run(); + } finally { + MDC.clear(); + } + }); + + thread.setName(threadPrefix + threadNumber.getAndIncrement()); + + return thread; + } +} diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/EncryptionConfig.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/EncryptionConfig.java index 826c2edadf..4b1e975c74 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/EncryptionConfig.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/EncryptionConfig.java @@ -11,6 +11,15 @@ public class EncryptionConfig { @JsonProperty("type") private EncryptionType type = EncryptionType.SSL; + @JsonProperty("certificate_content") + private String certificateContent; + + @JsonProperty("trust_store_file_path") + private String trustStoreFilePath; + + @JsonProperty("trust_store_password") + private String trustStorePassword; + @JsonProperty("insecure") private boolean insecure = false; @@ -18,6 +27,18 @@ public EncryptionType getType() { return type; } + public String getCertificateContent() { + return certificateContent; + } + + public String getTrustStoreFilePath() { + return trustStoreFilePath; + } + + public String getTrustStorePassword() { + return trustStorePassword; + } + public boolean getInsecure() { return insecure; } diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java index 099daf754f..581bf41a4e 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java @@ -67,6 +67,7 @@ public class KafkaCustomConsumer implements Runnable, ConsumerRebalanceListener private static final Logger LOG = LoggerFactory.getLogger(KafkaCustomConsumer.class); private static final Long COMMIT_OFFSET_INTERVAL_MS = 300000L; private static final int DEFAULT_NUMBER_OF_RECORDS_TO_ACCUMULATE = 1; + private static final int RETRY_ON_EXCEPTION_SLEEP_MS = 1000; static final String DEFAULT_KEY = "message"; private volatile long lastCommitTime; @@ -75,6 +76,7 @@ public class KafkaCustomConsumer implements Runnable, ConsumerRebalanceListener private final String topicName; private final TopicConsumerConfig topicConfig; private MessageFormat schema; + private boolean paused; private final BufferAccumulator> bufferAccumulator; private final Buffer> buffer; private static final ObjectMapper objectMapper = new ObjectMapper(); @@ -94,6 +96,7 @@ public class KafkaCustomConsumer implements Runnable, ConsumerRebalanceListener private long numRecordsCommitted = 0; private final LogRateLimiter errLogRateLimiter; private final ByteDecoder byteDecoder; + private final long maxRetriesOnException; public KafkaCustomConsumer(final KafkaConsumer consumer, final AtomicBoolean shutdownInProgress, @@ -110,8 +113,10 @@ public KafkaCustomConsumer(final KafkaConsumer consumer, this.shutdownInProgress = shutdownInProgress; this.consumer = consumer; this.buffer = buffer; + this.paused = false; this.byteDecoder = byteDecoder; this.topicMetrics = topicMetrics; + this.maxRetriesOnException = topicConfig.getMaxPollInterval().toMillis() / (2 * RETRY_ON_EXCEPTION_SLEEP_MS); this.pauseConsumePredicate = pauseConsumePredicate; this.topicMetrics.register(consumer); this.offsetsToCommit = new HashMap<>(); @@ -170,10 +175,15 @@ private AcknowledgementSet createAcknowledgementSet(Map void consumeRecords() throws Exception { - try { + ConsumerRecords doPoll() throws Exception { ConsumerRecords records = consumer.poll(Duration.ofMillis(topicConfig.getThreadWaitingTime().toMillis()/2)); + return records; + } + + void consumeRecords() throws Exception { + try { + ConsumerRecords records = doPoll(); if (Objects.nonNull(records) && !records.isEmpty() && records.count() > 0) { Map offsets = new HashMap<>(); AcknowledgementSet acknowledgementSet = null; @@ -338,9 +348,19 @@ public void run() { boolean retryingAfterException = false; while (!shutdownInProgress.get()) { try { - if (retryingAfterException || pauseConsumePredicate.pauseConsuming()) { - LOG.debug("Pause consuming from Kafka topic."); + if (retryingAfterException) { + LOG.debug("Pause consuming from Kafka topic due a previous exception."); + Thread.sleep(10000); + } else if (pauseConsumePredicate.pauseConsuming()) { + LOG.debug("Pause and skip consuming from Kafka topic due to an external condition: {}", pauseConsumePredicate); + paused = true; + consumer.pause(consumer.assignment()); Thread.sleep(10000); + continue; + } else if(paused) { + LOG.debug("Resume consuming from Kafka topic."); + paused = false; + consumer.resume(consumer.assignment()); } synchronized(this) { commitOffsets(false); @@ -361,6 +381,7 @@ public void run() { } private Record getRecord(ConsumerRecord consumerRecord, int partition) { + Instant now = Instant.now(); Map data = new HashMap<>(); Event event; Object value = consumerRecord.value(); @@ -419,21 +440,45 @@ private void processRecord(final AcknowledgementSet acknowledgementSet, final Re if (acknowledgementSet != null) { acknowledgementSet.add(record.getData()); } + long numRetries = 0; while (true) { try { - bufferAccumulator.add(record); + if (numRetries == 0) { + bufferAccumulator.add(record); + } else { + bufferAccumulator.flush(); + } break; } catch (Exception e) { + if (!paused && numRetries++ > maxRetriesOnException) { + paused = true; + consumer.pause(consumer.assignment()); + } if (e instanceof SizeOverflowException) { topicMetrics.getNumberOfBufferSizeOverflows().increment(); } else { LOG.debug("Error while adding record to buffer, retrying ", e); } try { - Thread.sleep(100); + Thread.sleep(RETRY_ON_EXCEPTION_SLEEP_MS); + if (paused) { + ConsumerRecords records = doPoll(); + if (records.count() > 0) { + LOG.warn("Unexpected records received while the consumer is paused. Resetting the partitions to retry from last read pointer"); + synchronized(this) { + partitionsToReset.addAll(consumer.assignment()); + }; + break; + } + } } catch (Exception ex) {} // ignore the exception because it only means the thread slept for shorter time } } + + if (paused) { + consumer.resume(consumer.assignment()); + paused = false; + } } private void iterateRecordPartitions(ConsumerRecords records, final AcknowledgementSet acknowledgementSet, @@ -452,7 +497,7 @@ private void iterateRecordPartitions(ConsumerRecords records, fin if (schema == MessageFormat.BYTES) { InputStream inputStream = new ByteArrayInputStream((byte[])consumerRecord.value()); if(byteDecoder != null) { - byteDecoder.parse(inputStream, (record) -> { + byteDecoder.parse(inputStream, Instant.ofEpochMilli(consumerRecord.timestamp()), (record) -> { processRecord(acknowledgementSet, record); }); } else { @@ -503,6 +548,9 @@ public void onPartitionsAssigned(Collection partitions) { LOG.info("Assigned partition {}", topicPartition); ownedPartitionsEpoch.put(topicPartition, epoch); } + if (paused) { + consumer.pause(consumer.assignment()); + } } dumpTopicPartitionOffsets(partitions); } @@ -520,6 +568,9 @@ public void onPartitionsRevoked(Collection partitions) { ownedPartitionsEpoch.remove(topicPartition); partitionCommitTrackerMap.remove(topicPartition.partition()); } + if (paused) { + consumer.pause(consumer.assignment()); + } } } diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerFactory.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerFactory.java index fab0a4f56e..961e0328d3 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerFactory.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerFactory.java @@ -10,7 +10,6 @@ import io.confluent.kafka.serializers.KafkaAvroDeserializer; import io.confluent.kafka.serializers.KafkaAvroDeserializerConfig; import io.confluent.kafka.serializers.KafkaJsonDeserializer; -import kafka.common.BrokerEndPointNotAvailableException; import org.apache.commons.lang3.StringUtils; import org.apache.kafka.clients.consumer.ConsumerConfig; import org.apache.kafka.clients.consumer.KafkaConsumer; @@ -108,8 +107,7 @@ public List createConsumersForTopic(final KafkaConsumerConf }); } catch (Exception e) { - if (e instanceof BrokerNotAvailableException || - e instanceof BrokerEndPointNotAvailableException || e instanceof TimeoutException) { + if (e instanceof BrokerNotAvailableException || e instanceof TimeoutException) { LOG.error("The Kafka broker is not available."); } else { LOG.error("Failed to setup the Kafka Source Plugin.", e); diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/PauseConsumePredicate.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/PauseConsumePredicate.java index 0d5260d92a..c1537de0b8 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/PauseConsumePredicate.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/PauseConsumePredicate.java @@ -30,7 +30,17 @@ public interface PauseConsumePredicate { static PauseConsumePredicate circuitBreakingPredicate(final CircuitBreaker circuitBreaker) { if(circuitBreaker == null) return noPause(); - return circuitBreaker::isOpen; + return new PauseConsumePredicate() { + @Override + public boolean pauseConsuming() { + return circuitBreaker.isOpen(); + } + + @Override + public String toString() { + return "Circuit Breaker"; + } + }; } /** diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducer.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducer.java index fe0291662b..56e9783d2b 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducer.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducer.java @@ -169,11 +169,11 @@ private void publishAvroMessage(final Record record, final String key) th } Future send(final String topicName, String key, final Object record) throws Exception { - if (Objects.isNull(key)) { - return producer.send(new ProducerRecord(topicName, record), callBack(record)); - } + ProducerRecord producerRecord = Objects.isNull(key) ? + new ProducerRecord(topicName, record) : + new ProducerRecord(topicName, key, record); - return producer.send(new ProducerRecord(topicName, key, record), callBack(record)); + return producer.send(producerRecord, callBack(record)); } private void publishJsonMessage(final Record record, final String key) throws IOException, ProcessingException, Exception { diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerFactory.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerFactory.java index ec65074a4a..72c0694b41 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerFactory.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerFactory.java @@ -11,8 +11,6 @@ import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; import org.opensearch.dataprepper.expression.ExpressionEvaluator; import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.configuration.PluginSetting; -import org.opensearch.dataprepper.model.plugin.PluginFactory; import org.opensearch.dataprepper.model.sink.SinkContext; import org.opensearch.dataprepper.plugins.kafka.common.KafkaDataConfig; import org.opensearch.dataprepper.plugins.kafka.common.KafkaDataConfigAdapter; @@ -22,13 +20,14 @@ import org.opensearch.dataprepper.plugins.kafka.common.serialization.SerializationFactory; import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerProperties; - -import org.opensearch.dataprepper.plugins.kafka.configuration.TopicProducerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicProducerConfig; import org.opensearch.dataprepper.plugins.kafka.consumer.KafkaCustomConsumerFactory; import org.opensearch.dataprepper.plugins.kafka.service.SchemaService; import org.opensearch.dataprepper.plugins.kafka.service.TopicService; +import org.opensearch.dataprepper.plugins.kafka.service.TopicServiceFactory; import org.opensearch.dataprepper.plugins.kafka.sink.DLQSink; +import org.opensearch.dataprepper.plugins.kafka.util.KafkaProducerMetrics; import org.opensearch.dataprepper.plugins.kafka.util.KafkaSecurityConfigurer; import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicProducerMetrics; import org.opensearch.dataprepper.plugins.kafka.util.RestUtils; @@ -43,13 +42,18 @@ public class KafkaCustomProducerFactory { private static final Logger LOG = LoggerFactory.getLogger(KafkaCustomConsumerFactory.class); private final SerializationFactory serializationFactory; private final AwsCredentialsSupplier awsCredentialsSupplier; + private final TopicServiceFactory topicServiceFactory; - public KafkaCustomProducerFactory(final SerializationFactory serializationFactory, AwsCredentialsSupplier awsCredentialsSupplier) { + public KafkaCustomProducerFactory( + final SerializationFactory serializationFactory, + final AwsCredentialsSupplier awsCredentialsSupplier, + final TopicServiceFactory topicServiceFactory) { this.serializationFactory = serializationFactory; this.awsCredentialsSupplier = awsCredentialsSupplier; + this.topicServiceFactory = topicServiceFactory; } - public KafkaCustomProducer createProducer(final KafkaProducerConfig kafkaProducerConfig, final PluginFactory pluginFactory, final PluginSetting pluginSetting, + public KafkaCustomProducer createProducer(final KafkaProducerConfig kafkaProducerConfig, final ExpressionEvaluator expressionEvaluator, final SinkContext sinkContext, final PluginMetrics pluginMetrics, final DLQSink dlqSink, final boolean topicNameInMetrics) { @@ -77,12 +81,14 @@ public KafkaCustomProducer createProducer(final KafkaProducerConfig kafkaProduce Serializer valueSerializer = (Serializer) serializationFactory.getSerializer(dataConfig); final KafkaProducer producer = new KafkaProducer<>(properties, keyDeserializer, valueSerializer); final KafkaTopicProducerMetrics topicMetrics = new KafkaTopicProducerMetrics(topic.getName(), pluginMetrics, topicNameInMetrics); + KafkaProducerMetrics.registerProducer(pluginMetrics, producer); final String topicName = ObjectUtils.isEmpty(kafkaProducerConfig.getTopic()) ? null : kafkaProducerConfig.getTopic().getName(); final SchemaService schemaService = new SchemaService.SchemaServiceBuilder().getFetchSchemaService(topicName, kafkaProducerConfig.getSchemaConfig()).build(); return new KafkaCustomProducer(producer, kafkaProducerConfig, dlqSink, expressionEvaluator, Objects.nonNull(sinkContext) ? sinkContext.getTagsTargetKey() : null, topicMetrics, schemaService); } + private void prepareTopicAndSchema(final KafkaProducerConfig kafkaProducerConfig, final Integer maxRequestSize) { checkTopicCreationCriteriaAndCreateTopic(kafkaProducerConfig, maxRequestSize); final SchemaConfig schemaConfig = kafkaProducerConfig.getSchemaConfig(); @@ -102,8 +108,8 @@ private void prepareTopicAndSchema(final KafkaProducerConfig kafkaProducerConfig private void checkTopicCreationCriteriaAndCreateTopic(final KafkaProducerConfig kafkaProducerConfig, final Integer maxRequestSize) { final TopicProducerConfig topic = kafkaProducerConfig.getTopic(); - if (!topic.isCreateTopic()) { - final TopicService topicService = new TopicService(kafkaProducerConfig); + if (topic.isCreateTopic()) { + final TopicService topicService = topicServiceFactory.createTopicService(kafkaProducerConfig); Long maxMessageBytes = null; if (maxRequestSize != null) { maxMessageBytes = Long.valueOf(maxRequestSize); @@ -111,7 +117,5 @@ private void checkTopicCreationCriteriaAndCreateTopic(final KafkaProducerConfig topicService.createTopic(kafkaProducerConfig.getTopic().getName(), topic.getNumberOfPartitions(), topic.getReplicationFactor(), maxMessageBytes); topicService.closeAdminClient(); } - - } } diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/service/TopicService.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/service/TopicService.java index 16b907f9bb..19edf8233d 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/service/TopicService.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/service/TopicService.java @@ -21,7 +21,7 @@ public class TopicService { private final AdminClient adminClient; - public TopicService(final KafkaProducerConfig kafkaProducerConfig) { + TopicService(final KafkaProducerConfig kafkaProducerConfig) { this.adminClient = AdminClient.create(SinkPropertyConfigurer.getPropertiesForAdminClient(kafkaProducerConfig)); } diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/service/TopicServiceFactory.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/service/TopicServiceFactory.java new file mode 100644 index 0000000000..c159334184 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/service/TopicServiceFactory.java @@ -0,0 +1,14 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.service; + +import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerConfig; + +public class TopicServiceFactory { + public TopicService createTopicService(final KafkaProducerConfig kafkaProducerConfig) { + return new TopicService(kafkaProducerConfig); + } +} diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSink.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSink.java index 1adb77643b..97c93d22de 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSink.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSink.java @@ -27,6 +27,7 @@ import org.opensearch.dataprepper.plugins.kafka.producer.ProducerWorker; import org.opensearch.dataprepper.plugins.kafka.service.SchemaService; import org.opensearch.dataprepper.plugins.kafka.service.TopicService; +import org.opensearch.dataprepper.plugins.kafka.service.TopicServiceFactory; import org.opensearch.dataprepper.plugins.kafka.util.RestUtils; //import org.opensearch.dataprepper.plugins.kafka.sink.DLQSink; import org.slf4j.Logger; @@ -50,6 +51,7 @@ public class KafkaSink extends AbstractSink> { private final KafkaSinkConfig kafkaSinkConfig; private final KafkaCustomProducerFactory kafkaCustomProducerFactory; + private final TopicServiceFactory topicServiceFactory; private volatile boolean sinkInitialized; @@ -86,7 +88,8 @@ public KafkaSink(final PluginSetting pluginSetting, final KafkaSinkConfig kafkaS this.sinkContext = sinkContext; SerializationFactory serializationFactory = new CommonSerializationFactory(); - kafkaCustomProducerFactory = new KafkaCustomProducerFactory(serializationFactory, awsCredentialsSupplier); + topicServiceFactory = new TopicServiceFactory(); + kafkaCustomProducerFactory = new KafkaCustomProducerFactory(serializationFactory, awsCredentialsSupplier, topicServiceFactory); } @@ -154,7 +157,7 @@ private void prepareTopicAndSchema() { private void checkTopicCreationCriteriaAndCreateTopic() { final TopicProducerConfig topic = kafkaSinkConfig.getTopic(); if (topic.isCreateTopic()) { - final TopicService topicService = new TopicService(kafkaSinkConfig); + final TopicService topicService = topicServiceFactory.createTopicService(kafkaSinkConfig); topicService.createTopic(kafkaSinkConfig.getTopic().getName(), topic.getNumberOfPartitions(), topic.getReplicationFactor(), null); topicService.closeAdminClient(); } @@ -164,7 +167,7 @@ private void checkTopicCreationCriteriaAndCreateTopic() { private KafkaCustomProducer createProducer() { final DLQSink dlqSink = new DLQSink(pluginFactory, kafkaSinkConfig, pluginSetting); - return kafkaCustomProducerFactory.createProducer(kafkaSinkConfig, pluginFactory, pluginSetting, expressionEvaluator, sinkContext, pluginMetrics, dlqSink, true); + return kafkaCustomProducerFactory.createProducer(kafkaSinkConfig, expressionEvaluator, sinkContext, pluginMetrics, dlqSink, true); } diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSource.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSource.java index ec27f1f370..fbdee41105 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSource.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSource.java @@ -11,7 +11,6 @@ import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; import io.confluent.kafka.serializers.KafkaAvroDeserializer; import io.confluent.kafka.serializers.json.KafkaJsonSchemaDeserializer; -import kafka.common.BrokerEndPointNotAvailableException; import org.apache.avro.generic.GenericRecord; import org.apache.commons.lang3.StringUtils; import org.apache.kafka.clients.consumer.ConsumerConfig; @@ -146,8 +145,7 @@ public void start(Buffer> buffer) { executorService.submit(consumer); }); } catch (Exception e) { - if (e instanceof BrokerNotAvailableException || - e instanceof BrokerEndPointNotAvailableException || e instanceof TimeoutException) { + if (e instanceof BrokerNotAvailableException || e instanceof TimeoutException) { LOG.error("The kafka broker is not available..."); } else { LOG.error("Failed to setup the Kafka Source Plugin.", e); diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/CustomClientSslEngineFactory.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/CustomClientSslEngineFactory.java new file mode 100644 index 0000000000..83c52d736b --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/CustomClientSslEngineFactory.java @@ -0,0 +1,82 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.util; + +import org.apache.kafka.common.security.auth.SslEngineFactory; +import org.opensearch.dataprepper.plugins.truststore.TrustStoreProvider; + +import javax.net.ssl.SSLContext; +import javax.net.ssl.SSLEngine; +import javax.net.ssl.TrustManager; +import java.security.KeyManagementException; +import java.security.KeyStore; +import java.security.NoSuchAlgorithmException; +import java.security.SecureRandom; +import java.util.Map; +import java.util.Objects; +import java.util.Set; + +public class CustomClientSslEngineFactory implements SslEngineFactory { + String certificateContent = null; + + @Override + public void configure(Map configs) { + certificateContent = configs.get("certificateContent").toString(); + } + + private TrustManager[] getTrustManager() { + final TrustManager[] trustManagers; + if (Objects.nonNull(certificateContent)) { + trustManagers = TrustStoreProvider.createTrustManager(certificateContent); + } else { + trustManagers = TrustStoreProvider.createTrustAllManager(); + } + return trustManagers; + } + + @Override + public SSLEngine createClientSslEngine(final String peerHost, final int peerPort, final String endpointIdentification) { + try { + final SSLContext sslContext = SSLContext.getInstance("SSL"); + sslContext.init(null, getTrustManager(), new SecureRandom()); + SSLEngine sslEngine = sslContext.createSSLEngine(peerHost, peerPort); + sslEngine.setUseClientMode(true); + return sslEngine; + } catch (NoSuchAlgorithmException | KeyManagementException e) { + throw new RuntimeException(e); + } + } + + @Override + public SSLEngine createServerSslEngine(String peerHost, int peerPort) { + return null; + } + + @Override + public boolean shouldBeRebuilt(Map nextConfigs) { + return false; + } + + @Override + public Set reconfigurableConfigs() { + return null; + } + + @Override + public KeyStore keystore() { + return null; + } + + @Override + public KeyStore truststore() { + return null; + } + + @Override + public void close() { + + } +} diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaProducerMetrics.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaProducerMetrics.java new file mode 100644 index 0000000000..478b743623 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaProducerMetrics.java @@ -0,0 +1,43 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.util; + +import org.apache.kafka.clients.producer.KafkaProducer; +import org.apache.kafka.common.Metric; +import org.apache.kafka.common.MetricName; +import org.opensearch.dataprepper.metrics.PluginMetrics; + +import java.util.Map; + +/** + * Metrics for a Kafka producer. These span all topics, which makes it distinct from + * the {@link KafkaTopicProducerMetrics}. + */ +public final class KafkaProducerMetrics { + static final Map METRICS_NAME_MAP = Map.of( + "record-queue-time-avg", "recordQueueTimeAvg", + "record-queue-time-max", "recordQueueTimeMax", + "buffer-exhausted-rate", "bufferExhaustedRate", + "buffer-available-bytes", "bufferAvailableBytes" + ); + + private KafkaProducerMetrics() { } + + public static void registerProducer(final PluginMetrics pluginMetrics, final KafkaProducer kafkaProducer) { + final Map kafkaProducerMetrics = kafkaProducer.metrics(); + for (final Map.Entry metricNameEntry : kafkaProducerMetrics.entrySet()) { + final MetricName kafkaMetricName = metricNameEntry.getKey(); + + final String dataPrepperMetricName = METRICS_NAME_MAP.get(kafkaMetricName.name()); + if(dataPrepperMetricName == null) + continue; + + final Metric metric = metricNameEntry.getValue(); + + pluginMetrics.gauge(dataPrepperMetricName, metric, m -> (double) m.metricValue()); + } + } +} diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaSecurityConfigurer.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaSecurityConfigurer.java index 779aefcab0..47daed1e24 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaSecurityConfigurer.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaSecurityConfigurer.java @@ -83,6 +83,10 @@ public class KafkaSecurityConfigurer { private static final String REGISTRY_BASIC_AUTH_USER_INFO = "schema.registry.basic.auth.user.info"; private static final int MAX_KAFKA_CLIENT_RETRIES = 360; // for one hour every 10 seconds + private static final String SSL_ENGINE_FACTORY_CLASS = "ssl.engine.factory.class"; + private static final String CERTIFICATE_CONTENT = "certificateContent"; + private static final String SSL_TRUSTSTORE_LOCATION = "ssl.truststore.location"; + private static final String SSL_TRUSTSTORE_PASSWORD = "ssl.truststore.password"; private static AwsCredentialsProvider credentialsProvider; private static GlueSchemaRegistryKafkaDeserializer glueDeserializer; @@ -108,18 +112,39 @@ public class KafkaSecurityConfigurer { properties.put(SASL_JAAS_CONFIG, String.format(PLAINTEXT_JAASCONFIG, username, password)); }*/ - private static void setPlainTextAuthProperties(Properties properties, final PlainTextAuthConfig plainTextAuthConfig, EncryptionType encryptionType) { - String username = plainTextAuthConfig.getUsername(); - String password = plainTextAuthConfig.getPassword(); + private static void setPlainTextAuthProperties(final Properties properties, final PlainTextAuthConfig plainTextAuthConfig, + final EncryptionConfig encryptionConfig) { + final String username = plainTextAuthConfig.getUsername(); + final String password = plainTextAuthConfig.getPassword(); properties.put(SASL_MECHANISM, "PLAIN"); properties.put(SASL_JAAS_CONFIG, "org.apache.kafka.common.security.plain.PlainLoginModule required username=\"" + username + "\" password=\"" + password + "\";"); - if (encryptionType == EncryptionType.NONE) { - properties.put(SECURITY_PROTOCOL, "SASL_PLAINTEXT"); - } else { // EncryptionType.SSL + if (checkEncryptionType(encryptionConfig, EncryptionType.SSL)) { properties.put(SECURITY_PROTOCOL, "SASL_SSL"); + setSecurityProtocolSSLProperties(properties, encryptionConfig); + } else { // EncryptionType.NONE + properties.put(SECURITY_PROTOCOL, "SASL_PLAINTEXT"); } } + private static void setSecurityProtocolSSLProperties(final Properties properties, final EncryptionConfig encryptionConfig) { + if (Objects.nonNull(encryptionConfig.getCertificateContent())) { + setCustomSslProperties(properties, encryptionConfig.getCertificateContent()); + } else if (Objects.nonNull(encryptionConfig.getTrustStoreFilePath()) && + Objects.nonNull(encryptionConfig.getTrustStorePassword())) { + setTruststoreProperties(properties, encryptionConfig); + } + } + + private static void setCustomSslProperties(final Properties properties, final String certificateContent) { + properties.put(CERTIFICATE_CONTENT, certificateContent); + properties.put(SSL_ENGINE_FACTORY_CLASS, CustomClientSslEngineFactory.class); + } + + private static void setTruststoreProperties(final Properties properties, final EncryptionConfig encryptionConfig) { + properties.put(SSL_TRUSTSTORE_LOCATION, encryptionConfig.getTrustStoreFilePath()); + properties.put(SSL_TRUSTSTORE_PASSWORD, encryptionConfig.getTrustStorePassword()); + } + public static void setOauthProperties(final KafkaClusterAuthConfig kafkaClusterAuthConfig, final Properties properties) { final OAuthConfig oAuthConfig = kafkaClusterAuthConfig.getAuthConfig().getSaslAuthConfig().getOAuthConfig(); @@ -258,27 +283,25 @@ public static String getBootStrapServersForMsk(final AwsIamAuthConfig awsIamAuth } } - public static void setAuthProperties(Properties properties, final KafkaClusterAuthConfig kafkaClusterAuthConfig, final Logger LOG) { + public static void setAuthProperties(final Properties properties, final KafkaClusterAuthConfig kafkaClusterAuthConfig, final Logger LOG) { final AwsConfig awsConfig = kafkaClusterAuthConfig.getAwsConfig(); final AuthConfig authConfig = kafkaClusterAuthConfig.getAuthConfig(); final EncryptionConfig encryptionConfig = kafkaClusterAuthConfig.getEncryptionConfig(); - final EncryptionType encryptionType = encryptionConfig.getType(); - credentialsProvider = DefaultCredentialsProvider.create(); String bootstrapServers = ""; if (Objects.nonNull(kafkaClusterAuthConfig.getBootstrapServers())) { bootstrapServers = String.join(",", kafkaClusterAuthConfig.getBootstrapServers()); } - AwsIamAuthConfig awsIamAuthConfig = null; + if (Objects.nonNull(authConfig)) { - AuthConfig.SaslAuthConfig saslAuthConfig = authConfig.getSaslAuthConfig(); + final AuthConfig.SaslAuthConfig saslAuthConfig = authConfig.getSaslAuthConfig(); if (Objects.nonNull(saslAuthConfig)) { - awsIamAuthConfig = saslAuthConfig.getAwsIamAuthConfig(); - PlainTextAuthConfig plainTextAuthConfig = saslAuthConfig.getPlainTextAuthConfig(); + final AwsIamAuthConfig awsIamAuthConfig = saslAuthConfig.getAwsIamAuthConfig(); + final PlainTextAuthConfig plainTextAuthConfig = saslAuthConfig.getPlainTextAuthConfig(); if (Objects.nonNull(awsIamAuthConfig)) { - if (encryptionType == EncryptionType.NONE) { + if (checkEncryptionType(encryptionConfig, EncryptionType.NONE)) { throw new RuntimeException("Encryption Config must be SSL to use IAM authentication mechanism"); } if (Objects.isNull(awsConfig)) { @@ -288,33 +311,39 @@ public static void setAuthProperties(Properties properties, final KafkaClusterAu bootstrapServers = getBootStrapServersForMsk(awsIamAuthConfig, awsConfig, LOG); } else if (Objects.nonNull(saslAuthConfig.getOAuthConfig())) { setOauthProperties(kafkaClusterAuthConfig, properties); - } else if (Objects.nonNull(plainTextAuthConfig)) { - setPlainTextAuthProperties(properties, plainTextAuthConfig, encryptionType); + } else if (Objects.nonNull(plainTextAuthConfig) && Objects.nonNull(kafkaClusterAuthConfig.getEncryptionConfig())) { + setPlainTextAuthProperties(properties, plainTextAuthConfig, kafkaClusterAuthConfig.getEncryptionConfig()); } else { throw new RuntimeException("No SASL auth config specified"); } } if (encryptionConfig.getInsecure()) { - properties.put("ssl.engine.factory.class", InsecureSslEngineFactory.class); + properties.put(SSL_ENGINE_FACTORY_CLASS, InsecureSslEngineFactory.class); } } if (Objects.isNull(authConfig) || Objects.isNull(authConfig.getSaslAuthConfig())) { - if (encryptionType == EncryptionType.SSL) { + if (checkEncryptionType(encryptionConfig, EncryptionType.SSL)) { properties.put(SECURITY_PROTOCOL, "SSL"); + setSecurityProtocolSSLProperties(properties, encryptionConfig); } } if (Objects.isNull(bootstrapServers) || bootstrapServers.isEmpty()) { throw new RuntimeException("Bootstrap servers are not specified"); } + properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers); } + private static boolean checkEncryptionType(final EncryptionConfig encryptionConfig, final EncryptionType encryptionType) { + return Objects.nonNull(encryptionConfig) && encryptionConfig.getType() == encryptionType; + } + public static GlueSchemaRegistryKafkaDeserializer getGlueSerializer(final KafkaConsumerConfig kafkaConsumerConfig) { SchemaConfig schemaConfig = kafkaConsumerConfig.getSchemaConfig(); if (Objects.isNull(schemaConfig) || schemaConfig.getType() != SchemaRegistryType.AWS_GLUE) { return null; } - Map configs = new HashMap(); + Map configs = new HashMap<>(); configs.put(AWSSchemaRegistryConstants.AWS_REGION, kafkaConsumerConfig.getAwsConfig().getRegion()); configs.put(AWSSchemaRegistryConstants.AVRO_RECORD_TYPE, AvroRecordType.GENERIC_RECORD.getName()); configs.put(AWSSchemaRegistryConstants.CACHE_TIME_TO_LIVE_MILLIS, "86400000"); diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferTest.java index 4ecb964b89..99f2afa76b 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferTest.java @@ -5,7 +5,9 @@ package org.opensearch.dataprepper.plugins.kafka.buffer; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; @@ -25,6 +27,8 @@ import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.buffer.blockingbuffer.BlockingBuffer; import org.opensearch.dataprepper.plugins.kafka.admin.KafkaAdminAccessor; +import org.opensearch.dataprepper.plugins.kafka.common.KafkaMdc; +import org.opensearch.dataprepper.plugins.kafka.common.thread.KafkaPluginThreadFactory; import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionType; @@ -35,6 +39,7 @@ import org.opensearch.dataprepper.plugins.kafka.producer.KafkaCustomProducerFactory; import org.opensearch.dataprepper.plugins.kafka.producer.ProducerWorker; import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; +import org.slf4j.MDC; import java.time.Duration; import java.util.Arrays; @@ -137,7 +142,7 @@ public KafkaBuffer createObjectUnderTest(final List consume final MockedConstruction producerFactoryMock = mockConstruction(KafkaCustomProducerFactory.class, (mock, context) -> { producerFactory = mock; - when(producerFactory.createProducer(any() ,any(), any(), isNull(), isNull(), any(), any(), anyBoolean())).thenReturn(producer); + when(producerFactory.createProducer(any(), isNull(), isNull(), any(), any(), anyBoolean())).thenReturn(producer); }); final MockedConstruction consumerFactoryMock = mockConstruction(KafkaCustomConsumerFactory.class, (mock, context) -> { @@ -151,8 +156,8 @@ public KafkaBuffer createObjectUnderTest(final List consume blockingBuffer = mock; })) { - executorsMockedStatic.when(() -> Executors.newFixedThreadPool(anyInt())).thenReturn(executorService); - return new KafkaBuffer(pluginSetting, bufferConfig, pluginFactory, acknowledgementSetManager, null, awsCredentialsSupplier, circuitBreaker); + executorsMockedStatic.when(() -> Executors.newFixedThreadPool(anyInt(), any(KafkaPluginThreadFactory.class))).thenReturn(executorService); + return new KafkaBuffer(pluginSetting, bufferConfig, acknowledgementSetManager, null, awsCredentialsSupplier, circuitBreaker); } } @@ -292,7 +297,7 @@ void test_kafkaBuffer_isEmpty_MultipleTopics_AllNotEmpty() { void test_kafkaBuffer_doCheckpoint() { kafkaBuffer = createObjectUnderTest(); kafkaBuffer.doCheckpoint(mock(CheckpointState.class)); - verify(blockingBuffer).doCheckpoint(any()); + verify(blockingBuffer).checkpoint(any()); } @Test @@ -353,4 +358,84 @@ public void testShutdown_InterruptedException() throws InterruptedException { verify(executorService).awaitTermination(eq(EXECUTOR_SERVICE_SHUTDOWN_TIMEOUT), eq(TimeUnit.SECONDS)); verify(executorService).shutdownNow(); } + + @Nested + class MdcTests { + private MockedStatic mdcMockedStatic; + + @BeforeEach + void setUp() { + mdcMockedStatic = mockStatic(MDC.class); + } + + @AfterEach + void tearDown() { + mdcMockedStatic.close(); + } + + @Test + void writeBytes_sets_and_clears_MDC() throws Exception { + createObjectUnderTest().writeBytes(new byte[] {}, UUID.randomUUID().toString(), 100); + + mdcMockedStatic.verify(() -> MDC.put(KafkaMdc.MDC_KAFKA_PLUGIN_KEY, "buffer")); + mdcMockedStatic.verify(() -> MDC.remove(KafkaMdc.MDC_KAFKA_PLUGIN_KEY)); + } + + @Test + void doWrite_sets_and_clears_MDC() throws Exception { + createObjectUnderTest().doWrite(mock(Record.class), 100); + + mdcMockedStatic.verify(() -> MDC.put(KafkaMdc.MDC_KAFKA_PLUGIN_KEY, "buffer")); + mdcMockedStatic.verify(() -> MDC.remove(KafkaMdc.MDC_KAFKA_PLUGIN_KEY)); + } + + @Test + void doWriteAll_sets_and_clears_MDC() throws Exception { + final List> records = Collections.singletonList(mock(Record.class)); + createObjectUnderTest().doWriteAll(records, 100); + + mdcMockedStatic.verify(() -> MDC.put(KafkaMdc.MDC_KAFKA_PLUGIN_KEY, "buffer")); + mdcMockedStatic.verify(() -> MDC.remove(KafkaMdc.MDC_KAFKA_PLUGIN_KEY)); + } + + @Test + void doRead_sets_and_clears_MDC() throws Exception { + createObjectUnderTest().doRead(100); + + mdcMockedStatic.verify(() -> MDC.put(KafkaMdc.MDC_KAFKA_PLUGIN_KEY, "buffer")); + mdcMockedStatic.verify(() -> MDC.remove(KafkaMdc.MDC_KAFKA_PLUGIN_KEY)); + } + + @Test + void doCheckpoint_sets_and_clears_MDC() throws Exception { + createObjectUnderTest().doCheckpoint(mock(CheckpointState.class)); + + mdcMockedStatic.verify(() -> MDC.put(KafkaMdc.MDC_KAFKA_PLUGIN_KEY, "buffer")); + mdcMockedStatic.verify(() -> MDC.remove(KafkaMdc.MDC_KAFKA_PLUGIN_KEY)); + } + + @Test + void postProcess_sets_and_clears_MDC() throws Exception { + createObjectUnderTest().postProcess(100L); + + mdcMockedStatic.verify(() -> MDC.put(KafkaMdc.MDC_KAFKA_PLUGIN_KEY, "buffer")); + mdcMockedStatic.verify(() -> MDC.remove(KafkaMdc.MDC_KAFKA_PLUGIN_KEY)); + } + + @Test + void isEmpty_sets_and_clears_MDC() throws Exception { + createObjectUnderTest().isEmpty(); + + mdcMockedStatic.verify(() -> MDC.put(KafkaMdc.MDC_KAFKA_PLUGIN_KEY, "buffer")); + mdcMockedStatic.verify(() -> MDC.remove(KafkaMdc.MDC_KAFKA_PLUGIN_KEY)); + } + + @Test + void shutdown_sets_and_clears_MDC() throws Exception { + createObjectUnderTest().shutdown(); + + mdcMockedStatic.verify(() -> MDC.put(KafkaMdc.MDC_KAFKA_PLUGIN_KEY, "buffer")); + mdcMockedStatic.verify(() -> MDC.remove(KafkaMdc.MDC_KAFKA_PLUGIN_KEY)); + } + } } diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/common/thread/KafkaPluginThreadFactoryTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/common/thread/KafkaPluginThreadFactoryTest.java new file mode 100644 index 0000000000..589f81a74c --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/common/thread/KafkaPluginThreadFactoryTest.java @@ -0,0 +1,107 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.common.thread; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.plugins.kafka.common.KafkaMdc; +import org.slf4j.MDC; + +import java.util.UUID; +import java.util.concurrent.ThreadFactory; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.notNullValue; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.not; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +class KafkaPluginThreadFactoryTest { + + @Mock + private ThreadFactory delegateThreadFactory; + @Mock + private Thread innerThread; + @Mock + private Runnable runnable; + private String pluginType; + + @BeforeEach + void setUp() { + pluginType = UUID.randomUUID().toString(); + + when(delegateThreadFactory.newThread(any(Runnable.class))).thenReturn(innerThread); + } + + + private KafkaPluginThreadFactory createObjectUnderTest() { + return new KafkaPluginThreadFactory(delegateThreadFactory, pluginType); + } + + @Test + void newThread_creates_thread_from_delegate() { + assertThat(createObjectUnderTest().newThread(runnable), equalTo(innerThread)); + } + + @Test + void newThread_creates_thread_with_name() { + final KafkaPluginThreadFactory objectUnderTest = createObjectUnderTest(); + + + final Thread thread1 = objectUnderTest.newThread(runnable); + assertThat(thread1, notNullValue()); + verify(thread1).setName(String.format("kafka-%s-1", pluginType)); + + final Thread thread2 = objectUnderTest.newThread(runnable); + assertThat(thread2, notNullValue()); + verify(thread2).setName(String.format("kafka-%s-2", pluginType)); + } + + @Test + void newThread_creates_thread_with_wrapping_runnable() { + createObjectUnderTest().newThread(runnable); + + final ArgumentCaptor actualRunnableCaptor = ArgumentCaptor.forClass(Runnable.class); + verify(delegateThreadFactory).newThread(actualRunnableCaptor.capture()); + + final Runnable actualRunnable = actualRunnableCaptor.getValue(); + + assertThat(actualRunnable, not(equalTo(runnable))); + + verifyNoInteractions(runnable); + actualRunnable.run(); + verify(runnable).run(); + } + + @Test + void newThread_creates_thread_that_calls_MDC_on_run() { + createObjectUnderTest().newThread(runnable); + + final ArgumentCaptor actualRunnableCaptor = ArgumentCaptor.forClass(Runnable.class); + verify(delegateThreadFactory).newThread(actualRunnableCaptor.capture()); + + final Runnable actualRunnable = actualRunnableCaptor.getValue(); + + final String[] actualKafkaPluginType = new String[1]; + doAnswer(a -> { + actualKafkaPluginType[0] = MDC.get(KafkaMdc.MDC_KAFKA_PLUGIN_KEY); + return null; + }).when(runnable).run(); + + actualRunnable.run(); + + assertThat(actualKafkaPluginType[0], equalTo(pluginType)); + } +} \ No newline at end of file diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerTest.java index 7d3a0f3fb9..968639f674 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerTest.java @@ -28,6 +28,7 @@ import org.opensearch.dataprepper.model.CheckpointState; import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; import org.opensearch.dataprepper.model.buffer.Buffer; +import org.opensearch.dataprepper.model.buffer.SizeOverflowException; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.record.Record; @@ -54,6 +55,7 @@ import static org.hamcrest.Matchers.hasEntry; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -70,6 +72,9 @@ public class KafkaCustomConsumerTest { private Buffer> buffer; + @Mock + private Buffer> mockBuffer; + @Mock private KafkaConsumerConfig sourceConfig; @@ -106,21 +111,31 @@ public class KafkaCustomConsumerTest { private Counter posCounter; @Mock private Counter negCounter; + @Mock + private Counter overflowCounter; private Duration delayTime; private double posCount; private double negCount; + private double overflowCount; + private boolean paused; + private boolean resumed; @BeforeEach public void setUp() { delayTime = Duration.ofMillis(10); + paused = false; + resumed = false; kafkaConsumer = mock(KafkaConsumer.class); topicMetrics = mock(KafkaTopicConsumerMetrics.class); counter = mock(Counter.class); posCounter = mock(Counter.class); + mockBuffer = mock(Buffer.class); negCounter = mock(Counter.class); + overflowCounter = mock(Counter.class); topicConfig = mock(TopicConsumerConfig.class); when(topicMetrics.getNumberOfPositiveAcknowledgements()).thenReturn(posCounter); when(topicMetrics.getNumberOfNegativeAcknowledgements()).thenReturn(negCounter); + when(topicMetrics.getNumberOfBufferSizeOverflows()).thenReturn(overflowCounter); when(topicMetrics.getNumberOfRecordsCommitted()).thenReturn(counter); when(topicMetrics.getNumberOfDeserializationErrors()).thenReturn(counter); when(topicConfig.getThreadWaitingTime()).thenReturn(Duration.ofSeconds(1)); @@ -128,6 +143,16 @@ public void setUp() { when(topicConfig.getAutoCommit()).thenReturn(false); when(kafkaConsumer.committed(any(TopicPartition.class))).thenReturn(null); + doAnswer((i)-> { + paused = true; + return null; + }).when(kafkaConsumer).pause(any()); + + doAnswer((i)-> { + resumed = true; + return null; + }).when(kafkaConsumer).resume(any()); + doAnswer((i)-> { posCount += 1.0; return null; @@ -136,6 +161,10 @@ public void setUp() { negCount += 1.0; return null; }).when(negCounter).increment(); + doAnswer((i)-> { + overflowCount += 1.0; + return null; + }).when(overflowCounter).increment(); doAnswer((i)-> {return posCount;}).when(posCounter).count(); doAnswer((i)-> {return negCount;}).when(negCounter).count(); callbackExecutor = Executors.newScheduledThreadPool(2); @@ -147,6 +176,11 @@ public void setUp() { when(topicConfig.getName()).thenReturn(TOPIC_NAME); } + public KafkaCustomConsumer createObjectUnderTestWithMockBuffer(String schemaType) { + return new KafkaCustomConsumer(kafkaConsumer, shutdownInProgress, mockBuffer, sourceConfig, topicConfig, schemaType, + acknowledgementSetManager, null, topicMetrics, pauseConsumePredicate); + } + public KafkaCustomConsumer createObjectUnderTest(String schemaType, boolean acknowledgementsEnabled) { when(sourceConfig.getAcknowledgementsEnabled()).thenReturn(acknowledgementsEnabled); return new KafkaCustomConsumer(kafkaConsumer, shutdownInProgress, buffer, sourceConfig, topicConfig, schemaType, @@ -162,6 +196,56 @@ private BlockingBuffer> getBuffer() { return new BlockingBuffer<>(pluginSetting); } + @Test + public void testBufferOverflowPauseResume() throws InterruptedException, Exception { + when(topicConfig.getMaxPollInterval()).thenReturn(Duration.ofMillis(4000)); + String topic = topicConfig.getName(); + consumerRecords = createPlainTextRecords(topic, 0L); + doAnswer((i)-> { + if (!paused && !resumed) + throw new SizeOverflowException("size overflow"); + buffer.writeAll(i.getArgument(0), i.getArgument(1)); + return null; + }).when(mockBuffer).writeAll(any(), anyInt()); + + doAnswer((i) -> { + if (paused && !resumed) + return List.of(); + return consumerRecords; + }).when(kafkaConsumer).poll(any(Duration.class)); + consumer = createObjectUnderTestWithMockBuffer("plaintext"); + try { + consumer.onPartitionsAssigned(List.of(new TopicPartition(topic, testPartition))); + consumer.consumeRecords(); + } catch (Exception e){} + assertTrue(paused); + assertTrue(resumed); + + final Map.Entry>, CheckpointState> bufferRecords = buffer.read(1000); + ArrayList> bufferedRecords = new ArrayList<>(bufferRecords.getKey()); + Assertions.assertEquals(consumerRecords.count(), bufferedRecords.size()); + Map offsetsToCommit = consumer.getOffsetsToCommit(); + Assertions.assertEquals(offsetsToCommit.size(), 1); + offsetsToCommit.forEach((topicPartition, offsetAndMetadata) -> { + Assertions.assertEquals(topicPartition.partition(), testPartition); + Assertions.assertEquals(topicPartition.topic(), topic); + Assertions.assertEquals(offsetAndMetadata.offset(), 2L); + }); + Assertions.assertEquals(consumer.getNumRecordsCommitted(), 2L); + + for (Record record: bufferedRecords) { + Event event = record.getData(); + String value1 = event.get(testKey1, String.class); + String value2 = event.get(testKey2, String.class); + assertTrue(value1 != null || value2 != null); + if (value1 != null) { + Assertions.assertEquals(value1, testValue1); + } + if (value2 != null) { + Assertions.assertEquals(value2, testValue2); + } + } + } @Test public void testPlainTextConsumeRecords() throws InterruptedException { String topic = topicConfig.getName(); diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/PauseConsumePredicateTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/PauseConsumePredicateTest.java index 0d2743b8ab..90efdfe02d 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/PauseConsumePredicateTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/PauseConsumePredicateTest.java @@ -51,4 +51,13 @@ void circuitBreakingPredicate_with_a_circuit_breaker_returns_predicate_with_paus assertThat(pauseConsumePredicate.pauseConsuming(), equalTo(isOpen)); } + + @Test + void circuitBreakingPredicate_with_a_circuit_breaker_returns_predicate_with_toString() { + final CircuitBreaker circuitBreaker = mock(CircuitBreaker.class); + + final PauseConsumePredicate pauseConsumePredicate = PauseConsumePredicate.circuitBreakingPredicate(circuitBreaker); + + assertThat(pauseConsumePredicate.toString(), equalTo("Circuit Breaker")); + } } \ No newline at end of file diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerFactoryTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerFactoryTest.java new file mode 100644 index 0000000000..4870574581 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerFactoryTest.java @@ -0,0 +1,118 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.producer; + +import org.apache.kafka.clients.producer.KafkaProducer; +import org.apache.kafka.common.serialization.Serializer; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.InOrder; +import org.mockito.Mock; +import org.mockito.MockedConstruction; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; +import org.opensearch.dataprepper.expression.ExpressionEvaluator; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.sink.SinkContext; +import org.opensearch.dataprepper.plugins.kafka.common.serialization.SerializationFactory; +import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicProducerConfig; +import org.opensearch.dataprepper.plugins.kafka.service.TopicService; +import org.opensearch.dataprepper.plugins.kafka.service.TopicServiceFactory; +import org.opensearch.dataprepper.plugins.kafka.sink.DLQSink; + +import java.util.Collections; +import java.util.Random; +import java.util.UUID; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.inOrder; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockConstruction; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +class KafkaCustomProducerFactoryTest { + private static final Random RANDOM = new Random(); + @Mock + private SerializationFactory serializationFactory; + @Mock + private AwsCredentialsSupplier awsCredentialsSupplier; + @Mock + private TopicServiceFactory topicServiceFactory; + + @Mock + private KafkaProducerConfig kafkaProducerConfig; + @Mock + private ExpressionEvaluator expressionEvaluator; + @Mock + private SinkContext sinkContext; + @Mock + private PluginMetrics pluginMetrics; + @Mock + private DLQSink dlqSink; + + @Mock + private TopicProducerConfig topicProducerConfig; + @Mock + private EncryptionConfig encryptionConfig; + + @BeforeEach + void setUp() { + when(kafkaProducerConfig.getTopic()).thenReturn(topicProducerConfig); + when(kafkaProducerConfig.getEncryptionConfig()).thenReturn(encryptionConfig); + when(kafkaProducerConfig.getBootstrapServers()).thenReturn(Collections.singletonList(UUID.randomUUID().toString())); + + final Serializer serializer = mock(Serializer.class); + when(serializationFactory.getSerializer(any())).thenReturn(serializer); + } + + private KafkaCustomProducerFactory createObjectUnderTest() { + return new KafkaCustomProducerFactory(serializationFactory, awsCredentialsSupplier, topicServiceFactory); + } + + @Test + void createProducer_does_not_create_TopicService_when_createTopic_is_false() { + final KafkaCustomProducerFactory objectUnderTest = createObjectUnderTest(); + try(final MockedConstruction ignored = mockConstruction(KafkaProducer.class)) { + objectUnderTest.createProducer(kafkaProducerConfig, expressionEvaluator, sinkContext, pluginMetrics, dlqSink, false); + } + + verify(topicServiceFactory, never()).createTopicService(any()); + } + + @Test + void createProducer_creates_TopicService_and_creates_topic_when_createTopic_is_true() { + final TopicService topicService = mock(TopicService.class); + when(topicServiceFactory.createTopicService(kafkaProducerConfig)).thenReturn(topicService); + + final String topicName = UUID.randomUUID().toString(); + final int numberOfPartitions = RANDOM.nextInt(1000); + final short replicationFactor = (short) RANDOM.nextInt(1000); + when(topicProducerConfig.getName()).thenReturn(topicName); + when(topicProducerConfig.getNumberOfPartitions()).thenReturn(numberOfPartitions); + when(topicProducerConfig.getReplicationFactor()).thenReturn(replicationFactor); + + + final KafkaCustomProducerFactory objectUnderTest = createObjectUnderTest(); + + when(topicProducerConfig.isCreateTopic()).thenReturn(true); + + try(final MockedConstruction ignored = mockConstruction(KafkaProducer.class)) { + objectUnderTest.createProducer(kafkaProducerConfig, expressionEvaluator, sinkContext, pluginMetrics, dlqSink, false); + } + + final InOrder inOrder = inOrder(topicService); + inOrder.verify(topicService).createTopic(topicName, numberOfPartitions, replicationFactor, null); + inOrder.verify(topicService).closeAdminClient(); + + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkTest.java index 54c58cb682..6d57c7803a 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkTest.java @@ -134,7 +134,7 @@ public void after() { private KafkaSink createObjectUnderTest() { final KafkaSink objectUnderTest; try(final MockedConstruction ignored = mockConstruction(KafkaCustomProducerFactory.class, (mock, context) -> { - when(mock.createProducer(any(), any(), any(), any(), any(), any(), any(), anyBoolean())).thenReturn(kafkaCustomProducer); + when(mock.createProducer(any(), any(), any(), any(), any(), anyBoolean())).thenReturn(kafkaCustomProducer); })) { objectUnderTest = new KafkaSink(pluginSetting, kafkaSinkConfig, pluginFactoryMock, pluginMetrics, mock(ExpressionEvaluator.class), sinkContext, awsCredentialsSupplier); } diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/CustomClientSslEngineFactoryTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/CustomClientSslEngineFactoryTest.java new file mode 100644 index 0000000000..a149fb840d --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/CustomClientSslEngineFactoryTest.java @@ -0,0 +1,41 @@ +package org.opensearch.dataprepper.plugins.kafka.util; + +import org.junit.jupiter.api.Test; +import javax.net.ssl.SSLEngine; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Collections; +import java.util.Objects; +import java.util.UUID; +import java.util.concurrent.ThreadLocalRandom; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.notNullValue; + +public class CustomClientSslEngineFactoryTest { + @Test + public void createClientSslEngine() { + try (final CustomClientSslEngineFactory customClientSslEngineFactory = new CustomClientSslEngineFactory()) { + final SSLEngine sslEngine = customClientSslEngineFactory.createClientSslEngine( + UUID.randomUUID().toString(), ThreadLocalRandom.current().nextInt(), UUID.randomUUID().toString()); + assertThat(sslEngine, is(notNullValue())); + } + } + + @Test + public void createClientSslEngineWithConfig() throws IOException { + try (final CustomClientSslEngineFactory customClientSslEngineFactory = new CustomClientSslEngineFactory()) { + final Path certFilePath = Path.of(Objects.requireNonNull(getClass().getClassLoader() + .getResource("test_cert.crt")).getPath()); + + final String certificateContent = Files.readString(certFilePath); + customClientSslEngineFactory.configure(Collections.singletonMap("certificateContent", certificateContent)); + final SSLEngine sslEngine = customClientSslEngineFactory.createClientSslEngine( + UUID.randomUUID().toString(), ThreadLocalRandom.current().nextInt(), UUID.randomUUID().toString()); + assertThat(sslEngine, is(notNullValue())); + } + } +} diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaProducerMetricsTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaProducerMetricsTest.java new file mode 100644 index 0000000000..adcd89d3d7 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaProducerMetricsTest.java @@ -0,0 +1,112 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.util; + +import org.apache.kafka.clients.producer.KafkaProducer; +import org.apache.kafka.common.Metric; +import org.apache.kafka.common.MetricName; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.api.extension.ExtensionContext; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.ArgumentsProvider; +import org.junit.jupiter.params.provider.ArgumentsSource; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.metrics.PluginMetrics; + +import java.util.HashMap; +import java.util.Map; +import java.util.Random; +import java.util.UUID; +import java.util.function.ToDoubleFunction; +import java.util.stream.IntStream; +import java.util.stream.Stream; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.junit.jupiter.params.provider.Arguments.arguments; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +class KafkaProducerMetricsTest { + @Mock + private PluginMetrics pluginMetrics; + @Mock + private KafkaProducer kafkaProducer; + + private Map kafkaMetricsMap; + private Map knownMetricsMap; + + @BeforeEach + void setUp() { + kafkaMetricsMap = new HashMap<>(); + knownMetricsMap = new HashMap<>(); + when(kafkaProducer.metrics()).thenReturn((Map) kafkaMetricsMap); + + KafkaProducerMetrics.METRICS_NAME_MAP.keySet() + .stream() + .map(KafkaProducerMetricsTest::createKafkaMetric) + .forEach(metricName -> { + final Metric metric = mock(Metric.class); + knownMetricsMap.put(metricName.name(), metric); + kafkaMetricsMap.put(metricName, metric); + }); + IntStream.range(0, 5) + .mapToObj(ignored -> UUID.randomUUID().toString()) + .map(KafkaProducerMetricsTest::createKafkaMetric) + .forEach(metricName -> kafkaMetricsMap.put(metricName, mock(Metric.class))); + } + + @Test + void registerProducer_creates_gauges_for_each_metric_from_the_map() { + KafkaProducerMetrics.registerProducer(pluginMetrics, kafkaProducer); + + verify(pluginMetrics, times(KafkaProducerMetrics.METRICS_NAME_MAP.size())).gauge(anyString(), any(), any()); + } + + @ParameterizedTest + @ArgumentsSource(RegisteredMetricsArgumentsProvider.class) + void registerProduct_creates_expected_gauge(final String kafkaName, final String expectedDataPrepperName) { + KafkaProducerMetrics.registerProducer(pluginMetrics, kafkaProducer); + + final Metric metric = knownMetricsMap.get(kafkaName); + final ArgumentCaptor metricFunctionArgumentCaptor = ArgumentCaptor.forClass(ToDoubleFunction.class); + + verify(pluginMetrics).gauge(eq(expectedDataPrepperName), eq(metric), metricFunctionArgumentCaptor.capture()); + + final ToDoubleFunction actualMetricDoubleFunction = metricFunctionArgumentCaptor.getValue(); + + final Random random = new Random(); + final double metricValue = random.nextDouble(); + when(metric.metricValue()).thenReturn(metricValue); + assertThat(actualMetricDoubleFunction.applyAsDouble(metric), equalTo(metricValue)); + } + + private static MetricName createKafkaMetric(final String name) { + final MetricName metricName = mock(MetricName.class); + when(metricName.name()).thenReturn(name); + return metricName; + } + + static class RegisteredMetricsArgumentsProvider implements ArgumentsProvider { + @Override + public Stream provideArguments(final ExtensionContext extensionContext) throws Exception { + return KafkaProducerMetrics.METRICS_NAME_MAP.entrySet() + .stream() + .map(entry -> arguments(entry.getKey(), entry.getValue())); + } + } +} \ No newline at end of file diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaSecurityConfigurerTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaSecurityConfigurerTest.java new file mode 100644 index 0000000000..e2506fafd1 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaSecurityConfigurerTest.java @@ -0,0 +1,107 @@ +package org.opensearch.dataprepper.plugins.kafka.util; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.plugins.kafka.source.KafkaSourceConfig; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.yaml.snakeyaml.Yaml; + +import java.io.FileReader; +import java.io.IOException; +import java.io.Reader; +import java.io.StringReader; +import java.util.Map; +import java.util.Objects; +import java.util.Properties; + +import static org.hamcrest.CoreMatchers.nullValue; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.CoreMatchers.is; + +public class KafkaSecurityConfigurerTest { + private static final Logger LOG = LoggerFactory.getLogger(KafkaSecurityConfigurerTest.class); + @Test + public void testSetAuthPropertiesWithSaslPlainCertificate() throws Exception { + final Properties props = new Properties(); + final KafkaSourceConfig kafkaSourceConfig = createKafkaSinkConfig("kafka-pipeline-sasl-ssl-certificate-content.yaml"); + KafkaSecurityConfigurer.setAuthProperties(props, kafkaSourceConfig, LOG); + assertThat(props.getProperty("sasl.mechanism"), is("PLAIN")); + assertThat(props.getProperty("security.protocol"), is("SASL_SSL")); + assertThat(props.getProperty("certificateContent"), is("CERTIFICATE_DATA")); + assertThat(props.getProperty("ssl.truststore.location"), is(nullValue())); + assertThat(props.getProperty("ssl.truststore.password"), is(nullValue())); + assertThat(props.get("ssl.engine.factory.class"), is(CustomClientSslEngineFactory.class)); + } + + @Test + public void testSetAuthPropertiesWithNoAuthSsl() throws Exception { + final Properties props = new Properties(); + final KafkaSourceConfig kafkaSourceConfig = createKafkaSinkConfig("kafka-pipeline-no-auth-ssl.yaml"); + KafkaSecurityConfigurer.setAuthProperties(props, kafkaSourceConfig, LOG); + assertThat(props.getProperty("sasl.mechanism"), is(nullValue())); + assertThat(props.getProperty("security.protocol"), is("SSL")); + assertThat(props.getProperty("certificateContent"), is("CERTIFICATE_DATA")); + assertThat(props.get("ssl.engine.factory.class"), is(CustomClientSslEngineFactory.class)); + } + @Test + public void testSetAuthPropertiesWithNoAuthSslNone() throws Exception { + final Properties props = new Properties(); + final KafkaSourceConfig kafkaSourceConfig = createKafkaSinkConfig("kafka-pipeline-no-auth-ssl-none.yaml"); + KafkaSecurityConfigurer.setAuthProperties(props, kafkaSourceConfig, LOG); + assertThat(props.getProperty("sasl.mechanism"), is(nullValue())); + assertThat(props.getProperty("security.protocol"), is(nullValue())); + assertThat(props.getProperty("certificateContent"), is(nullValue())); + assertThat(props.get("ssl.engine.factory.class"), is(nullValue())); + } + + @Test + public void testSetAuthPropertiesWithNoAuthInsecure() throws Exception { + final Properties props = new Properties(); + final KafkaSourceConfig kafkaSourceConfig = createKafkaSinkConfig("kafka-pipeline-auth-insecure.yaml"); + KafkaSecurityConfigurer.setAuthProperties(props, kafkaSourceConfig, LOG); + assertThat(props.getProperty("sasl.mechanism"), is("PLAIN")); + assertThat(props.getProperty("security.protocol"), is("SASL_PLAINTEXT")); + assertThat(props.getProperty("certificateContent"), is(nullValue())); + assertThat(props.get("ssl.engine.factory.class"), is(InsecureSslEngineFactory.class)); + } + @Test + public void testSetAuthPropertiesAuthSslWithTrustStore() throws Exception { + final Properties props = new Properties(); + final KafkaSourceConfig kafkaSourceConfig = createKafkaSinkConfig("kafka-pipeline-sasl-ssl-truststore.yaml"); + KafkaSecurityConfigurer.setAuthProperties(props, kafkaSourceConfig, LOG); + assertThat(props.getProperty("sasl.mechanism"), is("PLAIN")); + assertThat(props.getProperty("security.protocol"), is("SASL_SSL")); + assertThat(props.getProperty("certificateContent"), is(nullValue())); + assertThat(props.getProperty("ssl.truststore.location"), is("some-file-path")); + assertThat(props.getProperty("ssl.truststore.password"), is("some-password")); + assertThat(props.get("ssl.engine.factory.class"), is(nullValue())); + } + + @Test + public void testSetAuthPropertiesAuthSslWithNoCertContentNoTrustStore() throws Exception { + final Properties props = new Properties(); + final KafkaSourceConfig kafkaSourceConfig = createKafkaSinkConfig("kafka-pipeline-sasl-ssl-no-cert-content-no-truststore.yaml"); + KafkaSecurityConfigurer.setAuthProperties(props, kafkaSourceConfig, LOG); + assertThat(props.getProperty("sasl.mechanism"), is("PLAIN")); + assertThat(props.getProperty("security.protocol"), is("SASL_SSL")); + assertThat(props.getProperty("certificateContent"), is(nullValue())); + assertThat(props.getProperty("ssl.truststore.location"), is(nullValue())); + assertThat(props.getProperty("ssl.truststore.password"), is(nullValue())); + assertThat(props.get("ssl.engine.factory.class"), is(nullValue())); + } + + private KafkaSourceConfig createKafkaSinkConfig(final String fileName) throws IOException { + final Yaml yaml = new Yaml(); + final FileReader fileReader = new FileReader(Objects.requireNonNull(getClass().getClassLoader() + .getResource(fileName)).getFile()); + final Map>>> data = yaml.load(fileReader); + final Map>> logPipelineMap = data.get("log-pipeline"); + final Map> sourceMap = logPipelineMap.get("source"); + final Map kafkaConfigMap = sourceMap.get("kafka"); + final ObjectMapper mapper = new ObjectMapper(); + final String json = mapper.writeValueAsString(kafkaConfigMap); + final Reader reader = new StringReader(json); + return mapper.readValue(reader, KafkaSourceConfig.class); + } +} diff --git a/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-auth-insecure.yaml b/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-auth-insecure.yaml new file mode 100644 index 0000000000..b62ebfdf08 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-auth-insecure.yaml @@ -0,0 +1,19 @@ +log-pipeline : + source: + kafka: + bootstrap_servers: + - "localhost:9092" + encryption: + type: "NONE" + certificate_content: "CERTIFICATE_DATA" + insecure: "true" + authentication: + sasl: + plaintext: + username: username + password: password + topics: + - name: "quickstart-events" + group_id: "groupdID1" + sink: + stdout: \ No newline at end of file diff --git a/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-no-auth-ssl-none.yaml b/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-no-auth-ssl-none.yaml new file mode 100644 index 0000000000..4bae77dc5e --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-no-auth-ssl-none.yaml @@ -0,0 +1,13 @@ +log-pipeline : + source: + kafka: + bootstrap_servers: + - "localhost:9092" + encryption: + type: "NONE" + certificate_content: "CERTIFICATE_DATA" + topics: + - name: "quickstart-events" + group_id: "groupdID1" + sink: + stdout: \ No newline at end of file diff --git a/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-no-auth-ssl.yaml b/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-no-auth-ssl.yaml new file mode 100644 index 0000000000..5d894ab0cf --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-no-auth-ssl.yaml @@ -0,0 +1,13 @@ +log-pipeline : + source: + kafka: + bootstrap_servers: + - "localhost:9092" + encryption: + type: "SSL" + certificate_content: "CERTIFICATE_DATA" + topics: + - name: "quickstart-events" + group_id: "groupdID1" + sink: + stdout: \ No newline at end of file diff --git a/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-sasl-ssl-certificate-content.yaml b/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-sasl-ssl-certificate-content.yaml new file mode 100644 index 0000000000..d1fe45810d --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-sasl-ssl-certificate-content.yaml @@ -0,0 +1,18 @@ +log-pipeline : + source: + kafka: + bootstrap_servers: + - "localhost:9092" + encryption: + type: "SSL" + certificate_content: "CERTIFICATE_DATA" + authentication: + sasl: + plaintext: + username: username + password: password + topics: + - name: "quickstart-events" + group_id: "groupdID1" + sink: + stdout: \ No newline at end of file diff --git a/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-sasl-ssl-no-cert-content-no-truststore.yaml b/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-sasl-ssl-no-cert-content-no-truststore.yaml new file mode 100644 index 0000000000..4a16c918a1 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-sasl-ssl-no-cert-content-no-truststore.yaml @@ -0,0 +1,17 @@ +log-pipeline : + source: + kafka: + bootstrap_servers: + - "localhost:9092" + encryption: + type: "SSL" + authentication: + sasl: + plaintext: + username: username + password: password + topics: + - name: "quickstart-events" + group_id: "groupdID1" + sink: + stdout: \ No newline at end of file diff --git a/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-sasl-ssl-truststore.yaml b/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-sasl-ssl-truststore.yaml new file mode 100644 index 0000000000..fe499a6901 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-sasl-ssl-truststore.yaml @@ -0,0 +1,19 @@ +log-pipeline : + source: + kafka: + bootstrap_servers: + - "localhost:9092" + encryption: + type: "SSL" + trust_store_file_path: "some-file-path" + trust_store_password: "some-password" + authentication: + sasl: + plaintext: + username: username + password: password + topics: + - name: "quickstart-events" + group_id: "groupdID1" + sink: + stdout: \ No newline at end of file diff --git a/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-sasl-ssl.yaml b/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-sasl-ssl.yaml new file mode 100644 index 0000000000..d1fe45810d --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/resources/kafka-pipeline-sasl-ssl.yaml @@ -0,0 +1,18 @@ +log-pipeline : + source: + kafka: + bootstrap_servers: + - "localhost:9092" + encryption: + type: "SSL" + certificate_content: "CERTIFICATE_DATA" + authentication: + sasl: + plaintext: + username: username + password: password + topics: + - name: "quickstart-events" + group_id: "groupdID1" + sink: + stdout: \ No newline at end of file diff --git a/data-prepper-plugins/kafka-plugins/src/test/resources/test_cert.crt b/data-prepper-plugins/kafka-plugins/src/test/resources/test_cert.crt new file mode 100644 index 0000000000..26c78d1411 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/resources/test_cert.crt @@ -0,0 +1,14 @@ +-----BEGIN CERTIFICATE----- +MIICHTCCAYYCCQD4hqYeYDQZADANBgkqhkiG9w0BAQUFADBSMQswCQYDVQQGEwJV +UzELMAkGA1UECAwCVFgxDzANBgNVBAcMBkF1c3RpbjEPMA0GA1UECgwGQW1hem9u +MRQwEgYDVQQLDAtEYXRhcHJlcHBlcjAgFw0yMTA2MjUxOTIzMTBaGA8yMTIxMDYw +MTE5MjMxMFowUjELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAlRYMQ8wDQYDVQQHDAZB +dXN0aW4xDzANBgNVBAoMBkFtYXpvbjEUMBIGA1UECwwLRGF0YXByZXBwZXIwgZ8w +DQYJKoZIhvcNAQEBBQADgY0AMIGJAoGBAKrb3YhdKbQ5PtLHall10iLZC9ZdDVrq +HOvqVSM8NHlL8f82gJ3l0n9k7hYc5eKisutaS9eDTmJ+Dnn8xn/qPSKTIq9Wh+OZ +O+e9YEEpI/G4F9KpGULgMyRg9sJK0GlZdEt9o5GJNJIJUkptJU5eiLuE0IV+jyJo +Nvm8OE6EJPqxAgMBAAEwDQYJKoZIhvcNAQEFBQADgYEAjgnX5n/Tt7eo9uakIGAb +uBhvYdR8JqKXqF9rjFJ/MIK7FdQSF/gCdjnvBhzLlZFK/Nb6MGKoSKm5Lcr75LgC +FyhIwp3WlqQksiMFnOypYVY71vqDgj6UKdMaOBgthsYhngj8lC+wsVzWqQvkJ2Qg +/GAIzJwiZfXiaevQHRk79qI= +-----END CERTIFICATE----- diff --git a/data-prepper-plugins/key-value-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/keyvalue/KeyValueProcessor.java b/data-prepper-plugins/key-value-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/keyvalue/KeyValueProcessor.java index e7121456ce..0ccfa90baa 100644 --- a/data-prepper-plugins/key-value-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/keyvalue/KeyValueProcessor.java +++ b/data-prepper-plugins/key-value-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/keyvalue/KeyValueProcessor.java @@ -35,6 +35,8 @@ import java.util.Stack; import java.util.ArrayList; +import static org.opensearch.dataprepper.logging.DataPrepperMarkers.EVENT; + @DataPrepperPlugin(name = "key_value", pluginType = Processor.class, pluginConfigurationType = KeyValueProcessorConfig.class) public class KeyValueProcessor extends AbstractProcessor, Record> { private static final Logger LOG = LoggerFactory.getLogger(KeyValueProcessor.class); @@ -235,38 +237,44 @@ public Collection> doExecute(final Collection> recor for (final Record record : records) { final Map outputMap = new HashMap<>(); final Event recordEvent = record.getData(); - final String groupsRaw = recordEvent.get(keyValueProcessorConfig.getSource(), String.class); - if (groupsRaw == null) { - continue; - } - final String[] groups = fieldDelimiterPattern.split(groupsRaw, 0); - if (keyValueProcessorConfig.getRecursive()) { - try { - JsonNode recursedTree = recurse(groupsRaw, mapper); - outputMap.putAll(createRecursedMap(recursedTree, mapper)); - } catch (Exception e) { - LOG.error("Recursive parsing ran into an unexpected error, treating message as non-recursive", e); - recordEvent.getMetadata().addTags(tagsOnFailure); + try { + final String groupsRaw = recordEvent.get(keyValueProcessorConfig.getSource(), String.class); + if (groupsRaw == null) { + continue; } - } else { - try { - outputMap.putAll(createNonRecursedMap(groups)); - } catch (Exception e) { - LOG.error("Non-recursive parsing ran into an unexpected error", e); - recordEvent.getMetadata().addTags(tagsOnFailure); + final String[] groups = fieldDelimiterPattern.split(groupsRaw, 0); + + if (keyValueProcessorConfig.getRecursive()) { + try { + JsonNode recursedTree = recurse(groupsRaw, mapper); + outputMap.putAll(createRecursedMap(recursedTree, mapper)); + } catch (Exception e) { + LOG.error("Recursive parsing ran into an unexpected error, treating message as non-recursive", e); + recordEvent.getMetadata().addTags(tagsOnFailure); + } + } else { + try { + outputMap.putAll(createNonRecursedMap(groups)); + } catch (Exception e) { + LOG.error("Non-recursive parsing ran into an unexpected error", e); + recordEvent.getMetadata().addTags(tagsOnFailure); + } } - } - final Map processedMap = executeConfigs(outputMap); + final Map processedMap = executeConfigs(outputMap); - if (Objects.isNull(keyValueProcessorConfig.getDestination())) { - writeToRoot(recordEvent, processedMap); - } else { - if (keyValueProcessorConfig.getOverwriteIfDestinationExists() || - !recordEvent.containsKey(keyValueProcessorConfig.getDestination())) { - recordEvent.put(keyValueProcessorConfig.getDestination(), processedMap); + if (Objects.isNull(keyValueProcessorConfig.getDestination())) { + writeToRoot(recordEvent, processedMap); + } else { + if (keyValueProcessorConfig.getOverwriteIfDestinationExists() || + !recordEvent.containsKey(keyValueProcessorConfig.getDestination())) { + recordEvent.put(keyValueProcessorConfig.getDestination(), processedMap); + } } + } catch (final Exception e) { + LOG.error(EVENT, "There was an exception while processing on Event [{}]: ", recordEvent, e); + recordEvent.getMetadata().addTags(tagsOnFailure); } } diff --git a/data-prepper-plugins/mapdb-processor-state/build.gradle b/data-prepper-plugins/mapdb-processor-state/build.gradle index 56d216fb17..9c0c520346 100644 --- a/data-prepper-plugins/mapdb-processor-state/build.gradle +++ b/data-prepper-plugins/mapdb-processor-state/build.gradle @@ -7,10 +7,6 @@ plugins { id 'java' } -repositories { - mavenCentral() -} - dependencies { implementation project(':data-prepper-api') implementation project(':data-prepper-plugins:common') diff --git a/data-prepper-plugins/mutate-event-processors/README.md b/data-prepper-plugins/mutate-event-processors/README.md index b456552399..8acfd7edb7 100644 --- a/data-prepper-plugins/mutate-event-processors/README.md +++ b/data-prepper-plugins/mutate-event-processors/README.md @@ -577,14 +577,74 @@ The processed event will have the following data: } ``` +If we enable `convert_field_to_list` option: +```yaml +... +processor: + - map_to_list: + source: "my-map" + target: "my-list" + convert_field_to_list: true +... +``` +the processed event will have the following data: +```json +{ + "my-list": [ + ["key1", "value1"], + ["key2", "value2"], + ["key3", "value3"] + ], + "my-map": { + "key1": "value1", + "key2": "value2", + "key3": "value3" + } +} +``` + +If source is set to empty string (""), it will use the event root as source. +```yaml +... +processor: + - map_to_list: + source: "" + target: "my-list" + convert_field_to_list: true +... +``` +Input data like this: +```json +{ + "key1": "value1", + "key2": "value2", + "key3": "value3" +} +``` +will end up with this after processing: +```json +{ + "my-list": [ + ["key1", "value1"], + ["key2", "value2"], + ["key3", "value3"] + ], + "key1": "value1", + "key2": "value2", + "key3": "value3" +} +``` + ### Configuration -* `source` - (required): the source map to perform the operation +* `source` - (required): the source map to perform the operation; If set to empty string (""), it will use the event root as source. * `target` - (required): the target list to put the converted list * `key_name` - (optional): the key name of the field to hold the original key, default is "key" * `value_name` - (optional): the key name of the field to hold the original value, default is "value" * `exclude_keys` - (optional): the keys in source map that will be excluded from processing, default is empty list * `remove_processed_fields` - (optional): default is false; if true, will remove processed fields from source map * `map_to_list_when` - (optional): used to configure a condition for event processing based on certain property of the incoming event. Default is null (all events will be processed). +* `convert_field_to_list` - (optional): default to false; if true, will convert fields to lists instead of objects +* `tags_on_failure` - (optional): a list of tags to add to event metadata when the event fails to process ## Developer Guide diff --git a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/AddEntryProcessor.java b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/AddEntryProcessor.java index daaa739e84..b501eea7e0 100644 --- a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/AddEntryProcessor.java +++ b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/AddEntryProcessor.java @@ -17,10 +17,13 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.ArrayList; import java.util.Collection; import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.function.Consumer; +import java.util.function.Supplier; import static org.opensearch.dataprepper.logging.DataPrepperMarkers.EVENT; @@ -43,42 +46,49 @@ public Collection> doExecute(final Collection> recor for(final Record record : records) { final Event recordEvent = record.getData(); - for(AddEntryProcessorConfig.Entry entry : entries) { + try { + for (AddEntryProcessorConfig.Entry entry : entries) { - if (Objects.nonNull(entry.getAddWhen()) && !expressionEvaluator.evaluateConditional(entry.getAddWhen(), recordEvent)) { - continue; - } - - try { - final String key = entry.getKey(); - final String metadataKey = entry.getMetadataKey(); - Object value; - if (!Objects.isNull(entry.getValueExpression())) { - value = expressionEvaluator.evaluate(entry.getValueExpression(), recordEvent); - } else if (!Objects.isNull(entry.getFormat())) { - try { - value = recordEvent.formatString(entry.getFormat()); - } catch (final EventKeyNotFoundException e) { - value = null; - } - } else { - value = entry.getValue(); + if (Objects.nonNull(entry.getAddWhen()) && !expressionEvaluator.evaluateConditional(entry.getAddWhen(), recordEvent)) { + continue; } - if (!Objects.isNull(key)) { - if (!recordEvent.containsKey(key) || entry.getOverwriteIfKeyExists()) { - recordEvent.put(key, value); + + try { + final String key = (entry.getKey() == null) ? null : recordEvent.formatString(entry.getKey(), expressionEvaluator); + final String metadataKey = entry.getMetadataKey(); + Object value; + if (!Objects.isNull(entry.getValueExpression())) { + value = expressionEvaluator.evaluate(entry.getValueExpression(), recordEvent); + } else if (!Objects.isNull(entry.getFormat())) { + try { + value = recordEvent.formatString(entry.getFormat()); + } catch (final EventKeyNotFoundException e) { + value = null; + } + } else { + value = entry.getValue(); } - } else { - Map attributes = recordEvent.getMetadata().getAttributes(); - if (!attributes.containsKey(metadataKey) || entry.getOverwriteIfKeyExists()) { - recordEvent.getMetadata().setAttribute(metadataKey, value); - + if (!Objects.isNull(key)) { + if (!recordEvent.containsKey(key) || entry.getOverwriteIfKeyExists()) { + recordEvent.put(key, value); + } else if (recordEvent.containsKey(key) && entry.getAppendIfKeyExists()) { + mergeValueToEvent(recordEvent, key, value); + } + } else { + Map attributes = recordEvent.getMetadata().getAttributes(); + if (!attributes.containsKey(metadataKey) || entry.getOverwriteIfKeyExists()) { + recordEvent.getMetadata().setAttribute(metadataKey, value); + } else if (attributes.containsKey(metadataKey) && entry.getAppendIfKeyExists()) { + mergeValueToEventMetadata(recordEvent, metadataKey, value); + } } + } catch (Exception e) { + LOG.error(EVENT, "Error adding entry to record [{}] with key [{}], metadataKey [{}], value_expression [{}] format [{}], value [{}]", + recordEvent, entry.getKey(), entry.getMetadataKey(), entry.getValueExpression(), entry.getFormat(), entry.getValue(), e); } - } catch (Exception e) { - LOG.error(EVENT, "Error adding entry to record [{}] with key [{}], metadataKey [{}], value_expression [{}] format [{}], value [{}]", - recordEvent, entry.getKey(), entry.getMetadataKey(), entry.getValueExpression(), entry.getFormat(), entry.getValue(), e); } + } catch(final Exception e){ + LOG.error(EVENT, "There was an exception while processing Event [{}]", recordEvent, e); } } @@ -97,4 +107,25 @@ public boolean isReadyForShutdown() { @Override public void shutdown() { } + + private void mergeValueToEvent(final Event recordEvent, final String key, final Object value) { + mergeValue(value, () -> recordEvent.get(key, Object.class), newValue -> recordEvent.put(key, newValue)); + } + + private void mergeValueToEventMetadata(final Event recordEvent, final String key, final Object value) { + mergeValue(value, () -> recordEvent.getMetadata().getAttribute(key), newValue -> recordEvent.getMetadata().setAttribute(key, newValue)); + } + + private void mergeValue(final Object value, Supplier getter, Consumer setter) { + final Object currentValue = getter.get(); + final List mergedValue = new ArrayList<>(); + if (currentValue instanceof List) { + mergedValue.addAll((List) currentValue); + } else { + mergedValue.add(currentValue); + } + + mergedValue.add(value); + setter.accept(mergedValue); + } } diff --git a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/AddEntryProcessorConfig.java b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/AddEntryProcessorConfig.java index 90b114d62f..81f6bbab34 100644 --- a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/AddEntryProcessorConfig.java +++ b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/AddEntryProcessorConfig.java @@ -34,6 +34,9 @@ public static class Entry { @JsonProperty("overwrite_if_key_exists") private boolean overwriteIfKeyExists = false; + @JsonProperty("append_if_key_exists") + private boolean appendIfKeyExists = false; + public String getKey() { return key; } @@ -58,6 +61,10 @@ public boolean getOverwriteIfKeyExists() { return overwriteIfKeyExists; } + public boolean getAppendIfKeyExists() { + return appendIfKeyExists; + } + public String getAddWhen() { return addWhen; } @AssertTrue(message = "Either value or format or expression must be specified, and only one of them can be specified") @@ -65,12 +72,18 @@ public boolean hasValueOrFormatOrExpression() { return Stream.of(value, format, valueExpression).filter(n -> n!=null).count() == 1; } + @AssertTrue(message = "overwrite_if_key_exists and append_if_key_exists can not be set to true at the same time.") + boolean overwriteAndAppendNotBothSet() { + return !(overwriteIfKeyExists && appendIfKeyExists); + } + public Entry(final String key, final String metadataKey, final Object value, final String format, final String valueExpression, final boolean overwriteIfKeyExists, + final boolean appendIfKeyExists, final String addWhen) { if (key != null && metadataKey != null) { @@ -85,6 +98,7 @@ public Entry(final String key, this.format = format; this.valueExpression = valueExpression; this.overwriteIfKeyExists = overwriteIfKeyExists; + this.appendIfKeyExists = appendIfKeyExists; this.addWhen = addWhen; } diff --git a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/ConvertEntryTypeProcessor.java b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/ConvertEntryTypeProcessor.java index 8ba0400734..95cbd9f714 100644 --- a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/ConvertEntryTypeProcessor.java +++ b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/ConvertEntryTypeProcessor.java @@ -56,24 +56,30 @@ public Collection> doExecute(final Collection> recor for(final Record record : records) { final Event recordEvent = record.getData(); - if (Objects.nonNull(convertWhen) && !expressionEvaluator.evaluateConditional(convertWhen, recordEvent)) { - continue; - } + try { + + if (Objects.nonNull(convertWhen) && !expressionEvaluator.evaluateConditional(convertWhen, recordEvent)) { + continue; + } - for(final String key : convertEntryKeys) { - Object keyVal = recordEvent.get(key, Object.class); - if (keyVal != null) { - if (!nullValues.contains(keyVal.toString())) { - try { - recordEvent.put(key, converter.convert(keyVal)); - } catch (final RuntimeException e) { - LOG.error(EVENT, "Unable to convert key: {} with value: {} to {}", key, keyVal, type, e); - recordEvent.getMetadata().addTags(tagsOnFailure); + for (final String key : convertEntryKeys) { + Object keyVal = recordEvent.get(key, Object.class); + if (keyVal != null) { + if (!nullValues.contains(keyVal.toString())) { + try { + recordEvent.put(key, converter.convert(keyVal)); + } catch (final RuntimeException e) { + LOG.error(EVENT, "Unable to convert key: {} with value: {} to {}", key, keyVal, type, e); + recordEvent.getMetadata().addTags(tagsOnFailure); + } + } else { + recordEvent.delete(key); } - } else { - recordEvent.delete(key); } } + } catch (final Exception e) { + LOG.error(EVENT, "There was an exception while processing Event [{}]", recordEvent, e); + recordEvent.getMetadata().addTags(tagsOnFailure); } } return records; diff --git a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/CopyValueProcessor.java b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/CopyValueProcessor.java index 9883f71c29..845ae40e38 100644 --- a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/CopyValueProcessor.java +++ b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/CopyValueProcessor.java @@ -23,6 +23,8 @@ import java.util.Map; import java.util.Objects; +import static org.opensearch.dataprepper.logging.DataPrepperMarkers.EVENT; + @DataPrepperPlugin(name = "copy_values", pluginType = Processor.class, pluginConfigurationType = CopyValueProcessorConfig.class) public class CopyValueProcessor extends AbstractProcessor, Record> { private static final Logger LOG = LoggerFactory.getLogger(CopyValueProcessor.class); @@ -41,8 +43,9 @@ public CopyValueProcessor(final PluginMetrics pluginMetrics, final CopyValueProc @Override public Collection> doExecute(final Collection> records) { for(final Record record : records) { + final Event recordEvent = record.getData(); + try { - final Event recordEvent = record.getData(); if (config.getFromList() != null || config.getToList() != null) { // Copying entries between lists if (recordEvent.containsKey(config.getToList()) && !config.getOverwriteIfToListExists()) { @@ -80,9 +83,8 @@ public Collection> doExecute(final Collection> recor } } } - } catch (Exception e) { - LOG.error("Fail to perform copy values operation", e); - //TODO: add tagging on failure + } catch (final Exception e) { + LOG.error(EVENT, "There was an exception while processing Event [{}]", recordEvent, e); } } diff --git a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/DeleteEntryProcessor.java b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/DeleteEntryProcessor.java index 5bb01fa7a7..d7c902a32c 100644 --- a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/DeleteEntryProcessor.java +++ b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/DeleteEntryProcessor.java @@ -13,12 +13,18 @@ import org.opensearch.dataprepper.model.processor.AbstractProcessor; import org.opensearch.dataprepper.model.processor.Processor; import org.opensearch.dataprepper.model.record.Record; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.util.Collection; import java.util.Objects; +import static org.opensearch.dataprepper.logging.DataPrepperMarkers.EVENT; + @DataPrepperPlugin(name = "delete_entries", pluginType = Processor.class, pluginConfigurationType = DeleteEntryProcessorConfig.class) public class DeleteEntryProcessor extends AbstractProcessor, Record> { + + private static final Logger LOG = LoggerFactory.getLogger(DeleteEntryProcessor.class); private final String[] entries; private final String deleteWhen; @@ -37,13 +43,17 @@ public Collection> doExecute(final Collection> recor for(final Record record : records) { final Event recordEvent = record.getData(); - if (Objects.nonNull(deleteWhen) && !expressionEvaluator.evaluateConditional(deleteWhen, recordEvent)) { - continue; - } + try { + if (Objects.nonNull(deleteWhen) && !expressionEvaluator.evaluateConditional(deleteWhen, recordEvent)) { + continue; + } - for(String entry : entries) { - recordEvent.delete(entry); + for (String entry : entries) { + recordEvent.delete(entry); + } + } catch (final Exception e) { + LOG.error(EVENT, "There was an exception while processing Event [{}]", recordEvent, e); } } diff --git a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/ListToMapProcessor.java b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/ListToMapProcessor.java index de6757bace..d042f8fa28 100644 --- a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/ListToMapProcessor.java +++ b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/ListToMapProcessor.java @@ -45,38 +45,44 @@ public Collection> doExecute(final Collection> recor for (final Record record : records) { final Event recordEvent = record.getData(); - if (Objects.nonNull(config.getListToMapWhen()) && !expressionEvaluator.evaluateConditional(config.getListToMapWhen(), recordEvent)) { - continue; - } - - final List> sourceList; try { - sourceList = recordEvent.get(config.getSource(), List.class); - } catch (final Exception e) { - LOG.warn(EVENT, "Given source path [{}] is not valid on record [{}]", - config.getSource(), recordEvent, e); - recordEvent.getMetadata().addTags(config.getTagsOnFailure()); - continue; - } - final Map targetMap; - try { - targetMap = constructTargetMap(sourceList); - } catch (final IllegalArgumentException e) { - LOG.warn(EVENT, "Cannot find a list at the given source path [{}} on record [{}]", - config.getSource(), recordEvent, e); - recordEvent.getMetadata().addTags(config.getTagsOnFailure()); - continue; - } catch (final Exception e) { - LOG.error(EVENT, "Error converting source list to map on record [{}]", recordEvent, e); - recordEvent.getMetadata().addTags(config.getTagsOnFailure()); - continue; - } + if (Objects.nonNull(config.getListToMapWhen()) && !expressionEvaluator.evaluateConditional(config.getListToMapWhen(), recordEvent)) { + continue; + } - try { - updateEvent(recordEvent, targetMap); + final List> sourceList; + try { + sourceList = recordEvent.get(config.getSource(), List.class); + } catch (final Exception e) { + LOG.warn(EVENT, "Given source path [{}] is not valid on record [{}]", + config.getSource(), recordEvent, e); + recordEvent.getMetadata().addTags(config.getTagsOnFailure()); + continue; + } + + final Map targetMap; + try { + targetMap = constructTargetMap(sourceList); + } catch (final IllegalArgumentException e) { + LOG.warn(EVENT, "Cannot find a list at the given source path [{}} on record [{}]", + config.getSource(), recordEvent, e); + recordEvent.getMetadata().addTags(config.getTagsOnFailure()); + continue; + } catch (final Exception e) { + LOG.error(EVENT, "Error converting source list to map on record [{}]", recordEvent, e); + recordEvent.getMetadata().addTags(config.getTagsOnFailure()); + continue; + } + + try { + updateEvent(recordEvent, targetMap); + } catch (final Exception e) { + LOG.error(EVENT, "Error updating record [{}] after converting source list to map", recordEvent, e); + recordEvent.getMetadata().addTags(config.getTagsOnFailure()); + } } catch (final Exception e) { - LOG.error(EVENT, "Error updating record [{}] after converting source list to map", recordEvent, e); + LOG.error(EVENT, "There was an exception while processing Event [{}]", recordEvent, e); recordEvent.getMetadata().addTags(config.getTagsOnFailure()); } } diff --git a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/MapToListProcessor.java b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/MapToListProcessor.java index 49e74680fa..d911cd6194 100644 --- a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/MapToListProcessor.java +++ b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/MapToListProcessor.java @@ -5,6 +5,8 @@ package org.opensearch.dataprepper.plugins.processor.mutateevent; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; import org.opensearch.dataprepper.expression.ExpressionEvaluator; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; @@ -22,11 +24,15 @@ import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Set; +import static org.opensearch.dataprepper.logging.DataPrepperMarkers.EVENT; + @DataPrepperPlugin(name = "map_to_list", pluginType = Processor.class, pluginConfigurationType = MapToListProcessorConfig.class) public class MapToListProcessor extends AbstractProcessor, Record> { private static final Logger LOG = LoggerFactory.getLogger(MapToListProcessor.class); + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); private final MapToListProcessorConfig config; private final ExpressionEvaluator expressionEvaluator; private final Set excludeKeySet = new HashSet<>(); @@ -44,39 +50,79 @@ public Collection> doExecute(final Collection> recor for (final Record record : records) { final Event recordEvent = record.getData(); - if (config.getMapToListWhen() != null && !expressionEvaluator.evaluateConditional(config.getMapToListWhen(), recordEvent)) { - continue; - } - try { - final Map sourceMap = recordEvent.get(config.getSource(), Map.class); - final List> targetList = new ArrayList<>(); - - Map modifiedSourceMap = new HashMap<>(); - for (final Map.Entry entry : sourceMap.entrySet()) { - if (excludeKeySet.contains(entry.getKey())) { - if (config.getRemoveProcessedFields()) { - modifiedSourceMap.put(entry.getKey(), entry.getValue()); + + if (config.getMapToListWhen() != null && !expressionEvaluator.evaluateConditional(config.getMapToListWhen(), recordEvent)) { + continue; + } + + try { + final Map sourceMap = getSourceMap(recordEvent); + + if (config.getConvertFieldToList()) { + final List> targetNestedList = new ArrayList<>(); + + for (final Map.Entry entry : sourceMap.entrySet()) { + if (!excludeKeySet.contains(entry.getKey())) { + targetNestedList.add(List.of(entry.getKey(), entry.getValue())); + } + } - continue; + removeProcessedFields(sourceMap, recordEvent); + recordEvent.put(config.getTarget(), targetNestedList); + } else { + final List> targetList = new ArrayList<>(); + for (final Map.Entry entry : sourceMap.entrySet()) { + if (!excludeKeySet.contains(entry.getKey())) { + targetList.add(Map.of( + config.getKeyName(), entry.getKey(), + config.getValueName(), entry.getValue() + )); + } + } + removeProcessedFields(sourceMap, recordEvent); + recordEvent.put(config.getTarget(), targetList); } - targetList.add(Map.of( - config.getKeyName(), entry.getKey(), - config.getValueName(), entry.getValue() - )); + } catch (Exception e) { + LOG.error("Fail to perform Map to List operation", e); + recordEvent.getMetadata().addTags(config.getTagsOnFailure()); } + } catch (final Exception e) { + LOG.error(EVENT, "There was an exception while processing Event [{}]", recordEvent, e); + recordEvent.getMetadata().addTags(config.getTagsOnFailure()); + } + } + return records; + } - if (config.getRemoveProcessedFields()) { - recordEvent.put(config.getSource(), modifiedSourceMap); - } + private Map getSourceMap(Event recordEvent) throws JsonProcessingException { + final Map sourceMap; + sourceMap = recordEvent.get(config.getSource(), Map.class); + return sourceMap; + } - recordEvent.put(config.getTarget(), targetList); - } catch (Exception e) { - LOG.error("Fail to perform Map to List operation", e); - //TODO: add tagging on failure + private void removeProcessedFields(Map sourceMap, Event recordEvent) { + if (!config.getRemoveProcessedFields()) { + return; + } + + if (Objects.equals(config.getSource(), "")) { + // Source is root + for (final Map.Entry entry : sourceMap.entrySet()) { + if (excludeKeySet.contains(entry.getKey())) { + continue; + } + recordEvent.delete(entry.getKey()); } + } else { + Map modifiedSourceMap = new HashMap<>(); + for (final Map.Entry entry : sourceMap.entrySet()) { + if (excludeKeySet.contains(entry.getKey())) { + modifiedSourceMap.put(entry.getKey(), entry.getValue()); + } + } + recordEvent.put(config.getSource(), modifiedSourceMap); } - return records; } @Override diff --git a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/MapToListProcessorConfig.java b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/MapToListProcessorConfig.java index 3d863ea784..46a2ec79f0 100644 --- a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/MapToListProcessorConfig.java +++ b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/MapToListProcessorConfig.java @@ -18,7 +18,6 @@ public class MapToListProcessorConfig { private static final List DEFAULT_EXCLUDE_KEYS = new ArrayList<>(); private static final boolean DEFAULT_REMOVE_PROCESSED_FIELDS = false; - @NotEmpty @NotNull @JsonProperty("source") private String source; @@ -43,6 +42,12 @@ public class MapToListProcessorConfig { @JsonProperty("remove_processed_fields") private boolean removeProcessedFields = DEFAULT_REMOVE_PROCESSED_FIELDS; + @JsonProperty("convert_field_to_list") + private boolean convertFieldToList = false; + + @JsonProperty("tags_on_failure") + private List tagsOnFailure; + public String getSource() { return source; } @@ -70,4 +75,12 @@ public List getExcludeKeys() { public boolean getRemoveProcessedFields() { return removeProcessedFields; } + + public boolean getConvertFieldToList() { + return convertFieldToList; + } + + public List getTagsOnFailure() { + return tagsOnFailure; + } } diff --git a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/RenameKeyProcessor.java b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/RenameKeyProcessor.java index eff59f7751..05c1ad7530 100644 --- a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/RenameKeyProcessor.java +++ b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/RenameKeyProcessor.java @@ -13,13 +13,19 @@ import org.opensearch.dataprepper.model.processor.AbstractProcessor; import org.opensearch.dataprepper.model.processor.Processor; import org.opensearch.dataprepper.model.record.Record; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.util.Collection; import java.util.List; import java.util.Objects; +import static org.opensearch.dataprepper.logging.DataPrepperMarkers.EVENT; + @DataPrepperPlugin(name = "rename_keys", pluginType = Processor.class, pluginConfigurationType = RenameKeyProcessorConfig.class) public class RenameKeyProcessor extends AbstractProcessor, Record> { + + private static final Logger LOG = LoggerFactory.getLogger(RenameKeyProcessor.class); private final List entries; private final ExpressionEvaluator expressionEvaluator; @@ -36,20 +42,25 @@ public Collection> doExecute(final Collection> recor for(final Record record : records) { final Event recordEvent = record.getData(); - for(RenameKeyProcessorConfig.Entry entry : entries) { - if (Objects.nonNull(entry.getRenameWhen()) && !expressionEvaluator.evaluateConditional(entry.getRenameWhen(), recordEvent)) { - continue; - } + try { - if(entry.getFromKey().equals(entry.getToKey()) || !recordEvent.containsKey(entry.getFromKey())) { - continue; - } + for (RenameKeyProcessorConfig.Entry entry : entries) { + if (Objects.nonNull(entry.getRenameWhen()) && !expressionEvaluator.evaluateConditional(entry.getRenameWhen(), recordEvent)) { + continue; + } + + if (entry.getFromKey().equals(entry.getToKey()) || !recordEvent.containsKey(entry.getFromKey())) { + continue; + } - if (!recordEvent.containsKey(entry.getToKey()) || entry.getOverwriteIfToKeyExists()) { - final Object source = recordEvent.get(entry.getFromKey(), Object.class); - recordEvent.put(entry.getToKey(), source); - recordEvent.delete(entry.getFromKey()); + if (!recordEvent.containsKey(entry.getToKey()) || entry.getOverwriteIfToKeyExists()) { + final Object source = recordEvent.get(entry.getFromKey(), Object.class); + recordEvent.put(entry.getToKey(), source); + recordEvent.delete(entry.getFromKey()); + } } + } catch (final Exception e) { + LOG.error(EVENT, "There was an exception while processing Event [{}]", recordEvent, e); } } diff --git a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/SelectEntriesProcessor.java b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/SelectEntriesProcessor.java new file mode 100644 index 0000000000..af8de2fe7a --- /dev/null +++ b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/SelectEntriesProcessor.java @@ -0,0 +1,86 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.mutateevent; + +import org.opensearch.dataprepper.expression.ExpressionEvaluator; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; +import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; +import org.opensearch.dataprepper.model.plugin.InvalidPluginConfigurationException; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.processor.AbstractProcessor; +import org.opensearch.dataprepper.model.processor.Processor; +import org.opensearch.dataprepper.model.record.Record; + +import java.util.Collection; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +@DataPrepperPlugin(name = "select_entries", pluginType = Processor.class, pluginConfigurationType = SelectEntriesProcessorConfig.class) +public class SelectEntriesProcessor extends AbstractProcessor, Record> { + private final List keysToInclude; + private final String selectWhen; + + private final ExpressionEvaluator expressionEvaluator; + + @DataPrepperPluginConstructor + public SelectEntriesProcessor(final PluginMetrics pluginMetrics, final SelectEntriesProcessorConfig config, final ExpressionEvaluator expressionEvaluator) { + super(pluginMetrics); + this.selectWhen = config.getSelectWhen(); + if (selectWhen != null + && !expressionEvaluator.isValidExpressionStatement(selectWhen)) { + throw new InvalidPluginConfigurationException( + String.format("select_when value of %s is not a valid expression statement. " + + "See https://opensearch.org/docs/latest/data-prepper/pipelines/expression-syntax/ for valid expression syntax.", selectWhen)); + } + this.keysToInclude = config.getIncludeKeys(); + this.expressionEvaluator = expressionEvaluator; + } + + @Override + public Collection> doExecute(final Collection> records) { + for(final Record record : records) { + final Event recordEvent = record.getData(); + + if (Objects.nonNull(selectWhen) && !expressionEvaluator.evaluateConditional(selectWhen, recordEvent)) { + continue; + } + // To handle nested case, just get the values and store + // in a temporary map. + Map outMap = new HashMap<>(); + for (String keyToInclude: keysToInclude) { + Object value = recordEvent.get(keyToInclude, Object.class); + if (value != null) { + outMap.put(keyToInclude, value); + } + } + recordEvent.clear(); + + // add back only the keys selected + for (Map.Entry entry: outMap.entrySet()) { + recordEvent.put(entry.getKey(), entry.getValue()); + } + } + + return records; + } + + @Override + public void prepareForShutdown() { + } + + @Override + public boolean isReadyForShutdown() { + return true; + } + + @Override + public void shutdown() { + } +} + diff --git a/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/SelectEntriesProcessorConfig.java b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/SelectEntriesProcessorConfig.java new file mode 100644 index 0000000000..e19723f20d --- /dev/null +++ b/data-prepper-plugins/mutate-event-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutateevent/SelectEntriesProcessorConfig.java @@ -0,0 +1,31 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.mutateevent; + +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.constraints.NotEmpty; +import jakarta.validation.constraints.NotNull; + +import java.util.List; + +public class SelectEntriesProcessorConfig { + @NotEmpty + @NotNull + @JsonProperty("include_keys") + private List includeKeys; + + @JsonProperty("select_when") + private String selectWhen; + + public List getIncludeKeys() { + return includeKeys; + } + + public String getSelectWhen() { + return selectWhen; + } +} + diff --git a/data-prepper-plugins/mutate-event-processors/src/test/java/org/opensearch/dataprepper/plugins/processor/mutateevent/AddEntryProcessorTests.java b/data-prepper-plugins/mutate-event-processors/src/test/java/org/opensearch/dataprepper/plugins/processor/mutateevent/AddEntryProcessorTests.java index b3363aca53..36dc7ac5d4 100644 --- a/data-prepper-plugins/mutate-event-processors/src/test/java/org/opensearch/dataprepper/plugins/processor/mutateevent/AddEntryProcessorTests.java +++ b/data-prepper-plugins/mutate-event-processors/src/test/java/org/opensearch/dataprepper/plugins/processor/mutateevent/AddEntryProcessorTests.java @@ -15,14 +15,15 @@ import org.opensearch.dataprepper.model.event.JacksonEvent; import org.opensearch.dataprepper.model.record.Record; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.LinkedList; import java.util.List; import java.util.Map; -import java.util.UUID; import java.util.Random; +import java.util.UUID; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; @@ -47,7 +48,7 @@ public class AddEntryProcessorTests { @Test public void testSingleAddProcessorTests() { - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, 3, null, null, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, 3, null, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEvent("thisisamessage"); @@ -61,8 +62,8 @@ public void testSingleAddProcessorTests() { @Test public void testMultiAddProcessorTests() { - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, 3, null, null, false, null), - createEntry("message2", null, 4, null, null, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, 3, null, null, false, false,null), + createEntry("message2", null, 4, null, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEvent("thisisamessage"); @@ -78,7 +79,7 @@ public void testMultiAddProcessorTests() { @Test public void testSingleNoOverwriteAddProcessorTests() { - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, 3, null, null, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, 3, null, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEvent("thisisamessage"); @@ -93,7 +94,7 @@ public void testSingleNoOverwriteAddProcessorTests() { @Test public void testSingleOverwriteAddProcessorTests() { - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, 3, null, null, true, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, 3, null, null, true, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEvent("thisisamessage"); @@ -108,8 +109,8 @@ public void testSingleOverwriteAddProcessorTests() { @Test public void testMultiOverwriteMixedAddProcessorTests() { - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, 3, null, null, true, null), - createEntry("message", null, 4, null, null, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, 3, null, null, true, false,null), + createEntry("message", null, 4, null, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEvent("thisisamessage"); @@ -122,9 +123,37 @@ public void testMultiOverwriteMixedAddProcessorTests() { assertThat(editedRecords.get(0).getData().get("message", Object.class), equalTo("thisisamessage")); } + @Test + public void testAppendValueToExistingSimpleField() { + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("message", null, 3, null, null, false, true,null))); + + final AddEntryProcessor processor = createObjectUnderTest(); + final String currentValue = "old_message"; + final Record record = getEvent(currentValue); + final List> editedRecords = (List>) processor.doExecute(Collections.singletonList(record)); + + assertThat(editedRecords.get(0).getData().containsKey("message"), is(true)); + assertThat(editedRecords.get(0).getData().get("message", Object.class), equalTo(List.of(currentValue, 3))); + } + + @Test + public void testAppendValueToExistingListField() { + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("message", null, 3, null, null, false, true,null))); + + final AddEntryProcessor processor = createObjectUnderTest(); + final List listValue = new ArrayList<>(); + final String currentItem = "old_message"; + listValue.add(currentItem); + final Record record = getEvent(listValue); + final List> editedRecords = (List>) processor.doExecute(Collections.singletonList(record)); + + assertThat(editedRecords.get(0).getData().containsKey("message"), is(true)); + assertThat(editedRecords.get(0).getData().get("message", Object.class), equalTo(List.of(currentItem, 3))); + } + @Test public void testIntAddProcessorTests() { - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, 3, null, null, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, 3, null, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEvent("thisisamessage"); @@ -138,7 +167,7 @@ public void testIntAddProcessorTests() { @Test public void testBoolAddProcessorTests() { - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, true, null, null, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, true, null, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEvent("thisisamessage"); @@ -152,7 +181,7 @@ public void testBoolAddProcessorTests() { @Test public void testStringAddProcessorTests() { - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, "string", null, null, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, "string", null, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEvent("thisisamessage"); @@ -166,7 +195,7 @@ public void testStringAddProcessorTests() { @Test public void testNullAddProcessorTests() { - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, null, null, null, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, null, null, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEvent("thisisamessage"); @@ -192,7 +221,7 @@ public boolean equals(Object o) { public void testNestedAddProcessorTests() { TestObject obj = new TestObject(); obj.a = "test"; - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, obj, null, null, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, obj, null, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEvent("thisisamessage"); @@ -207,7 +236,7 @@ public void testNestedAddProcessorTests() { @Test public void testArrayAddProcessorTests() { Object[] array = new Object[] { 1, 1.2, "string", true, null }; - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, array, null, null, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, array, null, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEvent("thisisamessage"); @@ -222,7 +251,7 @@ public void testArrayAddProcessorTests() { @Test public void testFloatAddProcessorTests() { when(mockConfig.getEntries()) - .thenReturn(createListOfEntries(createEntry("newMessage", null, 1.2, null, null, false, null))); + .thenReturn(createListOfEntries(createEntry("newMessage", null, 1.2, null, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEvent("thisisamessage"); @@ -237,7 +266,7 @@ public void testFloatAddProcessorTests() { @Test public void testAddSingleFormatEntry() { when(mockConfig.getEntries()) - .thenReturn(createListOfEntries(createEntry("date-time", null, null, TEST_FORMAT, null, false, null))); + .thenReturn(createListOfEntries(createEntry("date-time", null, null, TEST_FORMAT, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getTestEventWithMultipleFields(); @@ -252,8 +281,8 @@ public void testAddSingleFormatEntry() { @Test public void testAddMultipleFormatEntries() { when(mockConfig.getEntries()) - .thenReturn(createListOfEntries(createEntry("date-time", null, null, TEST_FORMAT, null, false, null), - createEntry("date-time2", null, null, ANOTHER_TEST_FORMAT, null, false, null))); + .thenReturn(createListOfEntries(createEntry("date-time", null, null, TEST_FORMAT, null, false, false,null), + createEntry("date-time2", null, null, ANOTHER_TEST_FORMAT, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getTestEventWithMultipleFields(); @@ -270,7 +299,7 @@ public void testAddMultipleFormatEntries() { public void testFormatOverwritesExistingEntry() { when(mockConfig.getEntries()) .thenReturn( - createListOfEntries(createEntry("time", null, null, TEST_FORMAT, null, true, null))); + createListOfEntries(createEntry("time", null, null, TEST_FORMAT, null, true, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getTestEventWithMultipleFields(); @@ -285,7 +314,7 @@ public void testFormatOverwritesExistingEntry() { public void testFormatNotOverwriteExistingEntry() { when(mockConfig.getEntries()) .thenReturn( - createListOfEntries(createEntry("time", null, null, TEST_FORMAT, null, false, null))); + createListOfEntries(createEntry("time", null, null, TEST_FORMAT, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getTestEventWithMultipleFields(); @@ -297,11 +326,40 @@ public void testFormatNotOverwriteExistingEntry() { assertThat(event.containsKey("date-time"), equalTo(false)); } + @Test + public void testAppendFormatValueToExistingSimpleField() { + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("time", null, 3, TEST_FORMAT, null, false, true,null))); + + final AddEntryProcessor processor = createObjectUnderTest(); + final Record record = getTestEventWithMultipleFields(); + final List> editedRecords = (List>) processor.doExecute(Collections.singletonList(record)); + + Event event = editedRecords.get(0).getData(); + assertThat(event.get("date", Object.class), equalTo("date-value")); + assertThat(event.get("time", Object.class), equalTo(List.of("time-value", "date-value time-value"))); + } + + @Test + public void testAppendFormatValueToExistingListField() { + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("date-time", null, 3, TEST_FORMAT, null, false, true,null))); + + final AddEntryProcessor processor = createObjectUnderTest(); + final Record record = getTestEventWithMultipleFields(); + final List listValue = new ArrayList<>(); + final String currentItem = "date-time-value-1"; + listValue.add(currentItem); + record.getData().put("date-time", listValue); + final List> editedRecords = (List>) processor.doExecute(Collections.singletonList(record)); + + Event event = editedRecords.get(0).getData(); + assertThat(event.get("date-time", Object.class), equalTo(List.of(currentItem, "date-value time-value"))); + } + @Test public void testFormatPrecedesValue() { when(mockConfig.getEntries()) .thenReturn( - createListOfEntries(createEntry("date-time", null, "date-time-value", TEST_FORMAT, null, false, null))); + createListOfEntries(createEntry("date-time", null, "date-time-value", TEST_FORMAT, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getTestEventWithMultipleFields(); @@ -316,7 +374,7 @@ public void testFormatPrecedesValue() { @Test public void testFormatVariousDataTypes() { when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry( - "newField", null, null, "${number-key}-${boolean-key}-${string-key}", null, false, null))); + "newField", null, null, "${number-key}-${boolean-key}-${string-key}", null, false, false, null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getTestEventWithMultipleDataTypes(); @@ -328,7 +386,7 @@ public void testFormatVariousDataTypes() { @Test public void testBadFormatThenEntryNotAdded() { - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("data-time", null, null, BAD_TEST_FORMAT, null, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("data-time", null, null, BAD_TEST_FORMAT, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getTestEventWithMultipleFields(); @@ -342,7 +400,7 @@ public void testBadFormatThenEntryNotAdded() { @Test public void testMetadataKeySetWithBadFormatThenEntryNotAdded() { - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry(null,"data-time", null, BAD_TEST_FORMAT, null, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry(null,"data-time", null, BAD_TEST_FORMAT, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEventWithMetadata("message", Map.of("date", "date-value", "time", "time-value")); @@ -358,7 +416,7 @@ public void testMetadataKeySetWithBadFormatThenEntryNotAdded() { public void testKeyIsNotAdded_when_addWhen_condition_is_false() { final String addWhen = UUID.randomUUID().toString(); - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, 3, null, null, false, addWhen))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("newMessage", null, 3, null, null, false, false,addWhen))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEvent("thisisamessage"); @@ -375,7 +433,7 @@ public void testKeyIsNotAdded_when_addWhen_condition_is_false() { public void testMetadataKeyIsNotAdded_when_addWhen_condition_is_false() { final String addWhen = UUID.randomUUID().toString(); - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry(null, "newMessage", 3, null, null, false, addWhen))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry(null, "newMessage", 3, null, null, false, false,addWhen))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEventWithMetadata("thisisamessage", Map.of("key", "value")); @@ -393,9 +451,9 @@ public void testMetadataKeyIsNotAdded_when_addWhen_condition_is_false() { @Test public void testMetadataKeySetWithDifferentDataTypes() { when(mockConfig.getEntries()).thenReturn(createListOfEntries( - createEntry(null, "newField", "newValue", null, null, false, null), - createEntry(null, "newIntField", 123, null, null, false, null), - createEntry(null, "newBooleanField", true, null, null, false, null) + createEntry(null, "newField", "newValue", null, null, false, false,null), + createEntry(null, "newIntField", 123, null, null, false, false,null), + createEntry(null, "newBooleanField", true, null, null, false, false,null) )); final AddEntryProcessor processor = createObjectUnderTest(); @@ -412,7 +470,7 @@ public void testMetadataKeySetWithDifferentDataTypes() { public void testMetadataKeySetWithFormatNotOverwriteExistingEntry() { when(mockConfig.getEntries()) .thenReturn( - createListOfEntries(createEntry(null, "time", null, TEST_FORMAT, null, false, null))); + createListOfEntries(createEntry(null, "time", null, TEST_FORMAT, null, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEventWithMetadata("message", Map.of("date", "date-value", "time", "time-value")); @@ -428,7 +486,7 @@ public void testMetadataKeySetWithFormatNotOverwriteExistingEntry() { public void testMetadataKeySetWithFormatOverwriteExistingEntry() { when(mockConfig.getEntries()) .thenReturn( - createListOfEntries(createEntry(null, "time", null, TEST_FORMAT, null, true, null))); + createListOfEntries(createEntry(null, "time", null, TEST_FORMAT, null, true, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEventWithMetadata("message", Map.of("date", "date-value", "time", "time-value")); @@ -440,40 +498,72 @@ public void testMetadataKeySetWithFormatOverwriteExistingEntry() { assertThat(attributes.containsKey("date-time"), equalTo(false)); } + @Test + public void testMetadataKeySetAppendToExistingSimpleValue() { + when(mockConfig.getEntries()) + .thenReturn( + createListOfEntries(createEntry(null, "time", "time-value2", null, null, false, true,null))); + + final AddEntryProcessor processor = createObjectUnderTest(); + final String currentValue = "time-value1"; + final Record record = getEventWithMetadata("message", Map.of("time", currentValue)); + final List> editedRecords = (List>) processor.doExecute(Collections.singletonList(record)); + + Map attributes = editedRecords.get(0).getData().getMetadata().getAttributes(); + assertThat(attributes.get("time"), equalTo(List.of(currentValue, "time-value2"))); + } + + @Test + public void testMetadataKeySetAppendToExistingListValue() { + when(mockConfig.getEntries()) + .thenReturn( + createListOfEntries(createEntry(null, "time", "time-value2", null, null, false, true,null))); + + final AddEntryProcessor processor = createObjectUnderTest(); + final List listValue = new ArrayList<>(); + final String currentItem = "time-value1"; + listValue.add(currentItem); + final Record record = getEventWithMetadata("message", Map.of("time", listValue)); + final List> editedRecords = (List>) processor.doExecute(Collections.singletonList(record)); + + Map attributes = editedRecords.get(0).getData().getMetadata().getAttributes(); + assertThat(attributes.get("time"), equalTo(List.of(currentItem, "time-value2"))); + } + @Test public void testMetadataKeyAndKeyBothNotSetThrows() { - assertThrows(IllegalArgumentException.class, () -> createEntry(null, null, "newValue", null, null, false, null)); + assertThrows(IllegalArgumentException.class, () -> createEntry(null, null, "newValue", null, null, false, false,null)); } @Test public void testMetadataKeyAndKeyBothSetThrows() { - assertThrows(IllegalArgumentException.class, () -> createEntry("newKey", "newMetadataKey", "newValue", null, null, false, null)); + assertThrows(IllegalArgumentException.class, () -> createEntry("newKey", "newMetadataKey", "newValue", null, null, false, false,null)); } @Test public void testOnlyOneTypeOfValueIsSupported() { - assertThrows(RuntimeException.class, () -> createEntry("newKey", "newMetadataKey", "newValue", "/newFormat", null, false, null)); + assertThrows(RuntimeException.class, () -> createEntry("newKey", "newMetadataKey", "newValue", "/newFormat", null, false, false,null)); } @Test public void testOnlyOneTypeOfValueIsSupportedWithExpressionAndFormat() { - assertThrows(RuntimeException.class, () -> createEntry("newKey", "newMetadataKey", null, "/newFormat", "length(/message)", false, null)); + assertThrows(RuntimeException.class, () -> createEntry("newKey", "newMetadataKey", null, "/newFormat", "length(/message)", false, false,null)); } @Test public void testOnlyOneTypeOfValueIsSupportedWithValueAndExpressionAndFormat() { - assertThrows(RuntimeException.class, () -> createEntry("newKey", "newMetadataKey", "value", "/newFormat", "length(/message)", false, null)); + assertThrows(RuntimeException.class, () -> createEntry("newKey", "newMetadataKey", "value", "/newFormat", "length(/message)", false, false,null)); } @Test public void testWithAllValuesNull() { - assertThrows(RuntimeException.class, () -> createEntry("newKey", "newMetadataKey", null, null, null, false, null)); + assertThrows(RuntimeException.class, () -> createEntry("newKey", "newMetadataKey", null, null, null, false, false,null)); } @Test public void testValueExpressionWithArithmeticExpression() { String valueExpression = "/number-key"; - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("num_key", null, null, null, valueExpression, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("num_key", null, null, null, valueExpression, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getTestEventWithMultipleDataTypes(); Random random = new Random(); @@ -487,7 +577,7 @@ public void testValueExpressionWithArithmeticExpression() { @Test public void testValueExpressionWithStringExpression() { String valueExpression = "/string-key"; - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("num_key", null, null, null, valueExpression, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("num_key", null, null, null, valueExpression, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getTestEventWithMultipleDataTypes(); String randomString = UUID.randomUUID().toString(); @@ -500,7 +590,7 @@ public void testValueExpressionWithStringExpression() { @Test public void testValueExpressionWithBooleanExpression() { String valueExpression = "/number-key > 5"; - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("num_key", null, null, null, valueExpression, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("num_key", null, null, null, valueExpression, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getTestEventWithMultipleDataTypes(); when(expressionEvaluator.evaluate(valueExpression, record.getData())).thenReturn(false); @@ -512,7 +602,7 @@ public void testValueExpressionWithBooleanExpression() { @Test public void testValueExpressionWithIntegerFunctions() { String valueExpression = "length(/string-key)"; - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("length_key", null, null, null, valueExpression, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("length_key", null, null, null, valueExpression, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getTestEventWithMultipleDataTypes(); String randomString = UUID.randomUUID().toString(); @@ -525,7 +615,7 @@ public void testValueExpressionWithIntegerFunctions() { @Test public void testValueExpressionWithIntegerFunctionsAndMetadataKey() { String valueExpression = "length(/date)"; - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry(null, "length_key", null, null, valueExpression, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry(null, "length_key", null, null, valueExpression, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEventWithMetadata("message", Map.of("key", "value")); String randomString = UUID.randomUUID().toString(); @@ -538,7 +628,7 @@ public void testValueExpressionWithIntegerFunctionsAndMetadataKey() { @Test public void testValueExpressionWithStringExpressionWithMetadataKey() { String valueExpression = "/date"; - when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry(null, "newkey", null, null, valueExpression, false, null))); + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry(null, "newkey", null, null, valueExpression, false, false,null))); final AddEntryProcessor processor = createObjectUnderTest(); final Record record = getEventWithMetadata("message", Map.of("key", "value")); String randomString = UUID.randomUUID().toString(); @@ -548,20 +638,94 @@ public void testValueExpressionWithStringExpressionWithMetadataKey() { assertThat(attributes.get("newkey"), equalTo(randomString)); } + @Test + public void testAddSingleFieldWithDynamicKey() { + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("${message}", null, 3, null, null, false, false,null))); + + final AddEntryProcessor processor = createObjectUnderTest(); + final Record record = getEvent("value_as_name"); + final List> editedRecords = (List>) processor.doExecute(Collections.singletonList(record)); + + assertThat(editedRecords.get(0).getData().containsKey("message"), is(true)); + assertThat(editedRecords.get(0).getData().get("message", Object.class), equalTo("value_as_name")); + assertThat(editedRecords.get(0).getData().containsKey("value_as_name"), is(true)); + assertThat(editedRecords.get(0).getData().get("value_as_name", Object.class), equalTo(3)); + } + + @Test + public void testAddSingleFieldWithDynamicExpression() { + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("${message}_${getMetadata(\"id\")}", null, 3, null, null, false, false,null))); + + final AddEntryProcessor processor = createObjectUnderTest(); + final Record record = getEventWithMetadata("value_as_name", Map.of("id", 1)); + when(expressionEvaluator.isValidExpressionStatement("getMetadata(\"id\")")).thenReturn(true); + when(expressionEvaluator.evaluate("getMetadata(\"id\")", record.getData())).thenReturn(1); + final List> editedRecords = (List>) processor.doExecute(Collections.singletonList(record)); + + assertThat(editedRecords.get(0).getData().containsKey("message"), is(true)); + assertThat(editedRecords.get(0).getData().get("message", Object.class), equalTo("value_as_name")); + assertThat(editedRecords.get(0).getData().containsKey("value_as_name_1"), is(true)); + assertThat(editedRecords.get(0).getData().get("value_as_name_1", Object.class), equalTo(3)); + } + + @Test + public void testAddMultipleFieldsWithDynamicKeys() { + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("${message}", null, 3, null, null, false, false,null), + createEntry("${message}_2", null, 4, null, null, false, false,null))); + + final AddEntryProcessor processor = createObjectUnderTest(); + final Record record = getEvent("value_as_name"); + final List> editedRecords = (List>) processor.doExecute(Collections.singletonList(record)); + + assertThat(editedRecords.get(0).getData().containsKey("message"), is(true)); + assertThat(editedRecords.get(0).getData().get("message", Object.class), equalTo("value_as_name")); + assertThat(editedRecords.get(0).getData().containsKey("value_as_name"), is(true)); + assertThat(editedRecords.get(0).getData().get("value_as_name", Object.class), equalTo(3)); + assertThat(editedRecords.get(0).getData().containsKey("value_as_name_2"), is(true)); + assertThat(editedRecords.get(0).getData().get("value_as_name_2", Object.class), equalTo(4)); + } + + @Test + public void testAddFieldWithInvalidInputKeyThenNoChangeToEvent() { + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("${message", null, 3, null, null, false, false,null))); + + final AddEntryProcessor processor = createObjectUnderTest(); + final Record record = getEvent("value_as_name"); + final List> editedRecords = (List>) processor.doExecute(Collections.singletonList(record)); + + assertThat(editedRecords.get(0).getData().containsKey("message"), is(true)); + assertThat(editedRecords.get(0).getData().get("message", Object.class), equalTo("value_as_name")); + assertThat(editedRecords.get(0).getData().containsKey("value_as_name"), is(false)); + assertThat(editedRecords.get(0).getData().toMap().size(), is(1)); + } + + @Test + public void testAddFieldWithInvalidDynamicKeyThenNoChangeToEvent() { + when(mockConfig.getEntries()).thenReturn(createListOfEntries(createEntry("${message}", null, 3, null, null, false, false,null))); + + final AddEntryProcessor processor = createObjectUnderTest(); + final Record record = getEvent("name_with_invalid_chars|[$"); + final List> editedRecords = (List>) processor.doExecute(Collections.singletonList(record)); + + assertThat(editedRecords.get(0).getData().containsKey("message"), is(true)); + assertThat(editedRecords.get(0).getData().get("message", Object.class), equalTo("name_with_invalid_chars|[$")); + assertThat(editedRecords.get(0).getData().toMap().size(), is(1)); + } + private AddEntryProcessor createObjectUnderTest() { return new AddEntryProcessor(pluginMetrics, mockConfig, expressionEvaluator); } private AddEntryProcessorConfig.Entry createEntry( - final String key, final String metadataKey, final Object value, final String format, final String valueExpression, final boolean overwriteIfKeyExists, final String addWhen) { - return new AddEntryProcessorConfig.Entry(key, metadataKey, value, format, valueExpression, overwriteIfKeyExists, addWhen); + final String key, final String metadataKey, final Object value, final String format, final String valueExpression, final boolean overwriteIfKeyExists, final boolean appendIfKeyExists, final String addWhen) { + return new AddEntryProcessorConfig.Entry(key, metadataKey, value, format, valueExpression, overwriteIfKeyExists, appendIfKeyExists, addWhen); } private List createListOfEntries(final AddEntryProcessorConfig.Entry... entries) { return new LinkedList<>(Arrays.asList(entries)); } - private Record getEvent(String message) { + private Record getEvent(Object message) { final Map testData = new HashMap<>(); testData.put("message", message); return buildRecordWithEvent(testData); diff --git a/data-prepper-plugins/mutate-event-processors/src/test/java/org/opensearch/dataprepper/plugins/processor/mutateevent/MapToListProcessorTest.java b/data-prepper-plugins/mutate-event-processors/src/test/java/org/opensearch/dataprepper/plugins/processor/mutateevent/MapToListProcessorTest.java index c748e676f4..83d736ba21 100644 --- a/data-prepper-plugins/mutate-event-processors/src/test/java/org/opensearch/dataprepper/plugins/processor/mutateevent/MapToListProcessorTest.java +++ b/data-prepper-plugins/mutate-event-processors/src/test/java/org/opensearch/dataprepper/plugins/processor/mutateevent/MapToListProcessorTest.java @@ -18,6 +18,7 @@ import java.util.ArrayList; import java.util.Collections; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.UUID; @@ -48,6 +49,8 @@ void setUp() { lenient().when(mockConfig.getMapToListWhen()).thenReturn(null); lenient().when(mockConfig.getExcludeKeys()).thenReturn(new ArrayList<>()); lenient().when(mockConfig.getRemoveProcessedFields()).thenReturn(false); + lenient().when(mockConfig.getConvertFieldToList()).thenReturn(false); + lenient().when(mockConfig.getTagsOnFailure()).thenReturn(new ArrayList<>()); } @Test @@ -72,6 +75,47 @@ void testMapToListSuccessWithDefaultOptions() { assertSourceMapUnchanged(resultEvent); } + @Test + void testMapToListSuccessWithNestedMap() { + + final MapToListProcessor processor = createObjectUnderTest(); + final Record testRecord = createTestRecordWithNestedMap(); + final List> resultRecord = (List>) processor.doExecute(Collections.singletonList(testRecord)); + + assertThat(resultRecord.size(), is(1)); + + final Event resultEvent = resultRecord.get(0).getData(); + List> resultList = resultEvent.get("my-list", List.class); + + assertThat(resultList.size(), is(2)); + assertThat(resultList, containsInAnyOrder( + Map.of("key", "key1", "value", "value1"), + Map.of("key", "key2", "value", Map.of("key2-1", "value2")) + )); + } + + @Test + void testMapToListSuccessWithRootAsSource() { + when(mockConfig.getSource()).thenReturn(""); + + final MapToListProcessor processor = createObjectUnderTest(); + final Record testRecord = createFlatTestRecord(); + final List> resultRecord = (List>) processor.doExecute(Collections.singletonList(testRecord)); + + assertThat(resultRecord.size(), is(1)); + + final Event resultEvent = resultRecord.get(0).getData(); + List> resultList = resultEvent.get("my-list", List.class); + + assertThat(resultList.size(), is(3)); + assertThat(resultList, containsInAnyOrder( + Map.of("key", "key1", "value", "value1"), + Map.of("key", "key2", "value", "value2"), + Map.of("key", "key3", "value", "value3") + )); + assertSourceMapUnchangedForFlatRecord(resultEvent); + } + @Test void testMapToListSuccessWithCustomKeyNameValueName() { final String keyName = "custom-key-name"; @@ -130,6 +174,25 @@ void testExcludedKeysAreNotProcessed() { assertSourceMapUnchanged(resultEvent); } + @Test + void testExcludedKeysAreNotProcessedWithRootAsSource() { + when(mockConfig.getSource()).thenReturn(""); + when(mockConfig.getExcludeKeys()).thenReturn(List.of("key1", "key3", "key5")); + + final MapToListProcessor processor = createObjectUnderTest(); + final Record testRecord = createFlatTestRecord(); + final List> resultRecord = (List>) processor.doExecute(Collections.singletonList(testRecord)); + + assertThat(resultRecord.size(), is(1)); + + final Event resultEvent = resultRecord.get(0).getData(); + List> resultList = resultEvent.get("my-list", List.class); + + assertThat(resultList.size(), is(1)); + assertThat(resultList.get(0), is(Map.of("key", "key2", "value", "value2"))); + assertSourceMapUnchangedForFlatRecord(resultEvent); + } + @Test void testRemoveProcessedFields() { when(mockConfig.getExcludeKeys()).thenReturn(List.of("key1", "key3", "key5")); @@ -155,6 +218,96 @@ void testRemoveProcessedFields() { assertThat(resultEvent.get("my-map/key3", String.class), is("value3")); } + @Test + void testRemoveProcessedFieldsWithRootAsSource() { + when(mockConfig.getSource()).thenReturn(""); + when(mockConfig.getExcludeKeys()).thenReturn(List.of("key1", "key3", "key5")); + when(mockConfig.getRemoveProcessedFields()).thenReturn(true); + + final MapToListProcessor processor = createObjectUnderTest(); + final Record testRecord = createFlatTestRecord(); + final List> resultRecord = (List>) processor.doExecute(Collections.singletonList(testRecord)); + + assertThat(resultRecord.size(), is(1)); + + final Event resultEvent = resultRecord.get(0).getData(); + List> resultList = resultEvent.get("my-list", List.class); + + assertThat(resultList.size(), is(1)); + assertThat(resultList.get(0), is(Map.of("key", "key2", "value", "value2"))); + + assertThat(resultEvent.containsKey("key1"), is(true)); + assertThat(resultEvent.get("key1", String.class), is("value1")); + assertThat(resultEvent.containsKey("key2"), is(false)); + assertThat(resultEvent.containsKey("key3"), is(true)); + assertThat(resultEvent.get("key3", String.class), is("value3")); + } + + @Test + public void testConvertFieldToListSuccess() { + when(mockConfig.getConvertFieldToList()).thenReturn(true); + + final MapToListProcessor processor = createObjectUnderTest(); + final Record testRecord = createTestRecord(); + final List> resultRecord = (List>) processor.doExecute(Collections.singletonList(testRecord)); + + assertThat(resultRecord.size(), is(1)); + + final Event resultEvent = resultRecord.get(0).getData(); + List> resultList = resultEvent.get("my-list", List.class); + + assertThat(resultList.size(), is(3)); + assertThat(resultList, containsInAnyOrder( + List.of("key1", "value1"), + List.of("key2", "value2"), + List.of("key3", "value3") + )); + assertSourceMapUnchanged(resultEvent); + } + + @Test + public void testConvertFieldToListSuccessWithNestedMap() { + when(mockConfig.getConvertFieldToList()).thenReturn(true); + + final MapToListProcessor processor = createObjectUnderTest(); + final Record testRecord = createTestRecordWithNestedMap(); + final List> resultRecord = (List>) processor.doExecute(Collections.singletonList(testRecord)); + + assertThat(resultRecord.size(), is(1)); + + final Event resultEvent = resultRecord.get(0).getData(); + List> resultList = resultEvent.get("my-list", List.class); + + assertThat(resultList.size(), is(2)); + assertThat(resultList, containsInAnyOrder( + List.of("key1", "value1"), + List.of("key2", Map.of("key2-1", "value2")) + )); + } + + @Test + public void testConvertFieldToListSuccessWithRootAsSource() { + when(mockConfig.getSource()).thenReturn(""); + when(mockConfig.getConvertFieldToList()).thenReturn(true); + + final MapToListProcessor processor = createObjectUnderTest(); + final Record testRecord = createFlatTestRecord(); + final List> resultRecord = (List>) processor.doExecute(Collections.singletonList(testRecord)); + + assertThat(resultRecord.size(), is(1)); + + final Event resultEvent = resultRecord.get(0).getData(); + List> resultList = resultEvent.get("my-list", List.class); + + assertThat(resultList.size(), is(3)); + assertThat(resultList, containsInAnyOrder( + List.of("key1", "value1"), + List.of("key2", "value2"), + List.of("key3", "value3") + )); + assertSourceMapUnchangedForFlatRecord(resultEvent); + } + @Test public void testEventNotProcessedWhenTheWhenConditionIsFalse() { final String whenCondition = UUID.randomUUID().toString(); @@ -172,6 +325,25 @@ public void testEventNotProcessedWhenTheWhenConditionIsFalse() { assertSourceMapUnchanged(resultEvent); } + @Test + void testFailureTagsAreAddedWhenException() { + // non-existing source key + when(mockConfig.getSource()).thenReturn("my-other-map"); + final List testTags = List.of("tag1", "tag2"); + when(mockConfig.getTagsOnFailure()).thenReturn(testTags); + + final MapToListProcessor processor = createObjectUnderTest(); + final Record testRecord = createTestRecord(); + final List> resultRecord = (List>) processor.doExecute(Collections.singletonList(testRecord)); + + assertThat(resultRecord.size(), is(1)); + + final Event resultEvent = resultRecord.get(0).getData(); + assertThat(resultEvent.containsKey("my-list"), is(false)); + assertSourceMapUnchanged(resultEvent); + assertThat(resultEvent.getMetadata().getTags(), is(new HashSet<>(testTags))); + } + private MapToListProcessor createObjectUnderTest() { return new MapToListProcessor(pluginMetrics, mockConfig, expressionEvaluator); } @@ -188,10 +360,39 @@ private Record createTestRecord() { return new Record<>(event); } + private Record createFlatTestRecord() { + final Map data = Map.of( + "key1", "value1", + "key2", "value2", + "key3", "value3"); + final Event event = JacksonEvent.builder() + .withData(data) + .withEventType("event") + .build(); + return new Record<>(event); + } + + private Record createTestRecordWithNestedMap() { + final Map> data = Map.of("my-map", Map.of( + "key1", "value1", + "key2", Map.of("key2-1", "value2"))); + final Event event = JacksonEvent.builder() + .withData(data) + .withEventType("event") + .build(); + return new Record<>(event); + } + private void assertSourceMapUnchanged(final Event resultEvent) { assertThat(resultEvent.containsKey("my-map"), is(true)); assertThat(resultEvent.get("my-map/key1", String.class), is("value1")); assertThat(resultEvent.get("my-map/key2", String.class), is("value2")); assertThat(resultEvent.get("my-map/key3", String.class), is("value3")); } + + private void assertSourceMapUnchangedForFlatRecord(final Event resultEvent) { + assertThat(resultEvent.get("key1", String.class), is("value1")); + assertThat(resultEvent.get("key2", String.class), is("value2")); + assertThat(resultEvent.get("key3", String.class), is("value3")); + } } diff --git a/data-prepper-plugins/mutate-event-processors/src/test/java/org/opensearch/dataprepper/plugins/processor/mutateevent/SelectEntriesProcessorTests.java b/data-prepper-plugins/mutate-event-processors/src/test/java/org/opensearch/dataprepper/plugins/processor/mutateevent/SelectEntriesProcessorTests.java new file mode 100644 index 0000000000..048f304582 --- /dev/null +++ b/data-prepper-plugins/mutate-event-processors/src/test/java/org/opensearch/dataprepper/plugins/processor/mutateevent/SelectEntriesProcessorTests.java @@ -0,0 +1,148 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.processor.mutateevent; + +import org.opensearch.dataprepper.expression.ExpressionEvaluator; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.event.JacksonEvent; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.plugin.InvalidPluginConfigurationException; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +public class SelectEntriesProcessorTests { + @Mock + private PluginMetrics pluginMetrics; + + @Mock + private SelectEntriesProcessorConfig mockConfig; + + @Mock + private ExpressionEvaluator expressionEvaluator; + + @Test + public void testSelectEntriesProcessor() { + when(mockConfig.getIncludeKeys()).thenReturn(List.of("key1", "key2")); + when(mockConfig.getSelectWhen()).thenReturn(null); + final SelectEntriesProcessor processor = createObjectUnderTest(); + final Record record = getEvent("thisisamessage"); + record.getData().put("newMessage", "test"); + final String value1 = UUID.randomUUID().toString(); + final String value2 = UUID.randomUUID().toString(); + record.getData().put("key1", value1); + record.getData().put("key2", value2); + final List> editedRecords = (List>) processor.doExecute(Collections.singletonList(record)); + assertThat(editedRecords.get(0).getData().containsKey("key1"), is(true)); + assertThat(editedRecords.get(0).getData().containsKey("key2"), is(true)); + assertThat(editedRecords.get(0).getData().containsKey("newMessage"), is(false)); + assertThat(editedRecords.get(0).getData().containsKey("message"), is(false)); + assertThat(editedRecords.get(0).getData().get("key1", String.class), equalTo(value1)); + assertThat(editedRecords.get(0).getData().get("key2", String.class), equalTo(value2)); + } + + @Test + public void testWithKeyDneSelectEntriesProcessor() { + when(mockConfig.getIncludeKeys()).thenReturn(List.of("key1", "key2")); + when(mockConfig.getSelectWhen()).thenReturn(null); + final SelectEntriesProcessor processor = createObjectUnderTest(); + final Record record = getEvent("thisisamessage"); + record.getData().put("newMessage", "test"); + final List> editedRecords = (List>) processor.doExecute(Collections.singletonList(record)); + assertThat(editedRecords.get(0).getData().containsKey("key1"), is(false)); + assertThat(editedRecords.get(0).getData().containsKey("key2"), is(false)); + assertThat(editedRecords.get(0).getData().containsKey("message"), is(false)); + assertThat(editedRecords.get(0).getData().containsKey("newMessage"), is(false)); + } + + @Test + public void testSelectEntriesProcessorWithInvalidCondition() { + final String selectWhen = "/message == \""+UUID.randomUUID().toString()+"\""; + when(expressionEvaluator.isValidExpressionStatement(selectWhen)).thenReturn(false); + when(mockConfig.getSelectWhen()).thenReturn(selectWhen); + assertThrows(InvalidPluginConfigurationException.class, () -> createObjectUnderTest()); + final Record record = getEvent("thisisamessage"); + } + + @Test + public void testSelectEntriesProcessorWithCondition() { + when(mockConfig.getIncludeKeys()).thenReturn(List.of("key1", "key2")); + final String selectWhen = "/message == \""+UUID.randomUUID().toString()+"\""; + when(expressionEvaluator.isValidExpressionStatement(selectWhen)).thenReturn(true); + when(mockConfig.getSelectWhen()).thenReturn(selectWhen); + final SelectEntriesProcessor processor = createObjectUnderTest(); + final Record record = getEvent("thisisamessage"); + record.getData().put("newMessage", "test"); + final String value1 = UUID.randomUUID().toString(); + final String value2 = UUID.randomUUID().toString(); + record.getData().put("key1", value1); + record.getData().put("key2", value2); + when(expressionEvaluator.evaluateConditional(selectWhen, record.getData())).thenReturn(false); + final List> editedRecords = (List>) processor.doExecute(Collections.singletonList(record)); + assertThat(editedRecords.get(0).getData().containsKey("key1"), is(true)); + assertThat(editedRecords.get(0).getData().containsKey("key2"), is(true)); + assertThat(editedRecords.get(0).getData().containsKey("newMessage"), is(true)); + assertThat(editedRecords.get(0).getData().containsKey("message"), is(true)); + assertThat(editedRecords.get(0).getData().get("key1", String.class), equalTo(value1)); + assertThat(editedRecords.get(0).getData().get("key2", String.class), equalTo(value2)); + } + + @Test + public void testNestedSelectEntriesProcessor() { + when(mockConfig.getIncludeKeys()).thenReturn(List.of("nested/key1", "nested/nested2/key2")); + when(mockConfig.getSelectWhen()).thenReturn(null); + final String value1 = UUID.randomUUID().toString(); + final String value2 = UUID.randomUUID().toString(); + Map nested2 = Map.of("key2", value2, "key3", "value3"); + Map nested = Map.of("key1", value1, "fizz", 42, "nested2", nested2); + final SelectEntriesProcessor processor = createObjectUnderTest(); + final Record record = getEvent("thisisamessage"); + record.getData().put("nested", nested); + final List> editedRecords = (List>) processor.doExecute(Collections.singletonList(record)); + assertThat(editedRecords.get(0).getData().containsKey("nested/key1"), is(true)); + assertThat(editedRecords.get(0).getData().containsKey("nested/nested2/key2"), is(true)); + assertThat(editedRecords.get(0).getData().containsKey("nested/nested2/key3"), is(false)); + assertThat(editedRecords.get(0).getData().containsKey("nested/fizz"), is(false)); + assertThat(editedRecords.get(0).getData().containsKey("newMessage"), is(false)); + assertThat(editedRecords.get(0).getData().containsKey("message"), is(false)); + + assertThat(editedRecords.get(0).getData().get("nested/key1", String.class), equalTo(value1)); + assertThat(editedRecords.get(0).getData().get("nested/nested2/key2", String.class), equalTo(value2)); + } + + + private SelectEntriesProcessor createObjectUnderTest() { + return new SelectEntriesProcessor(pluginMetrics, mockConfig, expressionEvaluator); + } + + private Record getEvent(String message) { + final Map testData = new HashMap(); + testData.put("message", message); + return buildRecordWithEvent(testData); + } + + private static Record buildRecordWithEvent(final Map data) { + return new Record<>(JacksonEvent.builder() + .withData(data) + .withEventType("event") + .build()); + } +} diff --git a/data-prepper-plugins/mutate-string-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutatestring/AbstractStringProcessor.java b/data-prepper-plugins/mutate-string-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutatestring/AbstractStringProcessor.java index bcd0f21aa8..19d11daf62 100644 --- a/data-prepper-plugins/mutate-string-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutatestring/AbstractStringProcessor.java +++ b/data-prepper-plugins/mutate-string-processors/src/main/java/org/opensearch/dataprepper/plugins/processor/mutatestring/AbstractStringProcessor.java @@ -33,7 +33,11 @@ public AbstractStringProcessor(final PluginMetrics pluginMetrics, final StringPr public Collection> doExecute(final Collection> records) { for(final Record record : records) { final Event recordEvent = record.getData(); - performStringAction(recordEvent); + try { + performStringAction(recordEvent); + } catch (final Exception e) { + LOG.error(EVENT, "There was an exception while processing Event [{}]", recordEvent, e); + } } return records; diff --git a/data-prepper-plugins/newline-codecs/build.gradle b/data-prepper-plugins/newline-codecs/build.gradle index b4bedd67f1..d04e830851 100644 --- a/data-prepper-plugins/newline-codecs/build.gradle +++ b/data-prepper-plugins/newline-codecs/build.gradle @@ -2,11 +2,6 @@ plugins { id 'java' } - -repositories { - mavenCentral() -} - dependencies { implementation project(':data-prepper-api') implementation 'com.fasterxml.jackson.core:jackson-annotations' diff --git a/data-prepper-plugins/obfuscate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/obfuscation/ObfuscationProcessor.java b/data-prepper-plugins/obfuscate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/obfuscation/ObfuscationProcessor.java index 0ffc309797..21167bc747 100644 --- a/data-prepper-plugins/obfuscate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/obfuscation/ObfuscationProcessor.java +++ b/data-prepper-plugins/obfuscate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/obfuscation/ObfuscationProcessor.java @@ -29,6 +29,8 @@ import java.util.regex.Matcher; import java.util.regex.Pattern; +import static org.opensearch.dataprepper.logging.DataPrepperMarkers.EVENT; + @DataPrepperPlugin(name = "obfuscate", pluginType = Processor.class, pluginConfigurationType = ObfuscationProcessorConfig.class) public class ObfuscationProcessor extends AbstractProcessor, Record> { @@ -106,29 +108,34 @@ public Collection> doExecute(Collection> records) { for (final Record record : records) { final Event recordEvent = record.getData(); - if (obfuscationProcessorConfig.getObfuscateWhen() != null && !expressionEvaluator.evaluateConditional(obfuscationProcessorConfig.getObfuscateWhen(), recordEvent)) { - continue; - } + try { - if (!recordEvent.containsKey(source)) { - continue; - } + if (obfuscationProcessorConfig.getObfuscateWhen() != null && !expressionEvaluator.evaluateConditional(obfuscationProcessorConfig.getObfuscateWhen(), recordEvent)) { + continue; + } - String rawValue = recordEvent.get(source, String.class); + if (!recordEvent.containsKey(source)) { + continue; + } - // Call obfuscation action - String newValue = this.action.obfuscate(rawValue, patterns); + String rawValue = recordEvent.get(source, String.class); - // No changes means it does not match any patterns - if (rawValue.equals(newValue)) { - recordEvent.getMetadata().addTags(obfuscationProcessorConfig.getTagsOnMatchFailure()); - } + // Call obfuscation action + String newValue = this.action.obfuscate(rawValue, patterns); + + // No changes means it does not match any patterns + if (rawValue.equals(newValue)) { + recordEvent.getMetadata().addTags(obfuscationProcessorConfig.getTagsOnMatchFailure()); + } - // Update the event record. - if (target == null || target.isEmpty()) { - recordEvent.put(source, newValue); - } else { - recordEvent.put(target, newValue); + // Update the event record. + if (target == null || target.isEmpty()) { + recordEvent.put(source, newValue); + } else { + recordEvent.put(target, newValue); + } + } catch (final Exception e) { + LOG.error(EVENT, "There was an exception while processing Event [{}]", recordEvent, e); } } return records; diff --git a/data-prepper-plugins/opensearch/build.gradle b/data-prepper-plugins/opensearch/build.gradle index 13ea63adbe..72ad6b4488 100644 --- a/data-prepper-plugins/opensearch/build.gradle +++ b/data-prepper-plugins/opensearch/build.gradle @@ -11,6 +11,7 @@ dependencies { implementation project(':data-prepper-plugins:buffer-common') implementation project(':data-prepper-plugins:common') implementation project(':data-prepper-plugins:failures-common') + implementation project(':data-prepper-plugins:http-common') implementation libs.opensearch.client implementation libs.opensearch.rhlc implementation libs.opensearch.java diff --git a/data-prepper-plugins/opensearch/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkIT.java b/data-prepper-plugins/opensearch/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkIT.java index 38bb5d1caa..a9d91e78f0 100644 --- a/data-prepper-plugins/opensearch/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkIT.java +++ b/data-prepper-plugins/opensearch/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkIT.java @@ -896,7 +896,7 @@ public void testBulkActionUpdateWithActions() throws IOException, InterruptedExc @Test public void testBulkActionUpdateWithDocumentRootKey() throws IOException, InterruptedException { - final String testIndexAlias = "test-alias-upd1"; + final String testIndexAlias = "test-alias-update"; final String testTemplateFile = Objects.requireNonNull( getClass().getClassLoader().getResource(TEST_TEMPLATE_BULK_FILE)).getFile(); @@ -960,9 +960,40 @@ public void testBulkActionUpdateWithDocumentRootKey() throws IOException, Interr sink.shutdown(); } + @Test + public void testBulkActionUpsertWithActionsAndNoCreate() throws IOException, InterruptedException { + final String testIndexAlias = "test-alias-upsert-no-create2"; + final String testTemplateFile = Objects.requireNonNull( + getClass().getClassLoader().getResource(TEST_TEMPLATE_BULK_FILE)).getFile(); + + final String testIdField = "someId"; + final String testId = "foo"; + List> testRecords = Collections.singletonList(jsonStringToRecord(generateCustomRecordJson2(testIdField, testId, "key", "value"))); + + List> aList = new ArrayList<>(); + Map actionMap = new HashMap<>(); + actionMap.put("type", OpenSearchBulkActions.UPSERT.toString()); + aList.add(actionMap); + + final PluginSetting pluginSetting = generatePluginSetting(null, testIndexAlias, testTemplateFile); + pluginSetting.getSettings().put(IndexConfiguration.DOCUMENT_ID_FIELD, testIdField); + pluginSetting.getSettings().put(IndexConfiguration.ACTIONS, aList); + OpenSearchSink sink = createObjectUnderTest(pluginSetting, true); + + sink.output(testRecords); + List> retSources = getSearchResponseDocSources(testIndexAlias); + + assertThat(retSources.size(), equalTo(1)); + Map source = retSources.get(0); + assertThat((String) source.get("key"), equalTo("value")); + assertThat((String) source.get(testIdField), equalTo(testId)); + assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); + sink.shutdown(); + } + @Test public void testBulkActionUpsertWithActions() throws IOException, InterruptedException { - final String testIndexAlias = "test-alias-upd2"; + final String testIndexAlias = "test-alias-upsert"; final String testTemplateFile = Objects.requireNonNull( getClass().getClassLoader().getResource(TEST_TEMPLATE_BULK_FILE)).getFile(); @@ -1725,6 +1756,7 @@ private void wipeAllOpenSearchIndices() throws IOException { .filter(Predicate.not(indexName -> indexName.startsWith(".opendistro_"))) .filter(Predicate.not(indexName -> indexName.startsWith(".opensearch-"))) .filter(Predicate.not(indexName -> indexName.startsWith(".opensearch_"))) + .filter(Predicate.not(indexName -> indexName.startsWith(".ql"))) .filter(Predicate.not(indexName -> indexName.startsWith(".plugins-ml-config"))) .forEach(indexName -> { try { diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSink.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSink.java index d49949b946..7217e69c1c 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSink.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSink.java @@ -309,7 +309,7 @@ private BulkOperation getBulkOperationForAction(final String action, } - final UpdateOperation.Builder updateOperationBuilder = (action.toLowerCase() == OpenSearchBulkActions.UPSERT.toString()) ? + final UpdateOperation.Builder updateOperationBuilder = (StringUtils.equals(action.toLowerCase(), OpenSearchBulkActions.UPSERT.toString())) ? new UpdateOperation.Builder<>() .index(indexName) .document(filteredJsonNode) diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/ConnectionConfiguration.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/ConnectionConfiguration.java index 23f115923d..8e4dc3f4d6 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/ConnectionConfiguration.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/ConnectionConfiguration.java @@ -5,6 +5,7 @@ package org.opensearch.dataprepper.plugins.source.opensearch.configuration; import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.constraints.AssertTrue; import java.nio.file.Path; import java.time.Duration; @@ -14,6 +15,9 @@ public class ConnectionConfiguration { @JsonProperty("cert") private Path certPath; + @JsonProperty("certificate_content") + private String certificateContent; + @JsonProperty("socket_timeout") private Duration socketTimeout; @@ -27,6 +31,10 @@ public Path getCertPath() { return certPath; } + public String getCertificateContent() { + return certificateContent; + } + public Duration getSocketTimeout() { return socketTimeout; } @@ -38,4 +46,12 @@ public Duration getConnectTimeout() { public boolean isInsecure() { return insecure; } + + @AssertTrue(message = "Certificate file path and certificate content both are configured. " + + "Please use only one configuration.") + boolean certificateFileAndContentAreMutuallyExclusive() { + if(certPath == null && certificateContent == null) + return true; + return certPath != null ^ certificateContent != null; + } } diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchClientFactory.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchClientFactory.java index f15d70df90..977a627bb9 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchClientFactory.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchClientFactory.java @@ -15,13 +15,9 @@ import org.apache.http.auth.UsernamePasswordCredentials; import org.apache.http.client.CredentialsProvider; import org.apache.http.conn.ssl.NoopHostnameVerifier; -import org.apache.http.conn.ssl.TrustAllStrategy; -import org.apache.http.conn.ssl.TrustStrategy; import org.apache.http.impl.client.BasicCredentialsProvider; import org.apache.http.impl.nio.client.HttpAsyncClientBuilder; import org.apache.http.message.BasicHeader; -import org.apache.http.ssl.SSLContextBuilder; -import org.apache.http.ssl.SSLContexts; import org.opensearch.client.RestClient; import org.opensearch.client.RestClientBuilder; import org.opensearch.client.json.jackson.JacksonJsonpMapper; @@ -35,23 +31,17 @@ import org.opensearch.dataprepper.aws.api.AwsRequestSigningApache4Interceptor; import org.opensearch.dataprepper.plugins.source.opensearch.OpenSearchSourceConfiguration; import org.opensearch.dataprepper.plugins.source.opensearch.configuration.ConnectionConfiguration; +import org.opensearch.dataprepper.plugins.truststore.TrustStoreProvider; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; import software.amazon.awssdk.auth.signer.Aws4Signer; -import software.amazon.awssdk.http.apache.ApacheHttpClient; import software.amazon.awssdk.http.async.SdkAsyncHttpClient; import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient; import javax.net.ssl.SSLContext; import javax.net.ssl.TrustManager; -import javax.net.ssl.TrustManagerFactory; -import java.io.InputStream; -import java.nio.file.Files; import java.nio.file.Path; -import java.security.KeyStore; -import java.security.cert.Certificate; -import java.security.cert.CertificateFactory; import java.util.List; import java.util.Objects; @@ -267,20 +257,15 @@ private void setConnectAndSocketTimeout(final org.elasticsearch.client.RestClien }); } - private void attachSSLContext(final ApacheHttpClient.Builder apacheHttpClientBuilder, final OpenSearchSourceConfiguration openSearchSourceConfiguration) { - TrustManager[] trustManagers = createTrustManagers(openSearchSourceConfiguration.getConnectionConfiguration().getCertPath()); - apacheHttpClientBuilder.tlsTrustManagersProvider(() -> trustManagers); - } - private void attachSSLContext(final NettyNioAsyncHttpClient.Builder asyncClientBuilder, final OpenSearchSourceConfiguration openSearchSourceConfiguration) { - TrustManager[] trustManagers = createTrustManagers(openSearchSourceConfiguration.getConnectionConfiguration().getCertPath()); + TrustManager[] trustManagers = createTrustManagers(openSearchSourceConfiguration.getConnectionConfiguration()); asyncClientBuilder.tlsTrustManagersProvider(() -> trustManagers); } private void attachSSLContext(final HttpAsyncClientBuilder httpClientBuilder, final OpenSearchSourceConfiguration openSearchSourceConfiguration) { final ConnectionConfiguration connectionConfiguration = openSearchSourceConfiguration.getConnectionConfiguration(); - final SSLContext sslContext = Objects.nonNull(connectionConfiguration.getCertPath()) ? getCAStrategy(connectionConfiguration.getCertPath()) : getTrustAllStrategy(); + final SSLContext sslContext = getCAStrategy(connectionConfiguration); httpClientBuilder.setSSLContext(sslContext); if (connectionConfiguration.isInsecure()) { @@ -288,53 +273,25 @@ private void attachSSLContext(final HttpAsyncClientBuilder httpClientBuilder, fi } } - private static TrustManager[] createTrustManagers(final Path certPath) { - if (certPath != null) { - LOG.info("Using the cert provided in the config."); - try (InputStream certificateInputStream = Files.newInputStream(certPath)) { - final CertificateFactory factory = CertificateFactory.getInstance("X.509"); - final Certificate trustedCa = factory.generateCertificate(certificateInputStream); - final KeyStore trustStore = KeyStore.getInstance("pkcs12"); - trustStore.load(null, null); - trustStore.setCertificateEntry("ca", trustedCa); - - final TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance("X509"); - trustManagerFactory.init(trustStore); - return trustManagerFactory.getTrustManagers(); - } catch (Exception ex) { - throw new RuntimeException(ex.getMessage(), ex); - } + private TrustManager[] createTrustManagers(final ConnectionConfiguration connectionConfiguration) { + final Path certPath = connectionConfiguration.getCertPath(); + if (Objects.nonNull(certPath)) { + return TrustStoreProvider.createTrustManager(certPath); + } else if (Objects.nonNull(connectionConfiguration.getCertificateContent())) { + return TrustStoreProvider.createTrustManager(connectionConfiguration.getCertificateContent()); } else { - return new TrustManager[] { new X509TrustAllManager() }; - } - } - - private SSLContext getCAStrategy(final Path certPath) { - LOG.info("Using the cert provided in the config."); - try { - CertificateFactory factory = CertificateFactory.getInstance("X.509"); - Certificate trustedCa; - try (InputStream is = Files.newInputStream(certPath)) { - trustedCa = factory.generateCertificate(is); - } - KeyStore trustStore = KeyStore.getInstance("pkcs12"); - trustStore.load(null, null); - trustStore.setCertificateEntry("ca", trustedCa); - SSLContextBuilder sslContextBuilder = SSLContexts.custom() - .loadTrustMaterial(trustStore, null); - return sslContextBuilder.build(); - } catch (Exception ex) { - throw new RuntimeException(ex.getMessage(), ex); + return TrustStoreProvider.createTrustAllManager(); } } - private SSLContext getTrustAllStrategy() { - LOG.info("Using the trust all strategy"); - final TrustStrategy trustStrategy = new TrustAllStrategy(); - try { - return SSLContexts.custom().loadTrustMaterial(null, trustStrategy).build(); - } catch (Exception ex) { - throw new RuntimeException(ex.getMessage(), ex); + private SSLContext getCAStrategy(final ConnectionConfiguration connectionConfiguration) { + final Path certPath = connectionConfiguration.getCertPath(); + if (Objects.nonNull(certPath)) { + return TrustStoreProvider.createSSLContext(certPath); + } else if (Objects.nonNull(connectionConfiguration.getCertificateContent())) { + return TrustStoreProvider.createSSLContext(connectionConfiguration.getCertificateContent()); + } else { + return TrustStoreProvider.createSSLContextWithTrustAllStrategy(); } } } diff --git a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/ConnectionConfigurationTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/ConnectionConfigurationTest.java index 75e34a08dd..97d5152080 100644 --- a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/ConnectionConfigurationTest.java +++ b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/ConnectionConfigurationTest.java @@ -14,6 +14,7 @@ import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.CoreMatchers.is; public class ConnectionConfigurationTest { private final ObjectMapper objectMapper = new ObjectMapper(new YAMLFactory().enable(YAMLGenerator.Feature.USE_PLATFORM_LINE_BREAKS)); @@ -39,4 +40,20 @@ void connection_configuration_values_test() throws JsonProcessingException { assertThat(connectionConfig.getConnectTimeout(),equalTo(null)); assertThat(connectionConfig.isInsecure(),equalTo(true)); } + + @Test + void connection_configuration_certificate_values_test() throws JsonProcessingException { + + final String connectionYaml = + " cert: \"cert\"\n" + + " certificate_content: \"certificate content\"\n" + + " insecure: true\n"; + final ConnectionConfiguration connectionConfig = objectMapper.readValue(connectionYaml, ConnectionConfiguration.class); + assertThat(connectionConfig.getCertPath(),equalTo(Path.of("cert"))); + assertThat(connectionConfig.getCertificateContent(),equalTo("certificate content")); + assertThat(connectionConfig.certificateFileAndContentAreMutuallyExclusive(), is(false)); + assertThat(connectionConfig.getSocketTimeout(),equalTo(null)); + assertThat(connectionConfig.getConnectTimeout(),equalTo(null)); + assertThat(connectionConfig.isInsecure(),equalTo(true)); + } } diff --git a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchClientFactoryTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchClientFactoryTest.java index 097d0e1250..a357451147 100644 --- a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchClientFactoryTest.java +++ b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchClientFactoryTest.java @@ -11,6 +11,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.ArgumentCaptor; import org.mockito.Mock; +import org.mockito.MockedStatic; import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.client.opensearch.OpenSearchClient; import org.opensearch.dataprepper.aws.api.AwsCredentialsOptions; @@ -18,9 +19,14 @@ import org.opensearch.dataprepper.plugins.source.opensearch.OpenSearchSourceConfiguration; import org.opensearch.dataprepper.plugins.source.opensearch.configuration.AwsAuthenticationConfiguration; import org.opensearch.dataprepper.plugins.source.opensearch.configuration.ConnectionConfiguration; +import org.opensearch.dataprepper.plugins.truststore.TrustStoreProvider; import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.http.async.SdkAsyncHttpClient; import software.amazon.awssdk.regions.Region; +import javax.net.ssl.SSLContext; +import java.nio.file.Path; +import java.time.Duration; import java.util.Collections; import java.util.List; import java.util.UUID; @@ -28,7 +34,9 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.notNullValue; +import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; @@ -48,7 +56,7 @@ public class OpenSearchClientFactoryTest { @BeforeEach void setup() { - when(openSearchSourceConfiguration.getHosts()).thenReturn(List.of("http://localhost:9200")); + lenient().when(openSearchSourceConfiguration.getHosts()).thenReturn(List.of("http://localhost:9200")); when(openSearchSourceConfiguration.getConnectionConfiguration()).thenReturn(connectionConfiguration); } @@ -209,4 +217,59 @@ void provideOpenSearchClient_with_aws_auth_and_serverless_flag_true() { assertThat(awsCredentialsOptions.getStsHeaderOverrides(), equalTo(Collections.emptyMap())); assertThat(awsCredentialsOptions.getStsRoleArn(), equalTo(stsRoleArn)); } + + @Test + void provideOpenSearchClient_with_self_signed_certificate() { + final Path path = mock(Path.class); + final SSLContext sslContext = mock(SSLContext.class); + final String username = UUID.randomUUID().toString(); + final String password = UUID.randomUUID().toString(); + when(openSearchSourceConfiguration.getUsername()).thenReturn(username); + when(openSearchSourceConfiguration.getPassword()).thenReturn(password); + when(connectionConfiguration.getCertPath()).thenReturn(path); + try (MockedStatic trustStoreProviderMockedStatic = mockStatic(TrustStoreProvider.class)) { + trustStoreProviderMockedStatic.when(() -> TrustStoreProvider.createSSLContext(path)) + .thenReturn(sslContext); + final OpenSearchClient openSearchClient = createObjectUnderTest().provideOpenSearchClient(openSearchSourceConfiguration); + trustStoreProviderMockedStatic.verify(() -> TrustStoreProvider.createSSLContext(path)); + assertThat(openSearchClient, notNullValue()); + } + } + + @Test + void provideElasticSearchClient_with_self_signed_certificate() { + final Path path = mock(Path.class); + final SSLContext sslContext = mock(SSLContext.class); + final String username = UUID.randomUUID().toString(); + final String password = UUID.randomUUID().toString(); + when(openSearchSourceConfiguration.getUsername()).thenReturn(username); + when(openSearchSourceConfiguration.getPassword()).thenReturn(password); + when(connectionConfiguration.getCertPath()).thenReturn(path); + try (MockedStatic trustStoreProviderMockedStatic = mockStatic(TrustStoreProvider.class)) { + trustStoreProviderMockedStatic.when(() -> TrustStoreProvider.createSSLContext(path)) + .thenReturn(sslContext); + final ElasticsearchClient elasticsearchClient = createObjectUnderTest().provideElasticSearchClient(openSearchSourceConfiguration); + assertThat(elasticsearchClient, notNullValue()); + trustStoreProviderMockedStatic.verify(() -> TrustStoreProvider.createSSLContext(path)); + } + } + + + @Test + void createSdkAsyncHttpClient_with_self_signed_certificate() { + final Path path = mock(Path.class); + final Duration duration = mock(Duration.class); + final String username = UUID.randomUUID().toString(); + final String password = UUID.randomUUID().toString(); + lenient().when(openSearchSourceConfiguration.getUsername()).thenReturn(username); + lenient().when(openSearchSourceConfiguration.getPassword()).thenReturn(password); + lenient().when(connectionConfiguration.getConnectTimeout()).thenReturn(duration); + lenient().when(openSearchSourceConfiguration.getConnectionConfiguration()).thenReturn(connectionConfiguration); + lenient().when(connectionConfiguration.getCertPath()).thenReturn(path); + try (MockedStatic trustStoreProviderMockedStatic = mockStatic(TrustStoreProvider.class)) { + final SdkAsyncHttpClient sdkAsyncHttpClient = createObjectUnderTest().createSdkAsyncHttpClient(openSearchSourceConfiguration); + assertThat(sdkAsyncHttpClient, notNullValue()); + trustStoreProviderMockedStatic.verify(() -> TrustStoreProvider.createTrustManager(path)); + } + } } diff --git a/data-prepper-plugins/otel-logs-source/src/main/java/org/opensearch/dataprepper/plugins/source/otellogs/OTelLogsGrpcService.java b/data-prepper-plugins/otel-logs-source/src/main/java/org/opensearch/dataprepper/plugins/source/otellogs/OTelLogsGrpcService.java index 4a63dabac0..e800e48e74 100644 --- a/data-prepper-plugins/otel-logs-source/src/main/java/org/opensearch/dataprepper/plugins/source/otellogs/OTelLogsGrpcService.java +++ b/data-prepper-plugins/otel-logs-source/src/main/java/org/opensearch/dataprepper/plugins/source/otellogs/OTelLogsGrpcService.java @@ -25,6 +25,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.time.Instant; import java.util.List; import java.util.stream.Collectors; @@ -84,7 +85,7 @@ private void processRequest(final ExportLogsServiceRequest request, final Stream final List logs; try { - logs = oTelProtoDecoder.parseExportLogsServiceRequest(request); + logs = oTelProtoDecoder.parseExportLogsServiceRequest(request, Instant.now()); } catch (Exception e) { LOG.error("Failed to parse the request {} due to:", request, e); throw new BadRequestException(e.getMessage(), e); diff --git a/data-prepper-plugins/otel-logs-source/src/test/java/org/opensearch/dataprepper/plugins/source/otellogs/OTelLogsGrpcServiceTest.java b/data-prepper-plugins/otel-logs-source/src/test/java/org/opensearch/dataprepper/plugins/source/otellogs/OTelLogsGrpcServiceTest.java index a5a9353981..28b8cbeca4 100644 --- a/data-prepper-plugins/otel-logs-source/src/test/java/org/opensearch/dataprepper/plugins/source/otellogs/OTelLogsGrpcServiceTest.java +++ b/data-prepper-plugins/otel-logs-source/src/test/java/org/opensearch/dataprepper/plugins/source/otellogs/OTelLogsGrpcServiceTest.java @@ -35,6 +35,7 @@ import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoCodec; import java.io.IOException; +import java.time.Instant; import java.util.Collection; import java.util.Collections; import java.util.List; @@ -189,7 +190,7 @@ public void export_BufferTimeout_responseObserverOnError() throws Exception { public void export_BadRequest_responseObserverOnError() { final String testMessage = "test message"; final RuntimeException testException = new RuntimeException(testMessage); - when(mockOTelProtoDecoder.parseExportLogsServiceRequest(any())).thenThrow(testException); + when(mockOTelProtoDecoder.parseExportLogsServiceRequest(any(), any(Instant.class))).thenThrow(testException); objectUnderTest = generateOTelLogsGrpcService(mockOTelProtoDecoder); try (MockedStatic mockedStatic = mockStatic(ServiceRequestContext.class)) { diff --git a/data-prepper-plugins/otel-metrics-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/otelmetrics/OTelMetricsRawProcessor.java b/data-prepper-plugins/otel-metrics-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/otelmetrics/OTelMetricsRawProcessor.java index 679eef3224..87ec477a26 100644 --- a/data-prepper-plugins/otel-metrics-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/otelmetrics/OTelMetricsRawProcessor.java +++ b/data-prepper-plugins/otel-metrics-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/otelmetrics/OTelMetricsRawProcessor.java @@ -12,8 +12,8 @@ import static org.opensearch.dataprepper.model.metric.JacksonExponentialHistogram.POSITIVE_BUCKETS_KEY; import static org.opensearch.dataprepper.model.metric.JacksonExponentialHistogram.NEGATIVE_BUCKETS_KEY; import static org.opensearch.dataprepper.model.metric.JacksonHistogram.BUCKETS_KEY; +import org.opensearch.dataprepper.model.metric.JacksonMetric; import org.opensearch.dataprepper.model.metric.Metric; -import static org.opensearch.dataprepper.model.metric.JacksonMetric.ATTRIBUTES_KEY; import org.opensearch.dataprepper.model.processor.AbstractProcessor; import org.opensearch.dataprepper.model.processor.Processor; import org.opensearch.dataprepper.model.record.Record; @@ -23,9 +23,9 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.time.Instant; import java.util.ArrayList; import java.util.Collection; -import java.util.Map; import java.util.List; import java.util.concurrent.atomic.AtomicInteger; @@ -55,13 +55,8 @@ private void modifyRecord(Record record, boolean calcualteExponentialHistogramBuckets) { Event event = (Event)record.getData(); - if (flattenAttributes) { - Map attributes = event.get(ATTRIBUTES_KEY, Map.class); - - for (Map.Entry entry : attributes.entrySet()) { - event.put(entry.getKey(), entry.getValue()); - } - event.delete(ATTRIBUTES_KEY); + if (!flattenAttributes) { + ((JacksonMetric)record.getData()).setFlattenAttributes(false); } if (!calcualteHistogramBuckets && event.get(BUCKETS_KEY, List.class) != null) { event.delete(BUCKETS_KEY); @@ -85,7 +80,7 @@ public Collection> doExecute(Collection> reco for (Record rec : records) { if ((rec.getData() instanceof Event)) { Record newRecord = (Record)rec; - if (otelMetricsRawProcessorConfig.getFlattenAttributesFlag() || + if (!otelMetricsRawProcessorConfig.getFlattenAttributesFlag() || !otelMetricsRawProcessorConfig.getCalculateHistogramBuckets() || !otelMetricsRawProcessorConfig.getCalculateExponentialHistogramBuckets()) { modifyRecord(newRecord, otelMetricsRawProcessorConfig.getFlattenAttributesFlag(), otelMetricsRawProcessorConfig.getCalculateHistogramBuckets(), otelMetricsRawProcessorConfig.getCalculateExponentialHistogramBuckets()); @@ -98,7 +93,7 @@ public Collection> doExecute(Collection> reco } ExportMetricsServiceRequest request = ((Record)rec).getData(); - recordsOut.addAll(otelProtoDecoder.parseExportMetricsServiceRequest(request, droppedCounter, otelMetricsRawProcessorConfig.getExponentialHistogramMaxAllowedScale(), otelMetricsRawProcessorConfig.getCalculateHistogramBuckets(), otelMetricsRawProcessorConfig.getCalculateExponentialHistogramBuckets(), flattenAttributesFlag)); + recordsOut.addAll(otelProtoDecoder.parseExportMetricsServiceRequest(request, droppedCounter, otelMetricsRawProcessorConfig.getExponentialHistogramMaxAllowedScale(), Instant.now(), otelMetricsRawProcessorConfig.getCalculateHistogramBuckets(), otelMetricsRawProcessorConfig.getCalculateExponentialHistogramBuckets(), flattenAttributesFlag)); } recordsDroppedMetricsRawCounter.increment(droppedCounter.get()); return recordsOut; diff --git a/data-prepper-plugins/otel-metrics-source/src/main/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsGrpcService.java b/data-prepper-plugins/otel-metrics-source/src/main/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsGrpcService.java index b4c45e5a05..e88cf18b5f 100644 --- a/data-prepper-plugins/otel-metrics-source/src/main/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsGrpcService.java +++ b/data-prepper-plugins/otel-metrics-source/src/main/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsGrpcService.java @@ -16,48 +16,55 @@ import io.opentelemetry.proto.collector.metrics.v1.MetricsServiceGrpc; import org.opensearch.dataprepper.exceptions.BufferWriteException; import org.opensearch.dataprepper.exceptions.RequestCancelledException; +import static org.opensearch.dataprepper.plugins.otel.codec.OTelProtoCodec.DEFAULT_EXPONENTIAL_HISTOGRAM_MAX_ALLOWED_SCALE; +import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoCodec; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.buffer.Buffer; import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.metric.Metric; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.time.Instant; +import java.util.Collection; +import java.util.concurrent.atomic.AtomicInteger; + public class OTelMetricsGrpcService extends MetricsServiceGrpc.MetricsServiceImplBase { private static final Logger LOG = LoggerFactory.getLogger(OTelMetricsGrpcService.class); public static final String REQUESTS_RECEIVED = "requestsReceived"; public static final String SUCCESS_REQUESTS = "successRequests"; + public static final String RECORDS_CREATED = "recordsCreated"; + public static final String RECORDS_DROPPED = "recordsDropped"; public static final String PAYLOAD_SIZE = "payloadSize"; public static final String REQUEST_PROCESS_DURATION = "requestProcessDuration"; private final int bufferWriteTimeoutInMillis; - private final Buffer> buffer; + private final OTelProtoCodec.OTelProtoDecoder oTelProtoDecoder; + private final Buffer> buffer; private final Counter requestsReceivedCounter; private final Counter successRequestsCounter; + private final Counter recordsCreatedCounter; + private final Counter recordsDroppedCounter; private final DistributionSummary payloadSizeSummary; private final Timer requestProcessDuration; public OTelMetricsGrpcService(int bufferWriteTimeoutInMillis, - Buffer> buffer, + final OTelProtoCodec.OTelProtoDecoder oTelProtoDecoder, + Buffer> buffer, final PluginMetrics pluginMetrics) { this.bufferWriteTimeoutInMillis = bufferWriteTimeoutInMillis; this.buffer = buffer; requestsReceivedCounter = pluginMetrics.counter(REQUESTS_RECEIVED); successRequestsCounter = pluginMetrics.counter(SUCCESS_REQUESTS); + recordsCreatedCounter = pluginMetrics.counter(RECORDS_CREATED); + recordsDroppedCounter = pluginMetrics.counter(RECORDS_DROPPED); payloadSizeSummary = pluginMetrics.summary(PAYLOAD_SIZE); requestProcessDuration = pluginMetrics.timer(REQUEST_PROCESS_DURATION); - } - - public void rawExport(final ExportMetricsServiceRequest request) { - try { - if (buffer.isByteBuffer()) { - buffer.writeBytes(request.toByteArray(), null, bufferWriteTimeoutInMillis); - } - } catch (Exception e) { - } + this.oTelProtoDecoder = oTelProtoDecoder; } @Override @@ -81,7 +88,13 @@ private void processRequest(final ExportMetricsServiceRequest request, final Str if (buffer.isByteBuffer()) { buffer.writeBytes(request.toByteArray(), null, bufferWriteTimeoutInMillis); } else { - buffer.write(new Record<>(request), bufferWriteTimeoutInMillis); + Collection> metrics; + + AtomicInteger droppedCounter = new AtomicInteger(0); + metrics = oTelProtoDecoder.parseExportMetricsServiceRequest(request, droppedCounter, DEFAULT_EXPONENTIAL_HISTOGRAM_MAX_ALLOWED_SCALE, Instant.now(), true, true, true); + recordsDroppedCounter.increment(droppedCounter.get()); + recordsCreatedCounter.increment(metrics.size()); + buffer.writeAll(metrics, bufferWriteTimeoutInMillis); } } catch (Exception e) { if (ServiceRequestContext.current().isTimedOut()) { diff --git a/data-prepper-plugins/otel-metrics-source/src/main/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsSource.java b/data-prepper-plugins/otel-metrics-source/src/main/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsSource.java index 85e6982e23..337006cc71 100644 --- a/data-prepper-plugins/otel-metrics-source/src/main/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsSource.java +++ b/data-prepper-plugins/otel-metrics-source/src/main/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsSource.java @@ -31,9 +31,11 @@ import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.plugin.PluginFactory; import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.metric.Metric; import org.opensearch.dataprepper.model.source.Source; import org.opensearch.dataprepper.model.codec.ByteDecoder; import org.opensearch.dataprepper.plugins.otel.codec.OTelMetricDecoder; +import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoCodec; import org.opensearch.dataprepper.plugins.certificate.CertificateProvider; import org.opensearch.dataprepper.plugins.certificate.model.Certificate; import org.opensearch.dataprepper.plugins.health.HealthGrpcService; @@ -51,7 +53,7 @@ import java.util.function.Function; @DataPrepperPlugin(name = "otel_metrics_source", pluginType = Source.class, pluginConfigurationType = OTelMetricsSourceConfig.class) -public class OTelMetricsSource implements Source> { +public class OTelMetricsSource implements Source> { private static final Logger LOG = LoggerFactory.getLogger(OTelMetricsSource.class); private static final String HTTP_HEALTH_CHECK_PATH = "/health"; private static final String REGEX_HEALTH = "regex:^/(?!health$).*$"; @@ -91,15 +93,15 @@ public ByteDecoder getDecoder() { } @Override - public void start(Buffer> buffer) { + public void start(Buffer> buffer) { if (buffer == null) { throw new IllegalStateException("Buffer provided is null"); } if (server == null) { - final OTelMetricsGrpcService oTelMetricsGrpcService = new OTelMetricsGrpcService( (int) (oTelMetricsSourceConfig.getRequestTimeoutInMillis() * 0.8), + new OTelProtoCodec.OTelProtoDecoder(), buffer, pluginMetrics ); diff --git a/data-prepper-plugins/otel-metrics-source/src/test/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsGrpcServiceTest.java b/data-prepper-plugins/otel-metrics-source/src/test/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsGrpcServiceTest.java index f257f64f83..be5c1c817d 100644 --- a/data-prepper-plugins/otel-metrics-source/src/test/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsGrpcServiceTest.java +++ b/data-prepper-plugins/otel-metrics-source/src/test/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsGrpcServiceTest.java @@ -13,6 +13,7 @@ import io.opentelemetry.proto.collector.metrics.v1.ExportMetricsServiceRequest; import io.opentelemetry.proto.collector.metrics.v1.ExportMetricsServiceResponse; import io.opentelemetry.proto.metrics.v1.InstrumentationLibraryMetrics; +import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoCodec; import io.opentelemetry.proto.metrics.v1.Metric; import io.opentelemetry.proto.metrics.v1.ResourceMetrics; import org.junit.jupiter.api.BeforeEach; @@ -29,17 +30,20 @@ import org.opensearch.dataprepper.model.buffer.Buffer; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.event.Event; +import java.util.Collection; import java.util.Collections; +import java.util.Map; import java.util.concurrent.TimeoutException; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; -import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; +import static org.hamcrest.Matchers.hasEntry; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; @@ -48,15 +52,20 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; +import io.opentelemetry.proto.metrics.v1.NumberDataPoint; +import io.opentelemetry.proto.metrics.v1.Gauge; @ExtendWith(MockitoExtension.class) public class OTelMetricsGrpcServiceTest { + private static NumberDataPoint.Builder p1 = NumberDataPoint.newBuilder().setAsInt(4); + private static Gauge gauge = Gauge.newBuilder().addDataPoints(p1).build(); private static final ExportMetricsServiceRequest METRICS_REQUEST = ExportMetricsServiceRequest.newBuilder() .addResourceMetrics(ResourceMetrics.newBuilder() .addInstrumentationLibraryMetrics(InstrumentationLibraryMetrics.newBuilder() - .addMetrics(Metric.newBuilder().build()) + .addMetrics(Metric.newBuilder().setGauge(gauge).setUnit("seconds").setName("name").build()) .build())).build(); + private static Map expectedMetric = Map.of("unit", (Object)"seconds", "name", (Object)"name", "kind", (Object)"GAUGE"); private static PluginSetting pluginSetting; private final int bufferWriteTimeoutInMillis = 100000; @@ -65,6 +74,10 @@ public class OTelMetricsGrpcServiceTest { @Mock private Counter successRequestsCounter; @Mock + private Counter droppedCounter; + @Mock + private Counter createdCounter; + @Mock private DistributionSummary payloadSize; @Mock private Timer requestProcessDuration; @@ -76,7 +89,7 @@ public class OTelMetricsGrpcServiceTest { private ServiceRequestContext serviceRequestContext; @Captor - private ArgumentCaptor recordCaptor; + private ArgumentCaptor> recordCaptor; @Captor ArgumentCaptor bytesCaptor; @@ -92,6 +105,8 @@ public void setup() { when(mockPluginMetrics.counter(OTelMetricsGrpcService.REQUESTS_RECEIVED)).thenReturn(requestsReceivedCounter); when(mockPluginMetrics.counter(OTelMetricsGrpcService.SUCCESS_REQUESTS)).thenReturn(successRequestsCounter); + when(mockPluginMetrics.counter(OTelMetricsGrpcService.RECORDS_CREATED)).thenReturn(createdCounter); + when(mockPluginMetrics.counter(OTelMetricsGrpcService.RECORDS_DROPPED)).thenReturn(droppedCounter); when(mockPluginMetrics.summary(OTelMetricsGrpcService.PAYLOAD_SIZE)).thenReturn(payloadSize); when(mockPluginMetrics.timer(OTelMetricsGrpcService.REQUEST_PROCESS_DURATION)).thenReturn(requestProcessDuration); doAnswer(invocation -> { @@ -101,7 +116,7 @@ public void setup() { when(serviceRequestContext.isTimedOut()).thenReturn(false); - sut = new OTelMetricsGrpcService(bufferWriteTimeoutInMillis, buffer, mockPluginMetrics); + sut = new OTelMetricsGrpcService(bufferWriteTimeoutInMillis, new OTelProtoCodec.OTelProtoDecoder(), buffer, mockPluginMetrics); } @Test @@ -111,7 +126,7 @@ public void export_Success_responseObserverOnCompleted() throws Exception { sut.export(METRICS_REQUEST, responseObserver); } - verify(buffer, times(1)).write(recordCaptor.capture(), anyInt()); + verify(buffer, times(1)).writeAll(recordCaptor.capture(), anyInt()); verify(responseObserver, times(1)).onNext(ExportMetricsServiceResponse.newBuilder().build()); verify(responseObserver, times(1)).onCompleted(); verify(requestsReceivedCounter, times(1)).increment(); @@ -122,8 +137,11 @@ public void export_Success_responseObserverOnCompleted() throws Exception { assertThat(payloadLengthCaptor.getValue().intValue(), equalTo(METRICS_REQUEST.getSerializedSize())); verify(requestProcessDuration, times(1)).record(ArgumentMatchers.any()); - Record capturedRecord = recordCaptor.getValue(); - assertEquals(METRICS_REQUEST, capturedRecord.getData()); + Collection capturedRecords = recordCaptor.getValue(); + Record capturedRecord = (Record)(capturedRecords.toArray()[0]); + Map map = ((Event)capturedRecord.getData()).toMap(); + + expectedMetric.forEach((k, v) -> assertThat(map, hasEntry((String)k, (Object)v))); } @Test @@ -151,14 +169,14 @@ public void export_Success_with_ByteBuffer_responseObserverOnCompleted() throws @Test public void export_BufferTimeout_responseObserverOnError() throws Exception { - doThrow(new TimeoutException()).when(buffer).write(any(Record.class), anyInt()); + doThrow(new TimeoutException()).when(buffer).writeAll(any(Collection.class), anyInt()); try (MockedStatic mockedStatic = mockStatic(ServiceRequestContext.class)) { mockedStatic.when(ServiceRequestContext::current).thenReturn(serviceRequestContext); assertThrows(BufferWriteException.class, () -> sut.export(METRICS_REQUEST, responseObserver)); } - verify(buffer, times(1)).write(any(Record.class), anyInt()); + verify(buffer, times(1)).writeAll(any(Collection.class), anyInt()); verifyNoInteractions(responseObserver); verify(requestsReceivedCounter, times(1)).increment(); verifyNoInteractions(successRequestsCounter); diff --git a/data-prepper-plugins/otel-metrics-source/src/test/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsSourceTest.java b/data-prepper-plugins/otel-metrics-source/src/test/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsSourceTest.java index 66cab56203..daf6ae363a 100644 --- a/data-prepper-plugins/otel-metrics-source/src/test/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsSourceTest.java +++ b/data-prepper-plugins/otel-metrics-source/src/test/java/org/opensearch/dataprepper/plugins/source/otelmetrics/OTelMetricsSourceTest.java @@ -34,6 +34,11 @@ import io.opentelemetry.proto.collector.metrics.v1.ExportMetricsServiceRequest; import io.opentelemetry.proto.collector.metrics.v1.ExportMetricsServiceResponse; import io.opentelemetry.proto.collector.metrics.v1.MetricsServiceGrpc; +import io.opentelemetry.proto.metrics.v1.NumberDataPoint; + +import io.opentelemetry.proto.common.v1.InstrumentationLibrary; +import io.opentelemetry.proto.metrics.v1.Gauge; +import io.opentelemetry.proto.metrics.v1.InstrumentationLibraryMetrics; import io.opentelemetry.proto.common.v1.AnyValue; import io.opentelemetry.proto.common.v1.KeyValue; import io.opentelemetry.proto.metrics.v1.ResourceMetrics; @@ -62,6 +67,7 @@ import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.plugin.PluginFactory; import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.metric.Metric; import org.opensearch.dataprepper.model.types.ByteCount; import org.opensearch.dataprepper.plugins.GrpcBasicAuthenticationProvider; import org.opensearch.dataprepper.plugins.buffer.blockingbuffer.BlockingBuffer; @@ -79,6 +85,7 @@ import java.nio.file.Files; import java.nio.file.Path; import java.util.Base64; +import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.List; @@ -163,7 +170,7 @@ class OTelMetricsSourceTest { private OTelMetricsSourceConfig oTelMetricsSourceConfig; @Mock - private BlockingBuffer> buffer; + private BlockingBuffer> buffer; @Mock private HttpBasicAuthenticationConfig httpBasicAuthenticationConfig; @@ -901,12 +908,12 @@ void gRPC_request_writes_to_buffer_with_successful_response() throws Exception { final ExportMetricsServiceResponse exportResponse = client.export(createExportMetricsRequest()); assertThat(exportResponse, notNullValue()); - final ArgumentCaptor> bufferWriteArgumentCaptor = ArgumentCaptor.forClass(Record.class); - verify(buffer).write(bufferWriteArgumentCaptor.capture(), anyInt()); + final ArgumentCaptor>> bufferWriteArgumentCaptor = ArgumentCaptor.forClass(Collection.class); + verify(buffer).writeAll(bufferWriteArgumentCaptor.capture(), anyInt()); - final Record actualBufferWrites = bufferWriteArgumentCaptor.getValue(); + final Collection> actualBufferWrites = bufferWriteArgumentCaptor.getValue(); assertThat(actualBufferWrites, notNullValue()); - assertThat(actualBufferWrites.getData().getResourceMetricsCount(), equalTo(1)); + assertThat(actualBufferWrites.size(), equalTo(1)); } @Test @@ -935,12 +942,12 @@ void gRPC_with_auth_request_writes_to_buffer_with_successful_response() throws E final ExportMetricsServiceResponse exportResponse = client.export(createExportMetricsRequest()); assertThat(exportResponse, notNullValue()); - final ArgumentCaptor> bufferWriteArgumentCaptor = ArgumentCaptor.forClass(Record.class); - verify(buffer).write(bufferWriteArgumentCaptor.capture(), anyInt()); + final ArgumentCaptor>> bufferWriteArgumentCaptor = ArgumentCaptor.forClass(Collection.class); + verify(buffer).writeAll(bufferWriteArgumentCaptor.capture(), anyInt()); - final Record actualBufferWrites = bufferWriteArgumentCaptor.getValue(); + final Collection> actualBufferWrites = bufferWriteArgumentCaptor.getValue(); assertThat(actualBufferWrites, notNullValue()); - assertThat(actualBufferWrites.getData().getResourceMetricsCount(), equalTo(1)); + assertThat(actualBufferWrites.size(), equalTo(1)); } @Test @@ -971,7 +978,7 @@ void gRPC_request_returns_expected_status_for_exceptions_from_buffer( doThrow(bufferExceptionClass) .when(buffer) - .write(any(Record.class), anyInt()); + .writeAll(any(Collection.class), anyInt()); final ExportMetricsServiceRequest exportMetricsRequest = createExportMetricsRequest(); final StatusRuntimeException actualException = assertThrows(StatusRuntimeException.class, () -> client.export(exportMetricsRequest)); @@ -1014,9 +1021,26 @@ private ExportMetricsServiceRequest createExportMetricsRequest() { .setKey("service.name") .setValue(AnyValue.newBuilder().setStringValue("service").build()) ).build(); + NumberDataPoint.Builder p1 = NumberDataPoint.newBuilder().setAsInt(4); + Gauge gauge = Gauge.newBuilder().addDataPoints(p1).build(); + + io.opentelemetry.proto.metrics.v1.Metric.Builder metric = io.opentelemetry.proto.metrics.v1.Metric.newBuilder() + .setGauge(gauge) + .setUnit("seconds") + .setName("name") + .setDescription("description"); + InstrumentationLibraryMetrics isntLib = InstrumentationLibraryMetrics.newBuilder() + .addMetrics(metric) + .setInstrumentationLibrary(InstrumentationLibrary.newBuilder() + .setName("ilname") + .setVersion("ilversion") + .build()) + .build(); + final ResourceMetrics resourceMetrics = ResourceMetrics.newBuilder() .setResource(resource) + .addInstrumentationLibraryMetrics(isntLib) .build(); return ExportMetricsServiceRequest.newBuilder() diff --git a/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelLogsDecoder.java b/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelLogsDecoder.java index 6c491e7510..c140e9591d 100644 --- a/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelLogsDecoder.java +++ b/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelLogsDecoder.java @@ -16,6 +16,7 @@ import java.io.IOException; import java.io.InputStream; import java.util.function.Consumer; +import java.time.Instant; public class OTelLogsDecoder implements ByteDecoder { @@ -23,10 +24,10 @@ public class OTelLogsDecoder implements ByteDecoder { public OTelLogsDecoder() { otelProtoDecoder = new OTelProtoCodec.OTelProtoDecoder(); } - public void parse(InputStream inputStream, Consumer> eventConsumer) throws IOException { + public void parse(InputStream inputStream, Instant timeReceivedMs, Consumer> eventConsumer) throws IOException { ExportLogsServiceRequest request = ExportLogsServiceRequest.parseFrom(inputStream); AtomicInteger droppedCounter = new AtomicInteger(0); - List logs = otelProtoDecoder.parseExportLogsServiceRequest(request); + List logs = otelProtoDecoder.parseExportLogsServiceRequest(request, timeReceivedMs); for (OpenTelemetryLog log: logs) { eventConsumer.accept(new Record<>(log)); } diff --git a/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelMetricDecoder.java b/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelMetricDecoder.java index 1918842c22..d7bdd0c5fd 100644 --- a/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelMetricDecoder.java +++ b/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelMetricDecoder.java @@ -16,6 +16,7 @@ import java.util.concurrent.atomic.AtomicInteger; import java.io.IOException; import java.io.InputStream; +import java.time.Instant; import java.util.function.Consumer; @@ -24,11 +25,11 @@ public class OTelMetricDecoder implements ByteDecoder { public OTelMetricDecoder() { otelProtoDecoder = new OTelProtoCodec.OTelProtoDecoder(); } - public void parse(InputStream inputStream, Consumer> eventConsumer) throws IOException { + public void parse(InputStream inputStream, Instant timeReceivedMs, Consumer> eventConsumer) throws IOException { ExportMetricsServiceRequest request = ExportMetricsServiceRequest.parseFrom(inputStream); AtomicInteger droppedCounter = new AtomicInteger(0); Collection> records = - otelProtoDecoder.parseExportMetricsServiceRequest(request, droppedCounter, OTelProtoCodec.DEFAULT_EXPONENTIAL_HISTOGRAM_MAX_ALLOWED_SCALE, true, true, false); + otelProtoDecoder.parseExportMetricsServiceRequest(request, droppedCounter, OTelProtoCodec.DEFAULT_EXPONENTIAL_HISTOGRAM_MAX_ALLOWED_SCALE, timeReceivedMs, true, true, false); for (Record record: records) { eventConsumer.accept((Record)record); } diff --git a/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelProtoCodec.java b/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelProtoCodec.java index d4a5eab2ff..8875e90e53 100644 --- a/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelProtoCodec.java +++ b/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelProtoCodec.java @@ -163,9 +163,9 @@ public static long timeISO8601ToNanos(final String timeISO08601) { public static class OTelProtoDecoder { - public List parseExportTraceServiceRequest(final ExportTraceServiceRequest exportTraceServiceRequest) { + public List parseExportTraceServiceRequest(final ExportTraceServiceRequest exportTraceServiceRequest, final Instant timeReceived) { return exportTraceServiceRequest.getResourceSpansList().stream() - .flatMap(rs -> parseResourceSpans(rs).stream()).collect(Collectors.toList()); + .flatMap(rs -> parseResourceSpans(rs, timeReceived).stream()).collect(Collectors.toList()); } public Map splitExportTraceServiceRequestByTraceId(final ExportTraceServiceRequest exportTraceServiceRequest) { @@ -188,12 +188,12 @@ public Map splitExportTraceServiceRequestByTr return result; } - public List parseExportLogsServiceRequest(final ExportLogsServiceRequest exportLogsServiceRequest) { + public List parseExportLogsServiceRequest(final ExportLogsServiceRequest exportLogsServiceRequest, final Instant timeReceived) { return exportLogsServiceRequest.getResourceLogsList().stream() - .flatMap(rs -> parseResourceLogs(rs).stream()).collect(Collectors.toList()); + .flatMap(rs -> parseResourceLogs(rs, timeReceived).stream()).collect(Collectors.toList()); } - protected Collection parseResourceLogs(ResourceLogs rs) { + protected Collection parseResourceLogs(ResourceLogs rs, final Instant timeReceived) { final String serviceName = OTelProtoCodec.getServiceName(rs.getResource()).orElse(null); final Map resourceAttributes = OTelProtoCodec.getResourceAttributes(rs.getResource()); final String schemaUrl = rs.getSchemaUrl(); @@ -205,7 +205,8 @@ protected Collection parseResourceLogs(ResourceLogs rs) { serviceName, OTelProtoCodec.getInstrumentationLibraryAttributes(ils.getInstrumentationLibrary()), resourceAttributes, - schemaUrl)) + schemaUrl, + timeReceived)) .flatMap(Collection::stream); Stream mappedScopeListLogs = rs.getScopeLogsList() @@ -215,7 +216,8 @@ protected Collection parseResourceLogs(ResourceLogs rs) { serviceName, OTelProtoCodec.getInstrumentationScopeAttributes(sls.getScope()), resourceAttributes, - schemaUrl)) + schemaUrl, + timeReceived)) .flatMap(Collection::stream); return Stream.concat(mappedInstrumentationLibraryLogs, mappedScopeListLogs).collect(Collectors.toList()); @@ -260,26 +262,26 @@ protected Map splitResourceSpansByTraceId(final ResourceS return result; } - protected List parseResourceSpans(final ResourceSpans resourceSpans) { + protected List parseResourceSpans(final ResourceSpans resourceSpans, final Instant timeReceived) { final String serviceName = getServiceName(resourceSpans.getResource()).orElse(null); final Map resourceAttributes = getResourceAttributes(resourceSpans.getResource()); if (resourceSpans.getScopeSpansList().size() > 0) { - return parseScopeSpans(resourceSpans.getScopeSpansList(), serviceName, resourceAttributes); + return parseScopeSpans(resourceSpans.getScopeSpansList(), serviceName, resourceAttributes, timeReceived); } if (resourceSpans.getInstrumentationLibrarySpansList().size() > 0) { - return parseInstrumentationLibrarySpans(resourceSpans.getInstrumentationLibrarySpansList(), serviceName, resourceAttributes); + return parseInstrumentationLibrarySpans(resourceSpans.getInstrumentationLibrarySpansList(), serviceName, resourceAttributes, timeReceived); } LOG.debug("No spans found to parse from ResourceSpans object: {}", resourceSpans); return Collections.emptyList(); } - private List parseScopeSpans(final List scopeSpansList, final String serviceName, final Map resourceAttributes) { + private List parseScopeSpans(final List scopeSpansList, final String serviceName, final Map resourceAttributes, final Instant timeReceived) { return scopeSpansList.stream() .map(scopeSpans -> parseSpans(scopeSpans.getSpansList(), scopeSpans.getScope(), - OTelProtoCodec::getInstrumentationScopeAttributes, serviceName, resourceAttributes)) + OTelProtoCodec::getInstrumentationScopeAttributes, serviceName, resourceAttributes, timeReceived)) .flatMap(Collection::stream) .collect(Collectors.toList()); } @@ -305,11 +307,12 @@ private Map> splitScopeSpansByTraceId(final List parseInstrumentationLibrarySpans(final List instrumentationLibrarySpansList, - final String serviceName, final Map resourceAttributes) { + final String serviceName, final Map resourceAttributes, + final Instant timeReceived) { return instrumentationLibrarySpansList.stream() .map(instrumentationLibrarySpans -> parseSpans(instrumentationLibrarySpans.getSpansList(), instrumentationLibrarySpans.getInstrumentationLibrary(), this::getInstrumentationLibraryAttributes, - serviceName, resourceAttributes)) + serviceName, resourceAttributes, timeReceived)) .flatMap(Collection::stream) .collect(Collectors.toList()); } @@ -354,20 +357,22 @@ private Map> splitSpansByTrac private List parseSpans(final List spans, final T scope, final Function> scopeAttributesGetter, - final String serviceName, final Map resourceAttributes) { + final String serviceName, final Map resourceAttributes, + final Instant timeReceived) { return spans.stream() .map(span -> { final Map scopeAttributes = scopeAttributesGetter.apply(scope); - return parseSpan(span, scopeAttributes, serviceName, resourceAttributes); + return parseSpan(span, scopeAttributes, serviceName, resourceAttributes, timeReceived); }) .collect(Collectors.toList()); } protected List processLogsList(final List logsList, - final String serviceName, - final Map ils, - final Map resourceAttributes, - final String schemaUrl) { + final String serviceName, + final Map ils, + final Map resourceAttributes, + final String schemaUrl, + final Instant timeReceived) { return logsList.stream() .map(log -> JacksonOtelLog.builder() .withTime(OTelProtoCodec.convertUnixNanosToISO8601(log.getTimeUnixNano())) @@ -388,12 +393,13 @@ protected List processLogsList(final List logsList, .withSeverityText(log.getSeverityText()) .withDroppedAttributesCount(log.getDroppedAttributesCount()) .withBody(OTelProtoCodec.convertAnyValue(log.getBody())) + .withTimeReceived(timeReceived) .build()) .collect(Collectors.toList()); } protected Span parseSpan(final io.opentelemetry.proto.trace.v1.Span sp, final Map instrumentationScopeAttributes, - final String serviceName, final Map resourceAttributes) { + final String serviceName, final Map resourceAttributes, final Instant timeReceived) { return JacksonSpan.builder() .withSpanId(convertByteStringToString(sp.getSpanId())) .withTraceId(convertByteStringToString(sp.getTraceId())) @@ -420,6 +426,7 @@ protected Span parseSpan(final io.opentelemetry.proto.trace.v1.Span sp, final Ma .withTraceGroup(getTraceGroup(sp)) .withDurationInNanos(sp.getEndTimeUnixNano() - sp.getStartTimeUnixNano()) .withTraceGroupFields(getTraceGroupFields(sp)) + .withTimeReceived(timeReceived) .build(); } @@ -583,6 +590,7 @@ public Collection> parseExportMetricsServiceRequest( final ExportMetricsServiceRequest request, AtomicInteger droppedCounter, final Integer exponentialHistogramMaxAllowedScale, + final Instant timeReceived, final boolean calculateHistogramBuckets, final boolean calculateExponentialHistogramBuckets, final boolean flattenAttributes) { @@ -594,12 +602,12 @@ public Collection> parseExportMetricsServiceRequest( for (InstrumentationLibraryMetrics is : rs.getInstrumentationLibraryMetricsList()) { final Map ils = OTelProtoCodec.getInstrumentationLibraryAttributes(is.getInstrumentationLibrary()); - recordsOut.addAll(processMetricsList(is.getMetricsList(), serviceName, ils, resourceAttributes, schemaUrl, droppedCounter, exponentialHistogramMaxAllowedScale, calculateHistogramBuckets, calculateExponentialHistogramBuckets, flattenAttributes)); + recordsOut.addAll(processMetricsList(is.getMetricsList(), serviceName, ils, resourceAttributes, schemaUrl, droppedCounter, exponentialHistogramMaxAllowedScale, timeReceived, calculateHistogramBuckets, calculateExponentialHistogramBuckets, flattenAttributes)); } for (ScopeMetrics sm : rs.getScopeMetricsList()) { final Map ils = OTelProtoCodec.getInstrumentationScopeAttributes(sm.getScope()); - recordsOut.addAll(processMetricsList(sm.getMetricsList(), serviceName, ils, resourceAttributes, schemaUrl, droppedCounter, exponentialHistogramMaxAllowedScale, calculateHistogramBuckets, calculateExponentialHistogramBuckets, flattenAttributes)); + recordsOut.addAll(processMetricsList(sm.getMetricsList(), serviceName, ils, resourceAttributes, schemaUrl, droppedCounter, exponentialHistogramMaxAllowedScale, timeReceived, calculateHistogramBuckets, calculateExponentialHistogramBuckets, flattenAttributes)); } } return recordsOut; @@ -613,6 +621,7 @@ private List> processMetricsList( final String schemaUrl, AtomicInteger droppedCounter, final Integer exponentialHistogramMaxAllowedScale, + final Instant timeReceived, final boolean calculateHistogramBuckets, final boolean calculateExponentialHistogramBuckets, final boolean flattenAttributes) { @@ -620,15 +629,15 @@ private List> processMetricsList( for (io.opentelemetry.proto.metrics.v1.Metric metric : metricsList) { try { if (metric.hasGauge()) { - recordsOut.addAll(mapGauge(metric, serviceName, ils, resourceAttributes, schemaUrl, flattenAttributes)); + recordsOut.addAll(mapGauge(metric, serviceName, ils, resourceAttributes, schemaUrl, timeReceived, flattenAttributes)); } else if (metric.hasSum()) { - recordsOut.addAll(mapSum(metric, serviceName, ils, resourceAttributes, schemaUrl, flattenAttributes)); + recordsOut.addAll(mapSum(metric, serviceName, ils, resourceAttributes, schemaUrl, timeReceived, flattenAttributes)); } else if (metric.hasSummary()) { - recordsOut.addAll(mapSummary(metric, serviceName, ils, resourceAttributes, schemaUrl, flattenAttributes)); + recordsOut.addAll(mapSummary(metric, serviceName, ils, resourceAttributes, schemaUrl, timeReceived, flattenAttributes)); } else if (metric.hasHistogram()) { - recordsOut.addAll(mapHistogram(metric, serviceName, ils, resourceAttributes, schemaUrl, calculateHistogramBuckets, flattenAttributes)); + recordsOut.addAll(mapHistogram(metric, serviceName, ils, resourceAttributes, schemaUrl, timeReceived, calculateHistogramBuckets, flattenAttributes)); } else if (metric.hasExponentialHistogram()) { - recordsOut.addAll(mapExponentialHistogram(metric, serviceName, ils, resourceAttributes, schemaUrl, exponentialHistogramMaxAllowedScale, calculateExponentialHistogramBuckets, flattenAttributes)); + recordsOut.addAll(mapExponentialHistogram(metric, serviceName, ils, resourceAttributes, schemaUrl, exponentialHistogramMaxAllowedScale, timeReceived, calculateExponentialHistogramBuckets, flattenAttributes)); } } catch (Exception e) { LOG.warn("Error while processing metrics", e); @@ -644,6 +653,7 @@ private List> mapGauge( final Map ils, final Map resourceAttributes, final String schemaUrl, + final Instant timeReceived, final boolean flattenAttributes) { return metric.getGauge().getDataPointsList().stream() .map(dp -> JacksonGauge.builder() @@ -664,6 +674,7 @@ private List> mapGauge( .withSchemaUrl(schemaUrl) .withExemplars(OTelProtoCodec.convertExemplars(dp.getExemplarsList())) .withFlags(dp.getFlags()) + .withTimeReceived(timeReceived) .build(flattenAttributes)) .map(Record::new) .collect(Collectors.toList()); @@ -675,6 +686,7 @@ private List> mapSum( final Map ils, final Map resourceAttributes, final String schemaUrl, + final Instant timeReceived, final boolean flattenAttributes) { return metric.getSum().getDataPointsList().stream() .map(dp -> JacksonSum.builder() @@ -697,6 +709,7 @@ private List> mapSum( .withSchemaUrl(schemaUrl) .withExemplars(OTelProtoCodec.convertExemplars(dp.getExemplarsList())) .withFlags(dp.getFlags()) + .withTimeReceived(timeReceived) .build(flattenAttributes)) .map(Record::new) .collect(Collectors.toList()); @@ -708,6 +721,7 @@ private List> mapSummary( final Map ils, final Map resourceAttributes, final String schemaUrl, + final Instant timeReceived, final boolean flattenAttributes) { return metric.getSummary().getDataPointsList().stream() .map(dp -> JacksonSummary.builder() @@ -730,6 +744,7 @@ private List> mapSummary( )) .withSchemaUrl(schemaUrl) .withFlags(dp.getFlags()) + .withTimeReceived(timeReceived) .build(flattenAttributes)) .map(Record::new) .collect(Collectors.toList()); @@ -741,6 +756,7 @@ private List> mapHistogram( final Map ils, final Map resourceAttributes, final String schemaUrl, + final Instant timeReceived, final boolean calculateHistogramBuckets, final boolean flattenAttributes) { return metric.getHistogram().getDataPointsList().stream() @@ -768,6 +784,7 @@ private List> mapHistogram( )) .withSchemaUrl(schemaUrl) .withExemplars(OTelProtoCodec.convertExemplars(dp.getExemplarsList())) + .withTimeReceived(timeReceived) .withFlags(dp.getFlags()); if (calculateHistogramBuckets) { builder.withBuckets(OTelProtoCodec.createBuckets(dp.getBucketCountsList(), dp.getExplicitBoundsList())); @@ -787,6 +804,7 @@ private List> mapExponentialHistogram( final Map resourceAttributes, final String schemaUrl, final Integer exponentialHistogramMaxAllowedScale, + final Instant timeReceived, final boolean calculateExponentialHistogramBuckets, final boolean flattenAttributes) { return metric.getExponentialHistogram() @@ -827,6 +845,7 @@ private List> mapExponentialHistogram( )) .withSchemaUrl(schemaUrl) .withExemplars(OTelProtoCodec.convertExemplars(dp.getExemplarsList())) + .withTimeReceived(timeReceived) .withFlags(dp.getFlags()); if (calculateExponentialHistogramBuckets) { diff --git a/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelTraceDecoder.java b/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelTraceDecoder.java index 47c3fd03e9..302723021a 100644 --- a/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelTraceDecoder.java +++ b/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelTraceDecoder.java @@ -18,6 +18,7 @@ import java.io.IOException; import java.io.InputStream; import java.util.function.Consumer; +import java.time.Instant; public class OTelTraceDecoder implements ByteDecoder { @@ -25,11 +26,10 @@ public class OTelTraceDecoder implements ByteDecoder { public OTelTraceDecoder() { otelProtoDecoder = new OTelProtoCodec.OTelProtoDecoder(); } - public void parse(InputStream inputStream, Consumer> eventConsumer) throws IOException { + public void parse(InputStream inputStream, Instant timeReceivedMs, Consumer> eventConsumer) throws IOException { ExportTraceServiceRequest request = ExportTraceServiceRequest.parseFrom(inputStream); AtomicInteger droppedCounter = new AtomicInteger(0); - List spans = - otelProtoDecoder.parseExportTraceServiceRequest(request); + List spans = otelProtoDecoder.parseExportTraceServiceRequest(request, timeReceivedMs); for (Span span: spans) { eventConsumer.accept(new Record<>(span)); } diff --git a/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OTelLogsDecoderTest.java b/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OTelLogsDecoderTest.java index 8fef72b2fa..7c64d88a20 100644 --- a/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OTelLogsDecoderTest.java +++ b/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OTelLogsDecoderTest.java @@ -12,6 +12,7 @@ import java.io.InputStreamReader; import java.io.IOException; +import java.time.Instant; import java.util.Objects; import java.util.Map; import com.google.protobuf.util.JsonFormat; @@ -65,7 +66,7 @@ private void validateLog(OpenTelemetryLog logRecord) { public void testParse() throws Exception { final ExportLogsServiceRequest request = buildExportLogsServiceRequestFromJsonFile(TEST_REQUEST_LOGS_FILE); InputStream inputStream = new ByteArrayInputStream((byte[])request.toByteArray()); - createObjectUnderTest().parse(inputStream, (record) -> { + createObjectUnderTest().parse(inputStream, Instant.now(), (record) -> { validateLog((OpenTelemetryLog)record.getData()); }); diff --git a/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OTelMetricsDecoderTest.java b/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OTelMetricsDecoderTest.java index c6bdd9cc89..4dfb58dfa6 100644 --- a/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OTelMetricsDecoderTest.java +++ b/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OTelMetricsDecoderTest.java @@ -27,6 +27,7 @@ import org.opensearch.dataprepper.model.metric.JacksonSum; import org.opensearch.dataprepper.model.metric.JacksonHistogram; +import java.time.Instant; import java.util.Collections; import java.util.List; import java.util.Objects; @@ -87,7 +88,7 @@ private void validateMetric(Event event) { public void testParse() throws Exception { final ExportMetricsServiceRequest request = buildExportMetricsServiceRequestFromJsonFile(TEST_REQUEST_METRICS_FILE); InputStream inputStream = new ByteArrayInputStream((byte[])request.toByteArray()); - createObjectUnderTest().parse(inputStream, (record) -> { + createObjectUnderTest().parse(inputStream, Instant.now(), (record) -> { validateMetric((Event)record.getData()); }); diff --git a/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OTelProtoCodecTest.java b/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OTelProtoCodecTest.java index 0b75ea6f46..b2e42c6c20 100644 --- a/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OTelProtoCodecTest.java +++ b/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OTelProtoCodecTest.java @@ -211,28 +211,28 @@ public void testSplitExportTraceServiceRequestWithMultipleTraces() throws Except @Test public void testParseExportTraceServiceRequest() throws IOException { final ExportTraceServiceRequest exportTraceServiceRequest = buildExportTraceServiceRequestFromJsonFile(TEST_REQUEST_TRACE_JSON_FILE); - final List spans = decoderUnderTest.parseExportTraceServiceRequest(exportTraceServiceRequest); + final List spans = decoderUnderTest.parseExportTraceServiceRequest(exportTraceServiceRequest, Instant.now()); validateSpans(spans); } @Test public void testParseExportTraceServiceRequest_InstrumentationLibrarySpans() throws IOException { final ExportTraceServiceRequest exportTraceServiceRequest = buildExportTraceServiceRequestFromJsonFile(TEST_REQUEST_INSTRUMENTATION_LIBRARY_TRACE_JSON_FILE); - final List spans = decoderUnderTest.parseExportTraceServiceRequest(exportTraceServiceRequest); + final List spans = decoderUnderTest.parseExportTraceServiceRequest(exportTraceServiceRequest, Instant.now()); validateSpans(spans); } @Test public void testParseExportTraceServiceRequest_ScopeSpansTakesPrecedenceOverInstrumentationLibrarySpans() throws IOException { final ExportTraceServiceRequest exportTraceServiceRequest = buildExportTraceServiceRequestFromJsonFile(TEST_REQUEST_BOTH_SPAN_TYPES_JSON_FILE); - final List spans = decoderUnderTest.parseExportTraceServiceRequest(exportTraceServiceRequest); + final List spans = decoderUnderTest.parseExportTraceServiceRequest(exportTraceServiceRequest, Instant.now()); validateSpans(spans); } @Test public void testParseExportTraceServiceRequest_NoSpans() throws IOException { final ExportTraceServiceRequest exportTraceServiceRequest = buildExportTraceServiceRequestFromJsonFile(TEST_REQUEST_NO_SPANS_JSON_FILE); - final List spans = decoderUnderTest.parseExportTraceServiceRequest(exportTraceServiceRequest); + final List spans = decoderUnderTest.parseExportTraceServiceRequest(exportTraceServiceRequest, Instant.now()); assertThat(spans.size(), is(equalTo(0))); } @@ -492,7 +492,7 @@ public void testTraceGroupFields() { @Test public void testParseExportLogsServiceRequest_ScopedLogs() throws IOException { final ExportLogsServiceRequest exportLogsServiceRequest = buildExportLogsServiceRequestFromJsonFile(TEST_REQUEST_LOGS_JSON_FILE); - List logs = decoderUnderTest.parseExportLogsServiceRequest(exportLogsServiceRequest); + List logs = decoderUnderTest.parseExportLogsServiceRequest(exportLogsServiceRequest, Instant.now()); assertThat(logs.size() , is(equalTo(1))); validateLog(logs.get(0)); @@ -501,7 +501,7 @@ public void testParseExportLogsServiceRequest_ScopedLogs() throws IOException { @Test public void testParseExportLogsServiceRequest_InstrumentationLibraryLogs() throws IOException { final ExportLogsServiceRequest exportLogsServiceRequest = buildExportLogsServiceRequestFromJsonFile(TEST_REQUEST_LOGS_IS_JSON_FILE); - List logs = decoderUnderTest.parseExportLogsServiceRequest(exportLogsServiceRequest); + List logs = decoderUnderTest.parseExportLogsServiceRequest(exportLogsServiceRequest, Instant.now()); assertThat(logs.size() , is(equalTo(1))); validateLog(logs.get(0)); @@ -527,7 +527,7 @@ private void validateLog(OpenTelemetryLog logRecord) { @Test public void testParseExportLogsServiceRequest_InstrumentationLibrarySpans() throws IOException { final ExportTraceServiceRequest exportTraceServiceRequest = buildExportTraceServiceRequestFromJsonFile(TEST_REQUEST_INSTRUMENTATION_LIBRARY_TRACE_JSON_FILE); - final List spans = decoderUnderTest.parseExportTraceServiceRequest(exportTraceServiceRequest); + final List spans = decoderUnderTest.parseExportTraceServiceRequest(exportTraceServiceRequest, Instant.now()); validateSpans(spans); } @@ -535,7 +535,7 @@ public void testParseExportLogsServiceRequest_InstrumentationLibrarySpans() thro public void testParseExportMetricsServiceRequest_Guage() throws IOException { final ExportMetricsServiceRequest exportMetricsServiceRequest = buildExportMetricsServiceRequestFromJsonFile(TEST_REQUEST_GAUGE_METRICS_JSON_FILE); AtomicInteger droppedCount = new AtomicInteger(0); - final Collection> metrics = decoderUnderTest.parseExportMetricsServiceRequest(exportMetricsServiceRequest, droppedCount, 10, true, true, true); + final Collection> metrics = decoderUnderTest.parseExportMetricsServiceRequest(exportMetricsServiceRequest, droppedCount, 10, Instant.now(), true, true, true); validateGaugeMetricRequest(metrics); } @@ -544,7 +544,7 @@ public void testParseExportMetricsServiceRequest_Guage() throws IOException { public void testParseExportMetricsServiceRequest_Sum() throws IOException { final ExportMetricsServiceRequest exportMetricsServiceRequest = buildExportMetricsServiceRequestFromJsonFile(TEST_REQUEST_SUM_METRICS_JSON_FILE); AtomicInteger droppedCount = new AtomicInteger(0); - final Collection> metrics = decoderUnderTest.parseExportMetricsServiceRequest(exportMetricsServiceRequest, droppedCount, 10, true, true, true); + final Collection> metrics = decoderUnderTest.parseExportMetricsServiceRequest(exportMetricsServiceRequest, droppedCount, 10, Instant.now(), true, true, true); validateSumMetricRequest(metrics); } @@ -552,7 +552,7 @@ public void testParseExportMetricsServiceRequest_Sum() throws IOException { public void testParseExportMetricsServiceRequest_Histogram() throws IOException { final ExportMetricsServiceRequest exportMetricsServiceRequest = buildExportMetricsServiceRequestFromJsonFile(TEST_REQUEST_HISTOGRAM_METRICS_JSON_FILE); AtomicInteger droppedCount = new AtomicInteger(0); - final Collection> metrics = decoderUnderTest.parseExportMetricsServiceRequest(exportMetricsServiceRequest, droppedCount, 10, true, true, true); + final Collection> metrics = decoderUnderTest.parseExportMetricsServiceRequest(exportMetricsServiceRequest, droppedCount, 10, Instant.now(), true, true, true); validateHistogramMetricRequest(metrics); } @@ -560,7 +560,7 @@ public void testParseExportMetricsServiceRequest_Histogram() throws IOException public void testParseExportMetricsServiceRequest_Histogram_WithNoExplicitBounds() throws IOException { final ExportMetricsServiceRequest exportMetricsServiceRequest = buildExportMetricsServiceRequestFromJsonFile(TEST_REQUEST_HISTOGRAM_METRICS_NO_EXPLICIT_BOUNDS_JSON_FILE); AtomicInteger droppedCount = new AtomicInteger(0); - final Collection> metrics = decoderUnderTest.parseExportMetricsServiceRequest(exportMetricsServiceRequest, droppedCount, 10, true, true, true); + final Collection> metrics = decoderUnderTest.parseExportMetricsServiceRequest(exportMetricsServiceRequest, droppedCount, 10, Instant.now(), true, true, true); validateHistogramMetricRequestNoExplicitBounds(metrics); } @@ -944,13 +944,13 @@ public void testTimeCodec() { @Test public void testOTelProtoCodecConsistency() throws IOException, DecoderException { final ExportTraceServiceRequest request = buildExportTraceServiceRequestFromJsonFile(TEST_REQUEST_TRACE_JSON_FILE); - final List spansFirstDec = decoderUnderTest.parseExportTraceServiceRequest(request); + final List spansFirstDec = decoderUnderTest.parseExportTraceServiceRequest(request, Instant.now()); final List resourceSpansList = new ArrayList<>(); for (final Span span : spansFirstDec) { resourceSpansList.add(encoderUnderTest.convertToResourceSpans(span)); } final List spansSecondDec = resourceSpansList.stream() - .flatMap(rs -> decoderUnderTest.parseResourceSpans(rs).stream()).collect(Collectors.toList()); + .flatMap(rs -> decoderUnderTest.parseResourceSpans(rs, Instant.now()).stream()).collect(Collectors.toList()); assertThat(spansFirstDec.size(), equalTo(spansSecondDec.size())); for (int i = 0; i < spansFirstDec.size(); i++) { assertThat(spansFirstDec.get(i).toJsonString(), equalTo(spansSecondDec.get(i).toJsonString())); diff --git a/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OTelTraceDecoderTest.java b/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OTelTraceDecoderTest.java index a6cc6d122b..0e976a14e2 100644 --- a/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OTelTraceDecoderTest.java +++ b/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OTelTraceDecoderTest.java @@ -12,6 +12,7 @@ import java.io.InputStreamReader; import java.io.IOException; +import java.time.Instant; import java.util.Objects; import com.google.protobuf.util.JsonFormat; import org.opensearch.dataprepper.model.trace.Span; @@ -66,7 +67,7 @@ private void validateSpan(Span span) { public void testParse() throws Exception { final ExportTraceServiceRequest request = buildExportTraceServiceRequestFromJsonFile(TEST_REQUEST_TRACES_FILE); InputStream inputStream = new ByteArrayInputStream((byte[])request.toByteArray()); - createObjectUnderTest().parse(inputStream, (record) -> { + createObjectUnderTest().parse(inputStream, Instant.now(), (record) -> { validateSpan((Span)record.getData()); }); diff --git a/data-prepper-plugins/otel-trace-source/src/main/java/org/opensearch/dataprepper/plugins/source/oteltrace/OTelTraceGrpcService.java b/data-prepper-plugins/otel-trace-source/src/main/java/org/opensearch/dataprepper/plugins/source/oteltrace/OTelTraceGrpcService.java index 5ac38d8e4b..a5799605ff 100644 --- a/data-prepper-plugins/otel-trace-source/src/main/java/org/opensearch/dataprepper/plugins/source/oteltrace/OTelTraceGrpcService.java +++ b/data-prepper-plugins/otel-trace-source/src/main/java/org/opensearch/dataprepper/plugins/source/oteltrace/OTelTraceGrpcService.java @@ -26,6 +26,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.time.Instant; import java.util.Collection; import java.util.List; import java.util.Map; @@ -88,7 +89,7 @@ private void processRequest(final ExportTraceServiceRequest request, final Strea final Collection spans; try { - spans = oTelProtoDecoder.parseExportTraceServiceRequest(request); + spans = oTelProtoDecoder.parseExportTraceServiceRequest(request, Instant.now()); } catch (final Exception e) { LOG.warn(DataPrepperMarkers.SENSITIVE, "Failed to parse request with error '{}'. Request body: {}.", e.getMessage(), request); throw new BadRequestException(e.getMessage(), e); diff --git a/data-prepper-plugins/otel-trace-source/src/test/java/org/opensearch/dataprepper/plugins/source/oteltrace/OTelTraceGrpcServiceTest.java b/data-prepper-plugins/otel-trace-source/src/test/java/org/opensearch/dataprepper/plugins/source/oteltrace/OTelTraceGrpcServiceTest.java index d6aa5503c5..ff145a61a9 100644 --- a/data-prepper-plugins/otel-trace-source/src/test/java/org/opensearch/dataprepper/plugins/source/oteltrace/OTelTraceGrpcServiceTest.java +++ b/data-prepper-plugins/otel-trace-source/src/test/java/org/opensearch/dataprepper/plugins/source/oteltrace/OTelTraceGrpcServiceTest.java @@ -36,6 +36,7 @@ import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoCodec; import java.io.IOException; +import java.time.Instant; import java.util.Collection; import java.util.Collections; import java.util.List; @@ -198,7 +199,7 @@ public void export_BufferTimeout_responseObserverOnError() throws Exception { public void export_BadRequest_responseObserverOnError() throws Exception { final String testMessage = "test message"; final RuntimeException testException = new RuntimeException(testMessage); - when(mockOTelProtoDecoder.parseExportTraceServiceRequest(any())).thenThrow(testException); + when(mockOTelProtoDecoder.parseExportTraceServiceRequest(any(), any(Instant.class))).thenThrow(testException); objectUnderTest = generateOTelTraceGrpcService(mockOTelProtoDecoder); try (MockedStatic mockedStatic = mockStatic(ServiceRequestContext.class)) { diff --git a/data-prepper-plugins/parse-json-processor/build.gradle b/data-prepper-plugins/parse-json-processor/build.gradle index 7d4842850a..8ddb6af3c6 100644 --- a/data-prepper-plugins/parse-json-processor/build.gradle +++ b/data-prepper-plugins/parse-json-processor/build.gradle @@ -7,15 +7,12 @@ plugins { id 'java' } -repositories { - mavenCentral() -} - dependencies { implementation project(':data-prepper-api') implementation project(':data-prepper-plugins:common') implementation 'com.fasterxml.jackson.core:jackson-databind' implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-ion' + implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-xml' implementation 'org.apache.parquet:parquet-common:1.13.1' testImplementation project(':data-prepper-test-common') } diff --git a/data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/codec/json/JsonInputCodec.java b/data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/codec/json/JsonInputCodec.java index 724787879f..6222682e2a 100644 --- a/data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/codec/json/JsonInputCodec.java +++ b/data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/codec/json/JsonInputCodec.java @@ -8,11 +8,20 @@ import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; import org.opensearch.dataprepper.model.codec.InputCodec; import org.opensearch.dataprepper.model.codec.JsonDecoder; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.record.Record; + +import java.io.IOException; +import java.io.InputStream; +import java.util.function.Consumer; /** * An implementation of {@link InputCodec} which parses JSON Objects for arrays. */ @DataPrepperPlugin(name = "json", pluginType = InputCodec.class) public class JsonInputCodec extends JsonDecoder implements InputCodec { + public void parse(InputStream inputStream, Consumer> eventConsumer) throws IOException { + parse(inputStream, null, eventConsumer); + } } diff --git a/data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/parse/AbstractParseProcessor.java b/data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/parse/AbstractParseProcessor.java index b113534983..a2b984d070 100644 --- a/data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/parse/AbstractParseProcessor.java +++ b/data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/parse/AbstractParseProcessor.java @@ -72,7 +72,7 @@ public Collection> doExecute(final Collection> recor } final String message = event.get(source, String.class); - if (Objects.isNull(message)) { + if (Objects.isNull(message) && !doUsePointer) { continue; } diff --git a/data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/parse/xml/ParseXmlProcessor.java b/data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/parse/xml/ParseXmlProcessor.java new file mode 100644 index 0000000000..840515afb2 --- /dev/null +++ b/data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/parse/xml/ParseXmlProcessor.java @@ -0,0 +1,46 @@ +package org.opensearch.dataprepper.plugins.processor.parse.xml; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.dataformat.xml.XmlMapper; +import org.opensearch.dataprepper.expression.ExpressionEvaluator; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; +import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.processor.Processor; +import org.opensearch.dataprepper.plugins.processor.parse.AbstractParseProcessor; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.HashMap; +import java.util.Optional; + +import static org.opensearch.dataprepper.logging.DataPrepperMarkers.EVENT; + +@DataPrepperPlugin(name = "parse_xml", pluginType =Processor.class, pluginConfigurationType =ParseXmlProcessorConfig.class) +public class ParseXmlProcessor extends AbstractParseProcessor { + private static final Logger LOG = LoggerFactory.getLogger(ParseXmlProcessor.class); + + private final XmlMapper xmlMapper = new XmlMapper(); + + @DataPrepperPluginConstructor + public ParseXmlProcessor(final PluginMetrics pluginMetrics, + final ParseXmlProcessorConfig parseXmlProcessorConfig, + final ExpressionEvaluator expressionEvaluator) { + super(pluginMetrics, parseXmlProcessorConfig, expressionEvaluator); + } + + @Override + protected Optional> readValue(final String message, final Event context) { + try { + return Optional.of(xmlMapper.readValue(message, new TypeReference<>() {})); + } catch (JsonProcessingException e) { + LOG.error(EVENT, "An exception occurred due to invalid XML while reading event [{}]", context, e); + return Optional.empty(); + } catch (Exception e) { + LOG.error(EVENT, "An exception occurred while using the parse_xml processor on Event [{}]", context, e); + return Optional.empty(); + } + } +} diff --git a/data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/parse/xml/ParseXmlProcessorConfig.java b/data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/parse/xml/ParseXmlProcessorConfig.java new file mode 100644 index 0000000000..df4fabc397 --- /dev/null +++ b/data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/parse/xml/ParseXmlProcessorConfig.java @@ -0,0 +1,70 @@ +package org.opensearch.dataprepper.plugins.processor.parse.xml; + +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.constraints.AssertTrue; +import jakarta.validation.constraints.NotBlank; +import org.opensearch.dataprepper.plugins.processor.parse.CommonParseConfig; + +import java.util.List; +import java.util.Objects; + +public class ParseXmlProcessorConfig implements CommonParseConfig { + static final String DEFAULT_SOURCE = "message"; + + @NotBlank + @JsonProperty("source") + private String source = DEFAULT_SOURCE; + + @JsonProperty("destination") + private String destination; + + @JsonProperty("pointer") + private String pointer; + + @JsonProperty("parse_when") + private String parseWhen; + + @JsonProperty("tags_on_failure") + private List tagsOnFailure; + + @JsonProperty("overwrite_if_destination_exists") + private boolean overwriteIfDestinationExists = true; + + @Override + public String getSource() { + return source; + } + + @Override + public String getDestination() { + return destination; + } + + @Override + public String getPointer() { + return pointer; + } + + @Override + public List getTagsOnFailure() { + return tagsOnFailure; + } + + @Override + public String getParseWhen() { + return parseWhen; + } + + @Override + public boolean getOverwriteIfDestinationExists() { + return overwriteIfDestinationExists; + } + + @AssertTrue(message = "destination cannot be empty, whitespace, or a front slash (/)") + boolean isValidDestination() { + if (Objects.isNull(destination)) return true; + + final String trimmedDestination = destination.trim(); + return trimmedDestination.length() != 0 && !(trimmedDestination.equals("/")); + } +} diff --git a/data-prepper-plugins/parse-json-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/parse/xml/ParseXmlProcessorConfigTest.java b/data-prepper-plugins/parse-json-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/parse/xml/ParseXmlProcessorConfigTest.java new file mode 100644 index 0000000000..d5e7e1ec43 --- /dev/null +++ b/data-prepper-plugins/parse-json-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/parse/xml/ParseXmlProcessorConfigTest.java @@ -0,0 +1,57 @@ +package org.opensearch.dataprepper.plugins.processor.parse.xml; + +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.opensearch.dataprepper.test.helper.ReflectivelySetField.setField; + +public class ParseXmlProcessorConfigTest { + + private ParseXmlProcessorConfig createObjectUnderTest() { + return new ParseXmlProcessorConfig(); + } + + @Test + public void test_when_defaultParseXmlProcessorConfig_then_returns_default_values() { + final ParseXmlProcessorConfig objectUnderTest = createObjectUnderTest(); + + assertThat(objectUnderTest.getSource(), equalTo(ParseXmlProcessorConfig.DEFAULT_SOURCE)); + assertThat(objectUnderTest.getDestination(), equalTo(null)); + assertThat(objectUnderTest.getPointer(), equalTo(null)); + assertThat(objectUnderTest.getTagsOnFailure(), equalTo(null)); + assertThat(objectUnderTest.getOverwriteIfDestinationExists(), equalTo(true)); + } + + @Nested + class Validation { + final ParseXmlProcessorConfig config = createObjectUnderTest(); + + @Test + void test_when_destinationIsWhiteSpaceOrFrontSlash_then_isValidDestinationFalse() + throws NoSuchFieldException, IllegalAccessException { + setField(ParseXmlProcessorConfig.class, config, "destination", "good destination"); + + assertThat(config.isValidDestination(), equalTo(true)); + + setField(ParseXmlProcessorConfig.class, config, "destination", ""); + + assertThat(config.isValidDestination(), equalTo(false)); + + setField(ParseXmlProcessorConfig.class, config, "destination", " "); + + assertThat(config.isValidDestination(), equalTo(false)); + + setField(ParseXmlProcessorConfig.class, config, "destination", " / "); + + assertThat(config.isValidDestination(), equalTo(false)); + List tagsList = List.of("tag1", "tag2"); + setField(ParseXmlProcessorConfig.class, config, "tagsOnFailure", tagsList); + + assertThat(config.getTagsOnFailure(), equalTo(tagsList)); + } + } +} diff --git a/data-prepper-plugins/parse-json-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/parse/xml/ParseXmlProcessorTest.java b/data-prepper-plugins/parse-json-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/parse/xml/ParseXmlProcessorTest.java new file mode 100644 index 0000000000..51de35ca70 --- /dev/null +++ b/data-prepper-plugins/parse-json-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/parse/xml/ParseXmlProcessorTest.java @@ -0,0 +1,96 @@ +package org.opensearch.dataprepper.plugins.processor.parse.xml; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.expression.ExpressionEvaluator; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.event.JacksonEvent; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.plugins.processor.parse.AbstractParseProcessor; + +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.plugins.processor.parse.xml.ParseXmlProcessorConfig.DEFAULT_SOURCE; + + +@ExtendWith(MockitoExtension.class) +public class ParseXmlProcessorTest { + + @Mock + private ParseXmlProcessorConfig processorConfig; + + @Mock + private PluginMetrics pluginMetrics; + + @Mock + private ExpressionEvaluator expressionEvaluator; + + private AbstractParseProcessor parseXmlProcessor; + + @BeforeEach + public void setup() { + when(processorConfig.getSource()).thenReturn(DEFAULT_SOURCE); + when(processorConfig.getParseWhen()).thenReturn(null); + when(processorConfig.getOverwriteIfDestinationExists()).thenReturn(true); + } + + protected AbstractParseProcessor createObjectUnderTest() { + return new ParseXmlProcessor(pluginMetrics, processorConfig, expressionEvaluator); + } + + @Test + void test_when_using_xml_features_then_processorParsesCorrectly() { + parseXmlProcessor = createObjectUnderTest(); + + final String serializedMessage = "John Doe30"; + final Event parsedEvent = createAndParseMessageEvent(serializedMessage); + + assertThat(parsedEvent.get("name", String.class), equalTo("John Doe")); + assertThat(parsedEvent.get("age", String.class), equalTo("30")); + } + + @Test + void test_when_using_invalid_xml_tags_correctly() { + + final String tagOnFailure = UUID.randomUUID().toString(); + when(processorConfig.getTagsOnFailure()).thenReturn(List.of(tagOnFailure)); + + parseXmlProcessor = createObjectUnderTest(); + + final String serializedMessage = "invalidXml"; + final Event parsedEvent = createAndParseMessageEvent(serializedMessage); + + assertThat(parsedEvent.getMetadata().hasTags(List.of(tagOnFailure)), equalTo(true)); + } + + private Event createAndParseMessageEvent(final String message) { + final Record eventUnderTest = createMessageEvent(message); + final List> editedEvents = (List>) parseXmlProcessor.doExecute( + Collections.singletonList(eventUnderTest)); + return editedEvents.get(0).getData(); + } + + private Record createMessageEvent(final String message) { + final Map eventData = new HashMap<>(); + eventData.put(processorConfig.getSource(), message); + return buildRecordWithEvent(eventData); + } + + private Record buildRecordWithEvent(final Map data) { + return new Record<>(JacksonEvent.builder() + .withData(data) + .withEventType("event") + .build()); + } +} diff --git a/data-prepper-plugins/rss-source/build.gradle b/data-prepper-plugins/rss-source/build.gradle index df191a355e..eaf5c085b3 100644 --- a/data-prepper-plugins/rss-source/build.gradle +++ b/data-prepper-plugins/rss-source/build.gradle @@ -7,10 +7,6 @@ plugins { id 'java' } -repositories { - mavenCentral() -} - dependencies { implementation project(':data-prepper-api') implementation 'io.micrometer:micrometer-core' diff --git a/data-prepper-plugins/s3-sink/build.gradle b/data-prepper-plugins/s3-sink/build.gradle index bcf8dde821..8a792174de 100644 --- a/data-prepper-plugins/s3-sink/build.gradle +++ b/data-prepper-plugins/s3-sink/build.gradle @@ -12,7 +12,7 @@ dependencies { implementation 'com.fasterxml.jackson.core:jackson-databind' implementation libs.commons.compress implementation 'joda-time:joda-time:2.12.5' - implementation 'org.hibernate.validator:hibernate-validator:7.0.5.Final' + implementation 'org.hibernate.validator:hibernate-validator:8.0.1.Final' implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-csv' implementation 'software.amazon.awssdk:s3' implementation 'software.amazon.awssdk:sts' diff --git a/data-prepper-plugins/s3-sink/src/test/java/org/opensearch/dataprepper/plugins/codec/parquet/ParquetOutputCodecTest.java b/data-prepper-plugins/s3-sink/src/test/java/org/opensearch/dataprepper/plugins/codec/parquet/ParquetOutputCodecTest.java index 35b3128431..b441a7a6e3 100644 --- a/data-prepper-plugins/s3-sink/src/test/java/org/opensearch/dataprepper/plugins/codec/parquet/ParquetOutputCodecTest.java +++ b/data-prepper-plugins/s3-sink/src/test/java/org/opensearch/dataprepper/plugins/codec/parquet/ParquetOutputCodecTest.java @@ -97,7 +97,6 @@ void constructor_throws_if_schema_is_invalid() { RuntimeException actualException = assertThrows(RuntimeException.class, this::createObjectUnderTest); assertThat(actualException.getMessage(), notNullValue()); - assertThat(actualException.getMessage(), containsString(invalidSchema)); assertThat(actualException.getMessage(), containsString("was expecting comma")); } diff --git a/data-prepper-plugins/s3-source/README.md b/data-prepper-plugins/s3-source/README.md index 08f5437841..682d64cb6b 100644 --- a/data-prepper-plugins/s3-source/README.md +++ b/data-prepper-plugins/s3-source/README.md @@ -1,240 +1,9 @@ # S3 Source -This source allows Data Prepper to use S3 as a source. It uses SQS for notifications -of which S3 objects are new and loads those new objects to parse out events. -It supports scan pipeline to scan the data from s3 buckets and loads those new objects to parse out events. +This source ingests data into Data Prepper from [Amazon S3](https://aws.amazon.com/s3/). -## Basic Usage +See the [`s3` source documentation](https://opensearch.org/docs/latest/data-prepper/pipelines/configuration/sources/s3/) for details on usage. -This source requires an SQS queue which receives -[S3 Event Notifications](https://docs.aws.amazon.com/AmazonS3/latest/userguide/NotificationHowTo.html). -The S3 Source will load S3 objects that have Event notifications for Create events. -A user-specified codec parses the S3 Object and creates Events from them. - -Currently, there are three codecs: - -* `newline` - Parses files where each single line is a log event. -* `json` - Parses the file for a JSON array. Each object in the JSON array is a log event. -* `csv` - Parses a character separated file. Each line of data is a log event. - - - -The `compression` property defines how to handle compressed S3 objects. It has the following options. - -* `none` - The file is not compressed. -* `gzip` - Apply GZip de-compression on the S3 object. -* `automatic` - Attempts to automatically determine the compression. If the S3 object key name ends in`.gz`, then perform `gzip` compression. Otherwise, it is treated as `none`. - -### Example: Un-Compressed Logs - -The following configuration shows a minimum configuration for reading newline-delimited logs which -are not compressed. - -``` -source: - s3: - notification_type: "sqs" - codec: - newline: - compression: none - sqs: - queue_url: "https://sqs.us-east-1.amazonaws.com/123456789012/MyQueue" - aws: - region: "us-east-1" - sts_role_arn: "arn:aws:iam::123456789012:role/Data-Prepper" -``` - -The following configuration shows a minimum configuration for reading content using S3 select service and Scanning from S3 bucket which -are not compressed. - -``` -source-pipeline: - source: - s3: - notification_type: sqs - compression: none - codec: - newline: - s3_select: - expression: "select * from s3object s LIMIT 10000" - expression_type: SQL - input_serialization: csv - compression_type: none - csv: - file_header_info: use - quote_escape: - comments: - json: - type: document - sqs: - queue_url: "https://sqs.us-east-1.amazonaws.com/123456789012/MyQueue" - aws: - region: "us-east-1" - sts_role_arn: "arn:aws:iam::123456789012:role/Data-Prepper" - scan: - start_time: now - end_time: 2023-12-31T11:59:59 - buckets: - - bucket: - name: my-bucket-1 - filter: - include_prefix: - - bucket2/ - exclude_suffix: - - .jpeg - - .png -``` - -## Configuration Options - -All Duration values are a string that represents a duration. They support ISO_8601 notation string ("PT20.345S", "PT15M", etc.) as well as simple notation Strings for seconds ("60s") and milliseconds ("1500ms"). - -* `s3_select` : S3 Select Configuration. See [S3 Select Configuration](#s3_select_configuration) for details - -* `notification_type` (Optional) : Must be `sqs`. - -* `notification_source` (Optional): Provide how the notifications are generated. Must be either `s3` or `eventbridge`. Defaults to `s3`. `s3` represents notifications that are sent S3 to SQS directly or fanout pattern from S3 to SNS to SQS. `eventbridge` represents notifications which are sent from S3 to EventBridge to SQS. Only `Object Created` events are supported. - -* `compression` (Optional) : The compression algorithm to apply. May be one of: `none`, `gzip`, or `automatic`. Defaults to `none`. - -* `codec` (Required) : The codec to apply. Must be either `newline`, `csv` or `json`. - -* `sqs` (Optional) : The SQS configuration. See [SQS Configuration](#sqs_configuration) for details. - -* `scan` (Optional): S3 Scan Configuration. See [S3 Scan Configuration](#s3_scan_configuration) for details - -* `aws` (Optional) : AWS configurations. See [AWS Configuration](#aws_configuration) for details. - -* `acknowledgments` (Optional) : Enables End-to-end acknowledgments. If set to `true`, sqs message is deleted only after all events from the sqs message are successfully acknowledged by all sinks. Default value `false`. - -* `on_error` (Optional) : Determines how to handle errors in SQS. Can be either `retain_messages` or `delete_messages`. If `retain_messages`, then Data Prepper will leave the message in the SQS queue and try again. This is recommended for dead-letter queues. If `delete_messages`, then Data Prepper will delete failed messages. Defaults to `retain_messages`. - -* `buffer_timeout` (Optional) : Duration - The timeout for writing events to the Data Prepper buffer. Any Events which the S3 Source cannot write to the Buffer in this time will be discarded. Defaults to 10 seconds. - -* `records_to_accumulate` (Optional) : The number of messages to write to accumulate before writing to the Buffer. Defaults to 100. - -* `metadata_root_key` (Optional) : String - Sets the base key for adding S3 metadata to each Event. The metadata includes the `key` and `bucket` for each S3 object. Defaults to `s3/`. - -* `disable_bucket_ownership_validation` (Optional) : Boolean - If set to true, then the S3 Source will not attempt to validate that the bucket is owned by the expected account. The only expected account is the same account which owns the SQS queue. Defaults to `false`. - -* `delete_s3_objects_on_read` (Optional) : Boolean - If set to true, then the S3 Scan will attempt to delete S3 objects after all the events from the S3 object are successfully acknowledged by all sinks. `acknowledgments` should be enabled for deleting S3 objects. Defaults to `false`. - -### S3 Select Configuration - -* `expression` (Required if s3_select enabled) : Provide s3 select query to process the data using S3 select for the particular bucket. - -* `expression_type` (Optional if s3_select enabled) : Provide s3 select query type to process the data using S3 select for the particular bucket. - -* `compression_type` (Optional if s3_select enabled) : The compression algorithm to apply. May be one of: `none`, `gzip`. Defaults to `none`. - -* `input_serialization` (Required if s3_select enabled) : Provide the s3 select file format (csv/json/Apache Parquet) Amazon S3 uses this format to parse object data into records and returns only records that match the specified SQL expression. You must also specify the data serialization format for the response. - -* `csv` (Optional) : Provide the csv configuration to process the csv data. - - * `file_header_info` (Required if csv block is enabled) : Provide CSV Header example : `use` , `none` , `ignore`. Default is `use`. - - * `quote_escape` (Optional) : Provide quote_escape attribute example : `,` , `.`. - - * `comments` (Optional) : Provide comments attribute example : `#`. Default is `#`. - -* `json` (Optional) : Provide the json configuration to process the json data. - - * `type` (Optional) : Provide the type attribute to process the json type data example: `Lines` , `Document` Default is `Document`. - -### SQS Configuration - -* `queue_url` (Required) : The SQS queue URL of the queue to read from. -* `maximum_messages` (Optional) : Duration - The maximum number of messages to read from the queue in any request to the SQS queue. Defaults to 10. -* `visibility_timeout` (Optional) : Duration - The visibility timeout to apply to messages read from the SQS queue. This should be set to the amount of time that Data Prepper may take to read all the S3 objects in a batch. Defaults to 30 seconds. -* `wait_time` (Optional) : Duration - The time to wait for long-polling on the SQS API. Defaults to 20 seconds. -* `poll_delay` (Optional) : Duration - A delay to place between reading and processing a batch of SQS messages and making a subsequent request. Defaults to 0 seconds. - -### S3 Scan Configuration -* `start_time` (Optional) : Provide the start time to scan objects from all the buckets. This should follow [ISO LocalDateTime](https://docs.oracle.com/javase/8/docs/api/java/time/format/DateTimeFormatter.html#ISO_LOCAL_DATE_TIME) format, or it can be configured to `now` keyword which represents current LocalDateTime. This parameter defines a time range together with either end_time or range. Examples: `2023-01-23T10:00:00`, `now`. -* `end_time` (Optional) : Provide the end time to scan objects from all the buckets. This should follow [ISO LocalDateTime](https://docs.oracle.com/javase/8/docs/api/java/time/format/DateTimeFormatter.html#ISO_LOCAL_DATE_TIME) format, or it can be configured to `now` keyword which represents current LocalDateTime. This parameter defines a time range together with either start_time or range. Examples: `2023-01-23T10:00:00`, `now`. -* `range` (Optional) : Provide the duration to scan objects from all the buckets. This parameter defines a time range together with either `start_time` or `end_time`. -* `scheduling` (Optional): See [Scheduling Configuration](#scheduling_configuration) for details -* `bucket`: Provide S3 bucket information - * `name` (Required if bucket block is used): Provide S3 bucket name. - * `filter` (Optional) : Provide include and exclude list items to filter objects in bucket. - * `include_prefix` (Optional) : Provide the list of include key path prefix. For example `dlq/` - * `exclude_suffix` (Optional) : Provide the list of suffix to exclude items. For example `.csv`. - * `start_time` (Optional) : Provide the start time to scan objects from the current bucket. This parameter defines a time range together with either end_time or range. Example: `2023-01-23T10:00:00`. - * `end_time` (Optional) : Provide the end time to scan objects from the current bucket. This parameter defines a time range together with either start_time or range. Example: `2023-01-23T10:00:00`. - * `range` (Optional) : Provide the duration to scan objects from the current bucket. This parameter defines a time range together with either start_time or end_time. - -> Note: If a time range is not specified, all objects will be included by default. To set a time range, specify any two and only two configurations from start_time, end_time and range. The time range configured on a specific bucket will override the time range specified on the top level - -### Scheduling Configuration - -Schedule interval and amount of times a S3 bucket should be scanned when using S3 Scan. For example, -a `interval` of `PT1H` and a `count` of `3` would result in each bucket being scanned 3 times with 1 hour interval in between each scan, starting after source is ready -and then every hour after the first scan. -* `interval` (Optional) : A String that indicates the minimum interval between each scan. If objects from fist scan are not proceed within configured interval, scan will be done whenever there are no pending objects to process from previous scan. - Supports ISO_8601 notation Strings ("PT20.345S", "PT15M", etc.) as well as simple notation Strings for seconds ("60s") and milliseconds ("1500ms"). - Defaults to 8 hours, and is only applicable when `count` is greater than 1. -* `count` (Optional) : An Integer that specifies how many times bucket will be scanned. Defaults to 1. - - -### AWS Configuration - -The AWS configuration is the same for both SQS and S3. - -* `region` (Optional) : The AWS region to use for credentials. Defaults to [standard SDK behavior to determine the region](https://docs.aws.amazon.com/sdk-for-java/latest/developer-guide/region-selection.html). -* `sts_role_arn` (Optional) : The AWS STS role to assume for requests to SQS and S3. Defaults to null, which will use the [standard SDK behavior for credentials](https://docs.aws.amazon.com/sdk-for-java/latest/developer-guide/credentials.html). -* `sts_external_id` (Optional) : The external ID to attach to AssumeRole requests. - -The following policy shows the necessary permissions for S3 source. `kms:Decrypt` is required if SQS queue is encrypted with AWS [KMS](https://aws.amazon.com/kms/). -```json -{ - "Version": "2012-10-17", - "Statement": [ - { - "Sid": "s3policy", - "Effect": "Allow", - "Action": [ - "s3:GetObject", - "s3:ListBucket", - "s3:DeleteObject", - "sqs:DeleteMessage", - "sqs:ReceiveMessage", - "kms:Decrypt" - ], - "Resource": "*" - } - ] -} -``` -* `aws_sts_header_overrides` (Optional): A map of header overrides to make when assuming the IAM role for the sink plugin. - -## Metrics - -### Counters - -* `s3ObjectsFailed` - The number of S3 objects that the S3 Source failed to read. -* `s3ObjectsNotFound` - The number of S3 objects that the S3 Source failed to read due to a Not Found error from S3. These are also counted toward `s3ObjectsFailed`. -* `s3ObjectsAccessDenied` - The number of S3 objects that the S3 Source failed to read due to an Access Denied or Forbidden error. These are also counted toward `s3ObjectsFailed`. -* `s3ObjectNoRecordsFound` - The number of S3 objects that resulted in 0 records added to the buffer. -* `s3ObjectsSucceeded` - The number of S3 objects that the S3 Source successfully read. -* `sqsMessagesReceived` - The number of SQS messages received from the queue by the S3 Source. -* `sqsMessagesDeleted` - The number of SQS messages deleted from the queue by the S3 Source. -* `sqsMessagesFailed` - The number of SQS messages that the S3 Source failed to parse. -* `sqsMessagesDeleteFailed` - The number of SQS messages that the S3 Source failed to delete from the SQS queue. -* `s3ObjectsDeleted` - The number of S3 objects deleted by the S3 source. -* `s3ObjectsDeleteFailed` - The number of S3 objects that the S3 source failed to delete. -* `acknowledgementSetCallbackCounter` - The number of times End-to-end acknowledgments created an acknowledgment set. - - -### Timers - -* `s3ObjectReadTimeElapsed` - Measures the time the S3 Source takes to perform a request to GET an S3 object, parse it, and write Events to the buffer. -* `sqsMessageDelay` - Measures the time from when S3 records an event time for the creation of an object to when it was fully parsed. - -### Distribution Summaries - -* `s3ObjectSizeBytes` - Measures the size of S3 objects as reported by the S3 `Content-Length`. For compressed objects, this is the compressed size. -* `s3ObjectProcessedBytes` - Measures the bytes processed by the S3 source for a given object. For compressed objects, this is the un-compressed size. -* `s3ObjectsEvents` - Measures the number of events (sometimes called records) produced by an S3 object. ## Developer Guide diff --git a/data-prepper-plugins/s3-source/build.gradle b/data-prepper-plugins/s3-source/build.gradle index 6c5217f306..8f1f721809 100644 --- a/data-prepper-plugins/s3-source/build.gradle +++ b/data-prepper-plugins/s3-source/build.gradle @@ -7,10 +7,6 @@ plugins { id 'java' } -repositories { - mavenCentral() -} - dependencies { implementation project(':data-prepper-api') implementation project(':data-prepper-plugins:buffer-common') @@ -27,7 +23,7 @@ dependencies { implementation libs.commons.io implementation libs.commons.compress implementation 'joda-time:joda-time:2.12.6' - implementation 'org.hibernate.validator:hibernate-validator:7.0.5.Final' + implementation 'org.hibernate.validator:hibernate-validator:8.0.1.Final' implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-csv' implementation 'com.fasterxml.jackson.datatype:jackson-datatype-jsr310' implementation 'org.xerial.snappy:snappy-java:1.1.10.5' diff --git a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanPartitionCreationSupplier.java b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanPartitionCreationSupplier.java index 04da7cdaa1..367505a936 100644 --- a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanPartitionCreationSupplier.java +++ b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanPartitionCreationSupplier.java @@ -165,7 +165,7 @@ private boolean isLastModifiedTimeAfterMostRecentScanForBucket(final String buck final Instant lastProcessedObjectTimestamp = Instant.parse((String) globalStateMap.get(bucketName)); - return s3Object.lastModified().compareTo(lastProcessedObjectTimestamp) > 0; + return s3Object.lastModified().compareTo(lastProcessedObjectTimestamp.minusSeconds(1)) >= 0; } private Instant getMostRecentLastModifiedTimestamp(final ListObjectsV2Response listObjectsV2Response, diff --git a/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanPartitionCreationSupplierTest.java b/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanPartitionCreationSupplierTest.java index 77a110528c..a0ece6d988 100644 --- a/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanPartitionCreationSupplierTest.java +++ b/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanPartitionCreationSupplierTest.java @@ -204,7 +204,7 @@ void getNextPartition_supplier_with_scheduling_options_returns_expected_Partitio final S3Object invalidForFirstBucketSuffixObject = mock(S3Object.class); given(invalidForFirstBucketSuffixObject.key()).willReturn("test.invalid"); - given(invalidForFirstBucketSuffixObject.lastModified()).willReturn(Instant.now()); + given(invalidForFirstBucketSuffixObject.lastModified()).willReturn(Instant.now().minusSeconds(2)); s3ObjectsList.add(invalidForFirstBucketSuffixObject); expectedPartitionIdentifiers.add(PartitionIdentifier.builder().withPartitionKey(secondBucket + "|" + invalidForFirstBucketSuffixObject.key()).build()); @@ -223,7 +223,9 @@ void getNextPartition_supplier_with_scheduling_options_returns_expected_Partitio final List expectedPartitionIdentifiersSecondScan = new ArrayList<>(); expectedPartitionIdentifiersSecondScan.add(PartitionIdentifier.builder().withPartitionKey(firstBucket + "|" + secondScanObject.key()).build()); + expectedPartitionIdentifiersSecondScan.add(PartitionIdentifier.builder().withPartitionKey(firstBucket + "|" + validObject.key()).build()); expectedPartitionIdentifiersSecondScan.add(PartitionIdentifier.builder().withPartitionKey(secondBucket + "|" + secondScanObject.key()).build()); + expectedPartitionIdentifiersSecondScan.add(PartitionIdentifier.builder().withPartitionKey(secondBucket + "|" + validObject.key()).build()); final List secondScanObjects = new ArrayList<>(s3ObjectsList); secondScanObjects.add(secondScanObject); diff --git a/data-prepper-plugins/service-map-stateful/build.gradle b/data-prepper-plugins/service-map-stateful/build.gradle index 3c60b2ebd3..60b9512ed9 100644 --- a/data-prepper-plugins/service-map-stateful/build.gradle +++ b/data-prepper-plugins/service-map-stateful/build.gradle @@ -7,10 +7,6 @@ plugins { id 'java' } -repositories { - mavenCentral() -} - dependencies { implementation project(':data-prepper-api') implementation project(':data-prepper-plugins:common') diff --git a/data-prepper-plugins/split-event-processor/build.gradle b/data-prepper-plugins/split-event-processor/build.gradle new file mode 100644 index 0000000000..0271acc1a3 --- /dev/null +++ b/data-prepper-plugins/split-event-processor/build.gradle @@ -0,0 +1,22 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +jacocoTestCoverageVerification { + dependsOn jacocoTestReport + violationRules { + rule { + limit { + minimum = 1.0 + } + } + } +} + + +dependencies { + implementation project(':data-prepper-api') + implementation project(':data-prepper-plugins:common') + implementation 'com.fasterxml.jackson.core:jackson-databind' +} \ No newline at end of file diff --git a/data-prepper-plugins/split-event-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/splitevent/SplitEventProcessor.java b/data-prepper-plugins/split-event-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/splitevent/SplitEventProcessor.java new file mode 100644 index 0000000000..cadd463ae9 --- /dev/null +++ b/data-prepper-plugins/split-event-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/splitevent/SplitEventProcessor.java @@ -0,0 +1,126 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + */ + +package org.opensearch.dataprepper.plugins.processor.splitevent; + +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; +import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.event.JacksonEvent; +import org.opensearch.dataprepper.model.processor.AbstractProcessor; +import org.opensearch.dataprepper.model.processor.Processor; +import org.opensearch.dataprepper.model.record.Record; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.function.Function; +import java.util.regex.Pattern; + + +@DataPrepperPlugin(name = "split_event", pluginType = Processor.class, pluginConfigurationType = SplitEventProcessorConfig.class) +public class SplitEventProcessor extends AbstractProcessor, Record>{ + final String delimiter; + final String delimiterRegex; + final String field; + final Pattern pattern; + private final Function splitter; + + @DataPrepperPluginConstructor + public SplitEventProcessor(final PluginMetrics pluginMetrics, final SplitEventProcessorConfig config) { + super(pluginMetrics); + this.delimiter = config.getDelimiter(); + this.delimiterRegex = config.getDelimiterRegex(); + this.field = config.getField(); + + if(delimiterRegex != null && !delimiterRegex.isEmpty() + && delimiter != null && !delimiter.isEmpty()) { + throw new IllegalArgumentException("delimiter and delimiter_regex cannot be defined at the same time"); + } else if((delimiterRegex == null || delimiterRegex.isEmpty()) && + (delimiter == null || delimiter.isEmpty())) { + throw new IllegalArgumentException("delimiter or delimiter_regex needs to be defined"); + } + + if(delimiterRegex != null && !delimiterRegex.isEmpty()) { + pattern = Pattern.compile(delimiterRegex); + splitter = pattern::split; + } else { + splitter = inputString -> inputString.split(delimiter); + pattern = null; + } + } + + @Override + public Collection> doExecute(final Collection> records) { + Collection> newRecords = new ArrayList<>(); + for(final Record record : records) { + final Event recordEvent = record.getData(); + + if (!recordEvent.containsKey(field)) { + newRecords.add(record); + continue; + } + + final Object value = recordEvent.get(field, Object.class); + + //split record according to delimiter + final String[] splitValues = splitter.apply((String) value); + + // when no splits or empty value use the original record + if(splitValues.length <= 1) { + newRecords.add(record); + continue; + } + + //create new events for the splits + for (int i = 0; i < splitValues.length-1 ; i++) { + Record newRecord = createNewRecordFromEvent(recordEvent, splitValues[i]); + addToAcknowledgementSetFromOriginEvent((Event) newRecord.getData(), recordEvent); + newRecords.add(newRecord); + } + + // Modify original event to hold the last split + recordEvent.put(field, splitValues[splitValues.length-1]); + newRecords.add(record); + } + return newRecords; + } + + protected Record createNewRecordFromEvent(final Event recordEvent, String splitValue) { + Record newRecord; + JacksonEvent newRecordEvent; + + newRecordEvent = JacksonEvent.fromEvent(recordEvent); + newRecordEvent.put(field,(Object) splitValue); + newRecord = new Record<>(newRecordEvent); + return newRecord; + } + + protected void addToAcknowledgementSetFromOriginEvent(Event recordEvent, Event originRecordEvent) { + DefaultEventHandle eventHandle = (DefaultEventHandle) originRecordEvent.getEventHandle(); + if (eventHandle != null && eventHandle.getAcknowledgementSet() != null) { + eventHandle.getAcknowledgementSet().add(recordEvent); + } + } + + @Override + public void prepareForShutdown() { + } + + @Override + public boolean isReadyForShutdown() { + return true; + } + + @Override + public void shutdown() { + } +} diff --git a/data-prepper-plugins/split-event-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/splitevent/SplitEventProcessorConfig.java b/data-prepper-plugins/split-event-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/splitevent/SplitEventProcessorConfig.java new file mode 100644 index 0000000000..c4af96a3d4 --- /dev/null +++ b/data-prepper-plugins/split-event-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/splitevent/SplitEventProcessorConfig.java @@ -0,0 +1,42 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + */ + +package org.opensearch.dataprepper.plugins.processor.splitevent; + +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.constraints.NotEmpty; +import jakarta.validation.constraints.NotNull; +import jakarta.validation.constraints.Size; + + +public class SplitEventProcessorConfig { + @NotEmpty + @NotNull + @JsonProperty("field") + private String field; + + @JsonProperty("delimiter_regex") + private String delimiterRegex; + + @Size(min = 1, max = 1) + private String delimiter; + + public String getField() { + return field; + } + + public String getDelimiterRegex() { + return delimiterRegex; + } + + public String getDelimiter() { + return delimiter; + } +} diff --git a/data-prepper-plugins/split-event-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/splitevent/SplitEventProcessorTest.java b/data-prepper-plugins/split-event-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/splitevent/SplitEventProcessorTest.java new file mode 100644 index 0000000000..7fc126fdf5 --- /dev/null +++ b/data-prepper-plugins/split-event-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/splitevent/SplitEventProcessorTest.java @@ -0,0 +1,393 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + */ + +package org.opensearch.dataprepper.plugins.processor.splitevent; + + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.event.EventMetadata; +import org.opensearch.dataprepper.model.event.JacksonEvent; +import org.opensearch.dataprepper.model.record.Record; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Stream; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.hamcrest.Matchers.equalTo; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; +import static org.mockito.Mockito.verify; + + +@ExtendWith(MockitoExtension.class) +public class SplitEventProcessorTest { + @Mock + private PluginMetrics pluginMetrics; + + @Mock + private SplitEventProcessorConfig mockConfig; + + @Mock + private AcknowledgementSet mockAcknowledgementSet; + + private SplitEventProcessor splitEventProcessor; + + + private Record createTestRecord(final Map data) { + + Event event = JacksonEvent.builder() + .withData(data) + .withEventType("event") + .build(); + + DefaultEventHandle eventHandle = (DefaultEventHandle) event.getEventHandle(); + + eventHandle.setAcknowledgementSet(mockAcknowledgementSet); + return new Record<>(event); + } + + @BeforeEach + void setup() { + when(mockConfig.getField()).thenReturn("k1"); + when(mockConfig.getDelimiter()).thenReturn(" "); + + splitEventProcessor = new SplitEventProcessor(pluginMetrics, mockConfig); + } + + private static Stream provideMaps() { + return Stream.of( + Arguments.of( + Map.of( + "k1", "", + "k2", "v2" + ) + ), + Arguments.of( + Map.of( + "k1", "v1", + "k2", "v2" + ) + ), + Arguments.of( + Map.of("k1", "v1 v2", + "k2", "v2" + ) + ), + Arguments.of( + Map.of("k1", "v1 v2 v3", + "k2", "v2" + ) + ) + ); + } + + @Test + void testHappyPathWithSpaceDelimiter() { + final Map testData = new HashMap<>(); + testData.put("k1", "v1 v2"); + final Record record = createTestRecord(testData); + when(mockConfig.getDelimiter()).thenReturn(" "); + + final SplitEventProcessor objectUnderTest = new SplitEventProcessor(pluginMetrics, mockConfig); + final List> editedRecords = (List>) objectUnderTest.doExecute(Collections.singletonList(record)); + + for(Record r: editedRecords){ + Event event = (Event) r.getData(); + } + + assertThat(editedRecords.size(), equalTo(2)); + assertThat(editedRecords.get(0).getData().toMap(), equalTo(Map.of("k1","v1"))); + assertThat(editedRecords.get(1).getData().toMap(), equalTo(Map.of("k1","v2"))); + } + + @Test + void testHappyPathWithSpaceDelimiterForStringWithMultipleSpaces() { + final Map testData = new HashMap<>(); + testData.put("k1", "v1 v2"); + final Record record = createTestRecord(testData); + when(mockConfig.getDelimiter()).thenReturn(" "); + + final SplitEventProcessor objectUnderTest = new SplitEventProcessor(pluginMetrics, mockConfig); + final List> editedRecords = (List>) objectUnderTest.doExecute(Collections.singletonList(record)); + for(Record r: editedRecords){ + Event event = (Event) r.getData(); + } + + assertThat(editedRecords.size(), equalTo(3)); + assertThat(editedRecords.get(0).getData().toMap(), equalTo(Map.of("k1","v1"))); + assertThat(editedRecords.get(1).getData().toMap(), equalTo(Map.of("k1",""))); + assertThat(editedRecords.get(2).getData().toMap(), equalTo(Map.of("k1","v2"))); + } + + @Test + void testHappyPathWithSemiColonDelimiter() { + final Map testData = new HashMap<>(); + testData.put("k1", "v1;v2"); + final Record record = createTestRecord(testData); + when(mockConfig.getDelimiter()).thenReturn(";"); + + final SplitEventProcessor objectUnderTest = new SplitEventProcessor(pluginMetrics, mockConfig); + final List> editedRecords = (List>) objectUnderTest.doExecute(Collections.singletonList(record)); + + for(Record r: editedRecords){ + Event event = (Event) r.getData(); + } + + assertThat(editedRecords.size(), equalTo(2)); + assertThat(editedRecords.get(0).getData().toMap(), equalTo(Map.of("k1","v1"))); + assertThat(editedRecords.get(1).getData().toMap(), equalTo(Map.of("k1","v2"))); + } + + @Test + void testHappyPathWithSpaceDelimiterRegex() { + final Map testData = new HashMap<>(); + testData.put("k1", "v1 v2"); + final Record record = createTestRecord(testData); + when(mockConfig.getDelimiter()).thenReturn(null); + when(mockConfig.getDelimiterRegex()).thenReturn("\\s+"); + + final SplitEventProcessor objectUnderTest = new SplitEventProcessor(pluginMetrics, mockConfig); + final List> editedRecords = (List>) objectUnderTest.doExecute(Collections.singletonList(record)); + + for(Record r: editedRecords){ + Event event = (Event) r.getData(); + } + + assertThat(editedRecords.size(), equalTo(2)); + assertThat(editedRecords.get(0).getData().toMap(), equalTo(Map.of("k1","v1"))); + assertThat(editedRecords.get(1).getData().toMap(), equalTo(Map.of("k1","v2"))); + } + + @Test + void testHappyPathWithColonDelimiterRegex() { + final Map testData = new HashMap<>(); + testData.put("k1", "v1:v2"); + final Record record = createTestRecord(testData); + when(mockConfig.getDelimiter()).thenReturn(null); + when(mockConfig.getDelimiterRegex()).thenReturn(":"); + + final SplitEventProcessor objectUnderTest = new SplitEventProcessor(pluginMetrics, mockConfig); + final List> editedRecords = (List>) objectUnderTest.doExecute(Collections.singletonList(record)); + + for(Record r: editedRecords){ + Event event = (Event) r.getData(); + } + + assertThat(editedRecords.size(), equalTo(2)); + assertThat(editedRecords.get(0).getData().toMap(), equalTo(Map.of("k1","v1"))); + assertThat(editedRecords.get(1).getData().toMap(), equalTo(Map.of("k1","v2"))); + } + + @Test + void testFailureWithBothDelimiterRegexAndDelimiterDefined() { + when(mockConfig.getDelimiter()).thenReturn(" "); + when(mockConfig.getDelimiterRegex()).thenReturn("\\s+"); + assertThrows(IllegalArgumentException.class, () -> new SplitEventProcessor(pluginMetrics, mockConfig)); + } + + @Test + void testFailureWithDelimiterRegexAndDelimiterDefinedMissing() { + when(mockConfig.getDelimiter()).thenReturn(null); + when(mockConfig.getDelimiterRegex()).thenReturn(null); + assertThrows(IllegalArgumentException.class, () -> new SplitEventProcessor(pluginMetrics, mockConfig)); + } + + @ParameterizedTest + @MethodSource("provideMaps") + void testSplitEventsBelongToTheSameAcknowledgementSet(Map inputMap1) { + final Record testRecord = createTestRecord(inputMap1); + Event originalEvent = testRecord.getData(); + DefaultEventHandle originalEventHandle = (DefaultEventHandle) originalEvent.getEventHandle(); + AcknowledgementSet originalAcknowledgementSet= originalEventHandle.getAcknowledgementSet(); + + final SplitEventProcessor objectUnderTest = new SplitEventProcessor(pluginMetrics, mockConfig); + final List> editedRecords = (List>) objectUnderTest.doExecute(Collections.singletonList(testRecord)); + + DefaultEventHandle eventHandle = (DefaultEventHandle) originalEvent.getEventHandle(); + AcknowledgementSet acknowledgementSet; + for(Record record: editedRecords) { + Event event = testRecord.getData(); + eventHandle = (DefaultEventHandle) event.getEventHandle(); + acknowledgementSet = eventHandle.getAcknowledgementSet(); + assertEquals(originalAcknowledgementSet, acknowledgementSet); + } + } + + @Test + void testSplitEventsWhenNoSplits() { + final Map testData = new HashMap<>(); + List> records = new ArrayList<>(); + + testData.put("k1", "v1"); + testData.put("k2", "v2"); + final Record record1 = createTestRecord(testData); + + final Map testData2 = new HashMap<>(); + testData2.put("k1", ""); + testData2.put("k3", "v3"); + final Record record2 = createTestRecord(testData2); + + records.add(record1); + records.add(record2); + final List> editedRecords = (List>) splitEventProcessor.doExecute(records); + + assertThat(editedRecords.size(), equalTo(2)); + assertThat(editedRecords.get(0).getData().toMap(), equalTo(Map.of("k1","v1", "k2", "v2"))); + assertThat(editedRecords.get(1).getData().toMap(), equalTo(Map.of("k1","", "k3", "v3"))); + } + + + @Test + void testSplitEventsWhenOneSplit() { + List> records = new ArrayList<>(); + final Map testData = new HashMap<>(); + testData.put("k1", "v1 v2"); + testData.put("k2", "v3"); + final Record record = createTestRecord(testData); + records.add(record); + final List> editedRecords = (List>) splitEventProcessor.doExecute(records); + + assertThat(editedRecords.size(), equalTo(2)); + assertThat(editedRecords.get(0).getData().toMap(), equalTo(Map.of("k1","v1", "k2", "v3"))); + assertThat(editedRecords.get(1).getData().toMap(), equalTo(Map.of("k1","v2", "k2", "v3"))); + } + + + @Test + void testSplitEventsWhenMultipleSplits() { + List> records = new ArrayList<>(); + final Map testData = new HashMap<>(); + testData.put("k1", "v1 v2 v3 v4"); + testData.put("k2", "v5"); + final Record record = createTestRecord(testData); + records.add(record); + final List> editedRecords = (List>) splitEventProcessor.doExecute(records); + + assertThat(editedRecords.size(), equalTo(4)); + assertThat(editedRecords.get(0).getData().toMap(), equalTo(Map.of("k1","v1", "k2", "v5"))); + assertThat(editedRecords.get(1).getData().toMap(), equalTo(Map.of("k1","v2", "k2", "v5"))); + assertThat(editedRecords.get(2).getData().toMap(), equalTo(Map.of("k1","v3", "k2", "v5"))); + assertThat(editedRecords.get(3).getData().toMap(), equalTo(Map.of("k1","v4", "k2", "v5"))); + } + + @Test + void testSplitEventsWhenMultipleSplitsMultipleRecords() { + List> records = new ArrayList<>(); + final Map testData1 = new HashMap<>(); + testData1.put("k1", "v1 v2"); + testData1.put("k2", "v5"); + final Record record = createTestRecord(testData1); + records.add(record); + + final Map testData2 = new HashMap<>(); + testData2.put("k1", "v1"); + testData2.put("k2", "v3"); + final Record record2 = createTestRecord(testData2); + records.add(record2); + final List> editedRecords = (List>) splitEventProcessor.doExecute(records); + + assertThat(editedRecords.size(), equalTo(3)); + assertThat(editedRecords.get(0).getData().toMap(), equalTo(Map.of("k1","v1", "k2", "v5"))); + assertThat(editedRecords.get(1).getData().toMap(), equalTo(Map.of("k1","v2", "k2", "v5"))); + assertThat(editedRecords.get(2).getData().toMap(), equalTo(Map.of("k1","v1", "k2", "v3"))); + } + + @Test + void testSplitEventsWhenNoKeyPresentInEvent() { + List> records = new ArrayList<>(); + final Map testData1 = new HashMap<>(); + testData1.put("k2", "v5 v6"); + final Record record = createTestRecord(testData1); + records.add(record); + + final Map testData2 = new HashMap<>(); + testData2.put("k3", "v3 v5"); + final Record record2 = createTestRecord(testData2); + records.add(record2); + final List> editedRecords = (List>) splitEventProcessor.doExecute(records); + + assertThat(editedRecords.size(), equalTo(2)); + assertThat(editedRecords.get(0).getData().toMap(), equalTo(Map.of("k2", "v5 v6"))); + assertThat(editedRecords.get(1).getData().toMap(), equalTo(Map.of("k3", "v3 v5"))); + } + + @Test + public void testCreateNewRecordFromEvent() { + Event recordEvent = mock(Event.class); + + Map eventData = new HashMap<>(); + eventData.put("someField", "someValue"); + when(recordEvent.toMap()).thenReturn(eventData); + when(recordEvent.getMetadata()).thenReturn(mock(EventMetadata.class)); + String splitValue = "splitValue"; + + Record resultRecord = splitEventProcessor.createNewRecordFromEvent(recordEvent, splitValue); + Event editedEvent = (Event) resultRecord.getData(); + // Assertions + assertEquals(editedEvent.getMetadata(),recordEvent.getMetadata()); + } + + @Test + public void testAddToAcknowledgementSetFromOriginEvent() { + Map data = Map.of("k1","v1"); + EventMetadata eventMetadata = mock(EventMetadata.class); + Event originRecordEvent = JacksonEvent.builder() + .withEventMetadata(eventMetadata) + .withEventType("event") + .withData(data) + .build(); + Event spyEvent = spy(originRecordEvent); + + DefaultEventHandle mockEventHandle = mock(DefaultEventHandle.class); + when(spyEvent.getEventHandle()).thenReturn(mockEventHandle); + + Record record = splitEventProcessor + .createNewRecordFromEvent(spyEvent, "v1"); + + Event recordEvent = (Event) record.getData(); + splitEventProcessor.addToAcknowledgementSetFromOriginEvent(recordEvent, spyEvent); + + DefaultEventHandle spyEventHandle = (DefaultEventHandle) spyEvent.getEventHandle(); + // Verify that the add method is called on the acknowledgement set + verify(spyEventHandle).getAcknowledgementSet(); + + AcknowledgementSet spyAckSet = spyEventHandle.getAcknowledgementSet(); + DefaultEventHandle eventHandle = (DefaultEventHandle) recordEvent.getEventHandle(); + AcknowledgementSet ackSet1 = eventHandle.getAcknowledgementSet(); + + assertEquals(spyAckSet, ackSet1); + } + + @Test + void testIsReadyForShutdown() { + assertTrue(splitEventProcessor.isReadyForShutdown()); + } + +} diff --git a/data-prepper-plugins/truncate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/truncate/TruncateProcessor.java b/data-prepper-plugins/truncate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/truncate/TruncateProcessor.java index 8449675791..c9b71173ce 100644 --- a/data-prepper-plugins/truncate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/truncate/TruncateProcessor.java +++ b/data-prepper-plugins/truncate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/truncate/TruncateProcessor.java @@ -13,11 +13,15 @@ import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.processor.AbstractProcessor; import org.opensearch.dataprepper.model.processor.Processor; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.util.Collection; import java.util.ArrayList; import java.util.List; +import static org.opensearch.dataprepper.logging.DataPrepperMarkers.EVENT; + /** * This processor takes in a key and truncates its value to a string with * characters from the front or at the end or at both removed. @@ -25,6 +29,7 @@ */ @DataPrepperPlugin(name = "truncate", pluginType = Processor.class, pluginConfigurationType = TruncateProcessorConfig.class) public class TruncateProcessor extends AbstractProcessor, Record>{ + private static final Logger LOG = LoggerFactory.getLogger(TruncateProcessor.class); private final ExpressionEvaluator expressionEvaluator; private final List entries; @@ -48,34 +53,39 @@ private String getTruncatedValue(final String value, final int startIndex, final public Collection> doExecute(final Collection> records) { for(final Record record : records) { final Event recordEvent = record.getData(); - for (TruncateProcessorConfig.Entry entry: entries) { - final List sourceKeys = entry.getSourceKeys(); - final String truncateWhen = entry.getTruncateWhen(); - final int startIndex = entry.getStartAt() == null ? 0 : entry.getStartAt(); - final Integer length = entry.getLength(); - if (truncateWhen != null && !expressionEvaluator.evaluateConditional(truncateWhen, recordEvent)) { - continue; - } - for (String sourceKey: sourceKeys) { - if (!recordEvent.containsKey(sourceKey)) { + + try { + for (TruncateProcessorConfig.Entry entry : entries) { + final List sourceKeys = entry.getSourceKeys(); + final String truncateWhen = entry.getTruncateWhen(); + final int startIndex = entry.getStartAt() == null ? 0 : entry.getStartAt(); + final Integer length = entry.getLength(); + if (truncateWhen != null && !expressionEvaluator.evaluateConditional(truncateWhen, recordEvent)) { continue; } + for (String sourceKey : sourceKeys) { + if (!recordEvent.containsKey(sourceKey)) { + continue; + } - final Object value = recordEvent.get(sourceKey, Object.class); - if (value instanceof String) { - recordEvent.put(sourceKey, getTruncatedValue((String)value, startIndex, length)); - } else if (value instanceof List) { - List result = new ArrayList<>(); - for (Object listItem: (List)value) { - if (listItem instanceof String) { - result.add(getTruncatedValue((String)listItem, startIndex, length)); - } else { - result.add(listItem); + final Object value = recordEvent.get(sourceKey, Object.class); + if (value instanceof String) { + recordEvent.put(sourceKey, getTruncatedValue((String) value, startIndex, length)); + } else if (value instanceof List) { + List result = new ArrayList<>(); + for (Object listItem : (List) value) { + if (listItem instanceof String) { + result.add(getTruncatedValue((String) listItem, startIndex, length)); + } else { + result.add(listItem); + } } + recordEvent.put(sourceKey, result); } - recordEvent.put(sourceKey, result); } } + } catch (final Exception e) { + LOG.error(EVENT, "There was an exception while processing Event [{}]", recordEvent, e); } } diff --git a/data-prepper-plugins/user-agent-processor/README.md b/data-prepper-plugins/user-agent-processor/README.md index 950161f46b..a250472e0d 100644 --- a/data-prepper-plugins/user-agent-processor/README.md +++ b/data-prepper-plugins/user-agent-processor/README.md @@ -1,47 +1,4 @@ # User Agent Processor This processor parses User-Agent (UA) string in an event and add the parsing result to the event. -## Basic Example -An example configuration for the process is as follows: -```yaml -... - processor: - - user_agent: - source: "ua" - target: "user_agent" -... -``` - -Assume the event contains the following user agent string: -```json -{ - "ua": "Mozilla/5.0 (iPhone; CPU iPhone OS 13_5_1 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/13.1.1 Mobile/15E148 Safari/604.1" -} -``` - -The processor will parse the "ua" field and add the result to the specified target in the following format compatible with Elastic Common Schema (ECS): -``` -{ - "user_agent": { - "original": "Mozilla/5.0 (iPhone; CPU iPhone OS 13_5_1 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/13.1.1 Mobile/15E148 Safari/604.1", - "os": { - "version": "13.5.1", - "full": "iOS 13.5.1", - "name": "iOS" - }, - "name": "Mobile Safari", - "version": "13.1.1", - "device": { - "name": "iPhone" - } - }, - "ua": "Mozilla/5.0 (iPhone; CPU iPhone OS 13_5_1 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/13.1.1 Mobile/15E148 Safari/604.1" -} -``` - -## Configuration -* `source` (Required) — The key to the user agent string in the Event that will be parsed. -* `target` (Optional) — The key to put the parsing result in the Event. Defaults to `user_agent`. -* `exclude_original` (Optional) — Whether to exclude original user agent string from the parsing result. Defaults to false. -* `cache_size` (Optional) - Cache size to use in the parser. Should be a positive integer. Defaults to 1000. -* `tags_on_parse_failure` (Optional) - Tags to add to an event if the processor fails to parse the user agent string. +See the [`user_agent` processor documentation](https://opensearch.org/docs/latest/data-prepper/pipelines/configuration/processors/user-agent/). diff --git a/data-prepper-test-common/build.gradle b/data-prepper-test-common/build.gradle index a511e48891..be7559311d 100644 --- a/data-prepper-test-common/build.gradle +++ b/data-prepper-test-common/build.gradle @@ -7,10 +7,6 @@ plugins { id 'java' } -repositories { - mavenCentral() -} - dependencies { implementation testLibs.hamcrest testRuntimeOnly testLibs.junit.engine diff --git a/docs/expression_syntax.md b/docs/expression_syntax.md index 80191153e1..ad0392d408 100644 --- a/docs/expression_syntax.md +++ b/docs/expression_syntax.md @@ -137,7 +137,7 @@ Hard coded token that identifies the operation use in an _Expression_. ### Json Pointer A Literal used to reference a value within the Event provided as context for the _Expression String_. Json Pointers are identified by a -leading `/` containing alphanumeric character or underscores, delimited by `/`. Json Pointers can use an extended character set if wrapped +leading `/` containing one of more of characters from the set `a-zA-Z0-9_.@` , delimited by `/`. Json Pointers can use an extended character set if wrapped in double quotes (`"`) using the escape character `\`. Note, Json Pointer require `~` and `/` that should be used as part of the path and not a delimiter to be escaped. diff --git a/performance-test/build.gradle b/performance-test/build.gradle index b297cdbfcc..b8989f63a4 100644 --- a/performance-test/build.gradle +++ b/performance-test/build.gradle @@ -14,12 +14,8 @@ configurations.all { group 'org.opensearch.dataprepper.test.performance' -repositories { - mavenCentral() -} - dependencies { - gatlingImplementation 'software.amazon.awssdk:auth:2.23.13' + gatlingImplementation 'software.amazon.awssdk:auth:2.24.5' implementation 'com.fasterxml.jackson.core:jackson-core' testRuntimeOnly testLibs.junit.engine diff --git a/release/build-resources.gradle b/release/build-resources.gradle index 67e2ce8c61..492772c39d 100644 --- a/release/build-resources.gradle +++ b/release/build-resources.gradle @@ -9,7 +9,7 @@ ext { linux: ['x64'] ] jdkSources = [ - linux_x64: 'https://github.com/adoptium/temurin17-binaries/releases/download/jdk-17.0.8%2B7/OpenJDK17U-jdk_x64_linux_hotspot_17.0.8_7.tar.gz', + linux_x64: 'https://github.com/adoptium/temurin17-binaries/releases/download/jdk-17.0.10%2B7/OpenJDK17U-jdk_x64_linux_hotspot_17.0.10_7.tar.gz', linux_arm64: 'https://hg.openjdk.java.net/aarch64-port/jdk8/archive/tip.tar.gz' ] awsResources = [ diff --git a/release/release-notes/data-prepper.change-log-2.6.2.md b/release/release-notes/data-prepper.change-log-2.6.2.md new file mode 100644 index 0000000000..ef624f336d --- /dev/null +++ b/release/release-notes/data-prepper.change-log-2.6.2.md @@ -0,0 +1,302 @@ + +* __Adds release notes for Data Prepper 2.6.2 (#4149) (#4151)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Mon, 19 Feb 2024 08:33:17 -0800 + + EAD -> refs/heads/2.6, tag: refs/tags/2.6.2, refs/remotes/upstream/2.6 + * Adds release notes for Data Prepper 2.6.2 + Signed-off-by: David Venable <dlv@amazon.com> + (cherry picked from commit 119ccb6410400e51554be41c11b0a115dd176eac) + Co-authored-by: David Venable <dlv@amazon.com> + +* __Generated THIRD-PARTY file for 52f4697 (#4150)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Mon, 19 Feb 2024 08:28:41 -0800 + + + Signed-off-by: GitHub <noreply@github.com> + Co-authored-by: dlvenable + <dlvenable@users.noreply.github.com> + +* __Bump grpcio in /release/smoke-tests/otel-span-exporter (#4104) (#4148)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Mon, 19 Feb 2024 06:40:20 -0800 + + + Bumps [grpcio](https://github.com/grpc/grpc) from 1.53.0 to 1.53.2. + - [Release notes](https://github.com/grpc/grpc/releases) + - + [Changelog](https://github.com/grpc/grpc/blob/master/doc/grpc_release_schedule.md) + + - [Commits](https://github.com/grpc/grpc/compare/v1.53.0...v1.53.2) + + --- + updated-dependencies: + - dependency-name: grpcio + dependency-type: direct:production + ... + Signed-off-by: dependabot[bot] <support@github.com> + Co-authored-by: + dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> + (cherry picked from commit 30d68966a85366b244dd4db6c067707dd74764c2) + Co-authored-by: dependabot[bot] + <49699333+dependabot[bot]@users.noreply.github.com> + +* __Updates the JDK version of the release to jdk-17.0.10+7. (#4136) (#4141)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Sat, 17 Feb 2024 07:12:07 -0800 + + + Signed-off-by: David Venable <dlv@amazon.com> + (cherry picked from commit 8bf0daa4bb35d80a5c63e30f552eda04414c3e4b) + Co-authored-by: David Venable <dlv@amazon.com> + +* __FIX: plugin callback not loaded for secret refreshment (#4079) (#4140)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Fri, 16 Feb 2024 12:02:51 -0800 + + + Signed-off-by: George Chen <qchea@amazon.com> + (cherry picked from commit 2f4c8c9c7f8d4ec6e76c3653ef8446fcee35cd50) + Co-authored-by: Qi Chen <qchea@amazon.com> + +* __Bump com.github.seancfoley:ipaddress in /data-prepper-expression (#4060) (#4077)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Fri, 16 Feb 2024 09:57:16 -0600 + + + Bumps + [com.github.seancfoley:ipaddress](https://github.com/seancfoley/IPAddress) from + 5.4.0 to 5.4.2. + - [Release notes](https://github.com/seancfoley/IPAddress/releases) + - [Commits](https://github.com/seancfoley/IPAddress/compare/v5.4.0...v5.4.2) + + --- + updated-dependencies: + - dependency-name: com.github.seancfoley:ipaddress + dependency-type: direct:production + update-type: version-update:semver-patch + ... + Signed-off-by: dependabot[bot] <support@github.com> + Co-authored-by: + dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> + (cherry picked from commit 16d0d907a29483a5f7bdcf2c04c7bda87121366d) + Co-authored-by: dependabot[bot] + <49699333+dependabot[bot]@users.noreply.github.com> + +* __Require json-path 2.9.0 to fix CVE-2023-51074. Resolves #3919. (#4132) (#4133)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Fri, 16 Feb 2024 06:23:02 -0800 + + + Signed-off-by: David Venable <dlv@amazon.com> + (cherry picked from commit 838744c9f00fa9e75c4cd254a83f1fb493ca9ba9) + Co-authored-by: David Venable <dlv@amazon.com> + +* __Fix bug where s3 scan could skip when lastModifiedTimestamps are the same (#4124) (#4127)__ + + [Taylor Gray](mailto:tylgry@amazon.com) - Wed, 14 Feb 2024 13:31:29 -0600 + + + Signed-off-by: Taylor Gray <tylgry@amazon.com> + +* __Catch exception instead of shutting down in date processor (#4108) (#4117)__ + + [Taylor Gray](mailto:tylgry@amazon.com) - Mon, 12 Feb 2024 16:56:42 -0600 + + + Signed-off-by: Taylor Gray <tylgry@amazon.com> + (cherry picked from commit 0841ac17ae18a76610842be3c0ff3b0b0dc1e453) + +* __Updates jline to 3.25.0 to resolve CVE-2023-50572. (#4020) (#4029)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Tue, 30 Jan 2024 07:50:41 -0800 + + + Signed-off-by: David Venable <dlv@amazon.com> + (cherry picked from commit 8f0268bb4ac891467133096154acf42c39fd5aca) + Co-authored-by: David Venable <dlv@amazon.com> + +* __Cancel the existing grok task when a timeout occurs. Resolves #4026 (#4027) (#4028)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Mon, 29 Jan 2024 16:28:03 -0800 + + + Signed-off-by: David Venable <dlv@amazon.com> + (cherry picked from commit 76163969d6f030719f0a32dfdc2c4f4253dadd51) + Co-authored-by: David Venable <dlv@amazon.com> + +* __Mark the EventHandle as transient in the JacksonEvent to fix a serialization error with peer forwarding. Resolves #3981. (#3983) (#3987)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Fri, 19 Jan 2024 08:56:41 -0800 + + + Signed-off-by: David Venable <dlv@amazon.com> + (cherry picked from commit 564a749c8fec9a8c0cb7ec71d52fa0fc83760101) + Co-authored-by: David Venable <dlv@amazon.com> + +* __Release events that are not routed to any sinks (#3959) (#3966)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Wed, 17 Jan 2024 12:07:11 -0800 + + + * Release events that are not routed to any sinks + Signed-off-by: Krishna Kondaka + <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + + * Addressed review comments + Signed-off-by: Krishna Kondaka + <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + + * Fixed a bug in the code that's causing the test failures + Signed-off-by: Krishna Kondaka + <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + + * Modified to determine unrouted events after all routing is done + Signed-off-by: Krishna Kondaka + <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + + * Add test yaml files + Signed-off-by: Krishna Kondaka + <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + + * Addressed review comments + Signed-off-by: Krishna Kondaka + <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + + --------- + Signed-off-by: Krishna Kondaka + <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + Co-authored-by: + Krishna Kondaka <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + (cherry picked from commit f21937adc96e87e2dc932348b9f609afbf18f94c) + Co-authored-by: kkondaka <41027584+kkondaka@users.noreply.github.com> + +* __Updates wiremock to 3.3.1. This also involves changing the groupId to org.wiremock which is the new groupId as of 3.0.0. (#3969) (#3971)__ + + [David Venable](mailto:dlv@amazon.com) - Wed, 17 Jan 2024 11:51:49 -0800 + + + Signed-off-by: David Venable <dlv@amazon.com> + (cherry picked from commit e0ed91c11d9942867f89c29c08b37b52d2ce2652) + +* __Version bump to 2.6.2 (#3946)__ + + [David Venable](mailto:dlv@amazon.com) - Fri, 12 Jan 2024 14:13:14 -0800 + + + Signed-off-by: David Venable <dlv@amazon.com> + +* __Add your public modifier back to one of the AbstractBuffer constructors to attempt to fix the build. (#3947) (#3948)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Thu, 11 Jan 2024 10:45:05 -0800 + + + Signed-off-by: David Venable <dlv@amazon.com> + (cherry picked from commit 677643df66dc0cf62091586d8bb9a3417030bd5a) + Co-authored-by: David Venable <dlv@amazon.com> + +* __Corrects the bufferUsage metric by making it equal to the difference between the bufferCapacity and the available permits in the semaphore. Adds a new capacityUsed metric which tracks the actual capacity used by the semaphore which blocks. Resolves #3936. (#3937) (#3940)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Wed, 10 Jan 2024 15:17:41 -0800 + + + Signed-off-by: David Venable <dlv@amazon.com> + (cherry picked from commit d61b0c5d18210db202700a8a08ebcf5f6631768d) + Co-authored-by: David Venable <dlv@amazon.com> + +* __Fix for [BUG] Data Prepper is losing connections from S3 pool (#3836) (#3844)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Tue, 9 Jan 2024 14:00:00 -0800 + + + * Fix for [BUG] Data Prepper is losing connections from S3 pool + Signed-off-by: Krishna Kondaka + <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + + * Addressed review comments + Signed-off-by: Krishna Kondaka + <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + + * Fixed CheckStyle errors + Signed-off-by: Krishna Kondaka + <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + + --------- + Signed-off-by: Krishna Kondaka + <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + Co-authored-by: + Krishna Kondaka <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + (cherry picked from commit f9be56a65562e4e3b4906bc02d93e2fe5c9d4928) + Co-authored-by: kkondaka <41027584+kkondaka@users.noreply.github.com> + +* __Fix Null Pointer Exception in KeyValue Processor (#3927) (#3932)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Tue, 9 Jan 2024 13:59:45 -0800 + + + * Fix Null Pointer Exception in KeyValue Processor + Signed-off-by: Krishna Kondaka + <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + + * Added a test case + Signed-off-by: Krishna Kondaka + <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + + --------- + Signed-off-by: Krishna Kondaka + <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + Co-authored-by: + Krishna Kondaka <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + (cherry picked from commit 35a69489c2f8621c8aa258ddd8dda105cd67a9e4) + Co-authored-by: kkondaka <41027584+kkondaka@users.noreply.github.com> + +* __Add 4xx aggregate metric and shard progress metric for dynamodb source (#3913) (#3921)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Fri, 5 Jan 2024 14:42:44 -0600 + + + Signed-off-by: Taylor Gray <tylgry@amazon.com> + (cherry picked from commit e6df3eb2cd46ebd13dd1c7b808d288c2f3d6ee51) + Co-authored-by: Taylor Gray <tylgry@amazon.com> + +* __Updates Armeria to 1.26.4. This also updates io.grpc to 1.58.0 which has a slight breaking changing. This is fixed by explicitly adding io.grpc:grpc-inprocess to the build. (#3915) (#3917)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Fri, 5 Jan 2024 12:18:16 -0800 + + + Signed-off-by: David Venable <dlv@amazon.com> + (cherry picked from commit a243182f214583379249526f55215ae5b36d72d3) + Co-authored-by: David Venable <dlv@amazon.com> + +* __Updates opensearch library to 1.3.14. And run integration test against 2.11.1 and 1.3.14 as well. Resolves #3837. (#3838) (#3862)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Tue, 2 Jan 2024 14:12:22 -0800 + + + Signed-off-by: David Venable <dlv@amazon.com> + (cherry picked from commit df2bde6cc4d752013dc9a6f9266f651b43668b23) + Co-authored-by: David Venable <dlv@amazon.com> + +* __Require Mozilla Rhino 1.7.12 to fix SNYK-JAVA-ORGMOZILLA-1314295. (#3839) (#3842)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Tue, 19 Dec 2023 14:16:12 -0800 + + + Signed-off-by: David Venable <dlv@amazon.com> + (cherry picked from commit e09900a144753645be41cca7e1f618966e02ea58) + Co-authored-by: David Venable <dlv@amazon.com> + +* __rebasing to latest (#3846) (#3853)__ + + [opensearch-trigger-bot[bot]](mailto:98922864+opensearch-trigger-bot[bot]@users.noreply.github.com) - Wed, 13 Dec 2023 09:11:55 -0800 + + + Signed-off-by: Krishna Kondaka + <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + Co-authored-by: + Krishna Kondaka <krishkdk@dev-dsk-krishkdk-2c-bd29c437.us-west-2.amazon.com> + (cherry picked from commit f9d9e978bec8aad2dba6f7bf41f5ed07e9d68a3f) + Co-authored-by: kkondaka <41027584+kkondaka@users.noreply.github.com> + + diff --git a/release/release-notes/data-prepper.release-notes-2.6.2.md b/release/release-notes/data-prepper.release-notes-2.6.2.md new file mode 100644 index 0000000000..b1c2707f2a --- /dev/null +++ b/release/release-notes/data-prepper.release-notes-2.6.2.md @@ -0,0 +1,27 @@ +## 2024-02-19 Version 2.6.2 + +--- + +### Enhancements +* Add 4xx aggregate metric and shard progress metric for dynamodb source ([#3913](https://github.com/opensearch-project/data-prepper/pull/3913)) + + +### Bug Fixes +* S3 Scan has potential to filter out objects with the same timestamp ([#4123](https://github.com/opensearch-project/data-prepper/issues/4123)) +* Kafka buffer attempts to create a topic when disabled ([#4111](https://github.com/opensearch-project/data-prepper/issues/4111)) +* Grok processor match requests continue after timeout ([#4026](https://github.com/opensearch-project/data-prepper/issues/4026)) +* Serialization error during peer-forwarding ([#3981](https://github.com/opensearch-project/data-prepper/issues/3981)) +* BlockingBuffer.bufferUsage metric does not include records in-flight ([#3936](https://github.com/opensearch-project/data-prepper/issues/3936)) +* Null Pointer Exception in Key Value Processor ([#3928](https://github.com/opensearch-project/data-prepper/issues/3928)) +* Incomplete route set leads to duplicates when E2E ack is enabled. ([#3866](https://github.com/opensearch-project/data-prepper/issues/3866)) +* Data Prepper is losing connections from S3 pool ([#3809](https://github.com/opensearch-project/data-prepper/issues/3809)) +* Key value processor will throw NPE if source key does not exist in the Event ([#3496](https://github.com/opensearch-project/data-prepper/issues/3496)) +* Exception in substitute string processor shuts down processor work but not pipeline ([#2956](https://github.com/opensearch-project/data-prepper/issues/2956)) +* Add 4xx aggregate metric and shard progress metric for dynamodb source ([#3921](https://github.com/opensearch-project/data-prepper/pull/3921)) + +### Security +* Fix GHSA-6g3j-p5g6-992f from OpenSearch jar ([#3837](https://github.com/opensearch-project/data-prepper/issues/3837)) +* Fix CVE-2023-41329 (Medium) detected in wiremock-3.0.1.jar ([#3954](https://github.com/opensearch-project/data-prepper/issues/3954)) +* Fix CVE-2023-51074 (Medium) detected in json-path-2.8.0.jar ([#3919](https://github.com/opensearch-project/data-prepper/issues/3919)) +* Fix CVE-2023-50572 (Medium) detected in jline-3.9.0.jar, jline-3.22.0.jar ([#3871](https://github.com/opensearch-project/data-prepper/issues/3871)) +* Require Mozilla Rhino 1.7.12 to fix SNYK-JAVA-ORGMOZILLA-1314295. ([#3839](https://github.com/opensearch-project/data-prepper/pull/3839)) diff --git a/release/smoke-tests/otel-span-exporter/requirements.txt b/release/smoke-tests/otel-span-exporter/requirements.txt index 62e5b4548f..70fa887422 100644 --- a/release/smoke-tests/otel-span-exporter/requirements.txt +++ b/release/smoke-tests/otel-span-exporter/requirements.txt @@ -3,7 +3,7 @@ certifi==2023.7.22 charset-normalizer==2.0.9 Deprecated==1.2.13 googleapis-common-protos==1.53.0 -grpcio==1.53.0 +grpcio==1.53.2 idna==3.3 opentelemetry-api==1.7.1 opentelemetry-exporter-otlp==1.7.1 diff --git a/settings.gradle b/settings.gradle index c5d0d6c916..567e4aa40b 100644 --- a/settings.gradle +++ b/settings.gradle @@ -10,6 +10,19 @@ pluginManagement { } } +dependencyResolutionManagement { + repositories { + mavenCentral() + maven { + url 'https://packages.confluent.io/maven/' + } + maven { + url 'https://jitpack.io' + content { includeGroup 'com.github.fge' } + } + } +} + rootProject.name = 'opensearch-data-prepper' dependencyResolutionManagement { @@ -100,7 +113,6 @@ include 'data-prepper-plugins:blocking-buffer' include 'data-prepper-plugins:http-source' include 'data-prepper-plugins:drop-events-processor' include 'data-prepper-plugins:key-value-processor' -include 'data-prepper-plugins:type-conversion-processor' include 'data-prepper-plugins:mutate-event-processors' include 'data-prepper-plugins:geoip-processor' include 'data-prepper-plugins:grok-processor' @@ -134,7 +146,6 @@ include 'data-prepper-plugins:failures-common' include 'data-prepper-plugins:newline-codecs' include 'data-prepper-plugins:avro-codecs' include 'data-prepper-plugins:kafka-plugins' -include 'data-prepper-plugins:kafka-connect-plugins' include 'data-prepper-plugins:user-agent-processor' include 'data-prepper-plugins:in-memory-source-coordination-store' include 'data-prepper-plugins:aws-plugin-api' @@ -150,3 +161,7 @@ include 'data-prepper-plugins:buffer-common' //include 'data-prepper-plugins:prometheus-sink' include 'data-prepper-plugins:dissect-processor' include 'data-prepper-plugins:dynamodb-source' +include 'data-prepper-plugins:decompress-processor' +include 'data-prepper-plugins:split-event-processor' +include 'data-prepper-plugins:http-common' +include 'data-prepper-plugins:flatten-processor'