From 1b56c196b013d43266b1b6d47358bac0215f912c Mon Sep 17 00:00:00 2001 From: Adi Suresh Date: Mon, 23 Oct 2023 14:12:59 -0500 Subject: [PATCH 01/53] Add logic to create or update serverless network policy. (#3510) Signed-off-by: Adi Suresh --- data-prepper-plugins/opensearch/README.md | 6 + data-prepper-plugins/opensearch/build.gradle | 2 +- .../opensearch/ConnectionConfiguration.java | 69 ++++ .../sink/opensearch/OpenSearchSink.java | 21 + .../ServerlessNetworkPolicyUpdater.java | 159 ++++++++ .../ConnectionConfigurationTests.java | 30 ++ .../ServerlessNetworkPolicyUpdaterTest.java | 380 ++++++++++++++++++ 7 files changed, 666 insertions(+), 1 deletion(-) create mode 100644 data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/ServerlessNetworkPolicyUpdater.java create mode 100644 data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/ServerlessNetworkPolicyUpdaterTest.java diff --git a/data-prepper-plugins/opensearch/README.md b/data-prepper-plugins/opensearch/README.md index 3a76ae1db3..5bff437fa0 100644 --- a/data-prepper-plugins/opensearch/README.md +++ b/data-prepper-plugins/opensearch/README.md @@ -297,6 +297,12 @@ if `exclude_keys` is set to ["message", "status"], the document written to OpenS * `sts_role_arn` (Optional) : The STS role to assume for requests to AWS. 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_header_overrides` (Optional): A map of header overrides to make when assuming the IAM role for the sink plugin. * `serverless` (Optional): A boolean flag to indicate the OpenSearch backend is Amazon OpenSearch Serverless. Default to `false`. Notice that [ISM policies.](https://opensearch.org/docs/latest/im-plugin/ism/policies/) is not supported in Amazon OpenSearch Serverless and thus any ISM related configuration value has no effect, i.e. `ism_policy_file`. +* `serverless_options` (Optional): Additional options you can specify when using serverless. + +#### Serverless Configuration +* `network_policy_name` (Optional): The serverless network policy name being used. If both `collection_name` and `vpce_id` are specified, then this network policy will be attempted to be created or update. On the managed OpenSearch Ingestion Service, the `collection_name` and `vpce_id` fields are automatically set. +* `collection_name` (Optional): The serverless collection name. +* `vpce_id` (Optional): The VPCE ID connected to Amazon OpenSearch Serverless. ## Metrics ### Management Disabled Index Type diff --git a/data-prepper-plugins/opensearch/build.gradle b/data-prepper-plugins/opensearch/build.gradle index 8f26965ceb..e74771543b 100644 --- a/data-prepper-plugins/opensearch/build.gradle +++ b/data-prepper-plugins/opensearch/build.gradle @@ -28,6 +28,7 @@ dependencies { implementation 'software.amazon.awssdk:arns' implementation 'io.micrometer:micrometer-core' implementation 'software.amazon.awssdk:s3' + implementation 'software.amazon.awssdk:opensearchserverless' implementation libs.commons.lang3 implementation 'software.amazon.awssdk:apache-client' testImplementation testLibs.junit.vintage @@ -72,7 +73,6 @@ task integrationTest(type: Test) { } } - jacocoTestReport { dependsOn test reports { diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/ConnectionConfiguration.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/ConnectionConfiguration.java index 678d4fb00d..5600b699f6 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/ConnectionConfiguration.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/ConnectionConfiguration.java @@ -36,10 +36,14 @@ import software.amazon.awssdk.arns.Arn; import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; import software.amazon.awssdk.auth.signer.Aws4Signer; +import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration; +import software.amazon.awssdk.core.retry.RetryPolicy; +import software.amazon.awssdk.core.retry.backoff.FullJitterBackoffStrategy; import software.amazon.awssdk.http.SdkHttpClient; import software.amazon.awssdk.http.apache.ApacheHttpClient; import software.amazon.awssdk.http.apache.ProxyConfiguration; import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.opensearchserverless.OpenSearchServerlessClient; import javax.net.ssl.SSLContext; import javax.net.ssl.TrustManager; @@ -86,6 +90,10 @@ public class ConnectionConfiguration { public static final String AWS_STS_HEADER_OVERRIDES = "aws_sts_header_overrides"; public static final String PROXY = "proxy"; public static final String SERVERLESS = "serverless"; + public static final String SERVERLESS_OPTIONS = "serverless_options"; + public static final String COLLECTION_NAME = "collection_name"; + public static final String NETWORK_POLICY_NAME = "network_policy_name"; + public static final String VPCE_ID = "vpce_id"; public static final String REQUEST_COMPRESSION_ENABLED = "enable_request_compression"; /** @@ -109,6 +117,9 @@ public class ConnectionConfiguration { private final Optional proxy; private final String pipelineName; private final boolean serverless; + private final String serverlessNetworkPolicyName; + private final String serverlessCollectionName; + private final String serverlessVpceId; private final boolean requestCompressionEnabled; List getHosts() { @@ -155,6 +166,18 @@ boolean isServerless() { return serverless; } + String getServerlessNetworkPolicyName() { + return serverlessNetworkPolicyName; + } + + String getServerlessCollectionName() { + return serverlessCollectionName; + } + + String getServerlessVpceId() { + return serverlessVpceId; + } + boolean isRequestCompressionEnabled() { return requestCompressionEnabled; } @@ -174,6 +197,9 @@ private ConnectionConfiguration(final Builder builder) { this.awsStsHeaderOverrides = builder.awsStsHeaderOverrides; this.proxy = builder.proxy; this.serverless = builder.serverless; + this.serverlessNetworkPolicyName = builder.serverlessNetworkPolicyName; + this.serverlessCollectionName = builder.serverlessCollectionName; + this.serverlessVpceId = builder.serverlessVpceId; this.requestCompressionEnabled = builder.requestCompressionEnabled; this.pipelineName = builder.pipelineName; } @@ -212,6 +238,13 @@ public static ConnectionConfiguration readConnectionConfiguration(final PluginSe builder.withAwsStsHeaderOverrides((Map)awsOption.get(AWS_STS_HEADER_OVERRIDES.substring(4))); builder.withServerless(OBJECT_MAPPER.convertValue( awsOption.getOrDefault(SERVERLESS, false), Boolean.class)); + + Map serverlessOptions = (Map) awsOption.get(SERVERLESS_OPTIONS); + if (serverlessOptions != null && !serverlessOptions.isEmpty()) { + builder.withServerlessNetworkPolicyName((String)(serverlessOptions.getOrDefault(NETWORK_POLICY_NAME, null))); + builder.withServerlessCollectionName((String)(serverlessOptions.getOrDefault(COLLECTION_NAME, null))); + builder.withServerlessVpceId((String)(serverlessOptions.getOrDefault(VPCE_ID, null))); + } } else { builder.withServerless(false); } @@ -407,6 +440,24 @@ private OpenSearchTransport createOpenSearchTransport(final RestHighLevelClient } } + public OpenSearchServerlessClient createOpenSearchServerlessClient(final AwsCredentialsSupplier awsCredentialsSupplier) { + final AwsCredentialsOptions awsCredentialsOptions = createAwsCredentialsOptions(); + + return OpenSearchServerlessClient.builder() + .credentialsProvider(awsCredentialsSupplier.getProvider(awsCredentialsOptions)) + .region(Region.of(awsRegion)) + .overrideConfiguration(ClientOverrideConfiguration.builder() + .retryPolicy(RetryPolicy.builder() + .backoffStrategy(FullJitterBackoffStrategy.builder() + .baseDelay(Duration.ofSeconds(10)) + .maxBackoffTime(Duration.ofSeconds(60)) + .build()) + .numRetries(10) + .build()) + .build()) + .build(); + } + private SdkHttpClient createSdkHttpClient() { ApacheHttpClient.Builder apacheHttpClientBuilder = ApacheHttpClient.builder(); if (connectTimeout != null) { @@ -475,6 +526,9 @@ public static class Builder { private Optional proxy = Optional.empty(); private String pipelineName; private boolean serverless; + private String serverlessNetworkPolicyName; + private String serverlessCollectionName; + private String serverlessVpceId; private boolean requestCompressionEnabled; private void validateStsRoleArn(final String awsStsRoleArn) { @@ -585,6 +639,21 @@ public Builder withServerless(boolean serverless) { return this; } + public Builder withServerlessNetworkPolicyName(final String serverlessNetworkPolicyName) { + this.serverlessNetworkPolicyName = serverlessNetworkPolicyName; + return this; + } + + public Builder withServerlessCollectionName(final String serverlessCollectionName) { + this.serverlessCollectionName = serverlessCollectionName; + return this; + } + + public Builder withServerlessVpceId(final String serverlessVpceId) { + this.serverlessVpceId = serverlessVpceId; + return this; + } + public Builder withRequestCompressionEnabled(final boolean requestCompressionEnabled) { this.requestCompressionEnabled = requestCompressionEnabled; return this; 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 ea7a46b44c..8820de4afe 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 @@ -61,6 +61,7 @@ import org.opensearch.dataprepper.plugins.sink.opensearch.index.TemplateStrategy; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.opensearchserverless.OpenSearchServerlessClient; import java.io.BufferedWriter; import java.io.IOException; @@ -243,6 +244,9 @@ private void doInitializeInternal() throws IOException { bulkRequestSupplier, pluginSetting); + // Attempt to update the serverless network policy if required argument are given. + maybeUpdateServerlessNetworkPolicy(); + objectMapper = new ObjectMapper(); this.initialized = true; LOG.info("Initialized OpenSearch sink"); @@ -505,4 +509,21 @@ public void shutdown() { super.shutdown(); closeFiles(); } + + private void maybeUpdateServerlessNetworkPolicy() { + final ConnectionConfiguration connectionConfiguration = openSearchSinkConfig.getConnectionConfiguration(); + LOG.info(connectionConfiguration.toString()); + if (connectionConfiguration.isServerless() && + !StringUtils.isBlank(connectionConfiguration.getServerlessNetworkPolicyName()) && + !StringUtils.isBlank(connectionConfiguration.getServerlessCollectionName()) && + !StringUtils.isBlank(connectionConfiguration.getServerlessVpceId()) + ) { + final OpenSearchServerlessClient openSearchServerlessClient = connectionConfiguration.createOpenSearchServerlessClient(awsCredentialsSupplier); + final ServerlessNetworkPolicyUpdater networkPolicyUpdater = new ServerlessNetworkPolicyUpdater(openSearchServerlessClient); + networkPolicyUpdater.updateNetworkPolicy( + connectionConfiguration.getServerlessNetworkPolicyName(), + connectionConfiguration.getServerlessCollectionName(), + connectionConfiguration.getServerlessVpceId()); + } + } } diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/ServerlessNetworkPolicyUpdater.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/ServerlessNetworkPolicyUpdater.java new file mode 100644 index 0000000000..ef427b0026 --- /dev/null +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/ServerlessNetworkPolicyUpdater.java @@ -0,0 +1,159 @@ +package org.opensearch.dataprepper.plugins.sink.opensearch; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.awssdk.core.document.Document; +import software.amazon.awssdk.services.opensearchserverless.OpenSearchServerlessClient; +import software.amazon.awssdk.services.opensearchserverless.model.CreateSecurityPolicyRequest; +import software.amazon.awssdk.services.opensearchserverless.model.GetSecurityPolicyRequest; +import software.amazon.awssdk.services.opensearchserverless.model.GetSecurityPolicyResponse; +import software.amazon.awssdk.services.opensearchserverless.model.ResourceNotFoundException; +import software.amazon.awssdk.services.opensearchserverless.model.SecurityPolicyDetail; +import software.amazon.awssdk.services.opensearchserverless.model.SecurityPolicyType; +import software.amazon.awssdk.services.opensearchserverless.model.UpdateSecurityPolicyRequest; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.regex.Pattern; + +public class ServerlessNetworkPolicyUpdater { + + static final String COLLECTION = "collection"; + static final String CREATED_BY_DATA_PREPPER = "Created by Data Prepper"; + static final String DESCRIPTION = "Description"; + static final String RESOURCE = "Resource"; + static final String RESOURCE_TYPE = "ResourceType"; + static final String RULES = "Rules"; + static final String SOURCE_VPCES = "SourceVPCEs"; + + private static final Logger LOG = LoggerFactory.getLogger(ServerlessNetworkPolicyUpdater.class); + + private final OpenSearchServerlessClient client; + + public ServerlessNetworkPolicyUpdater(OpenSearchServerlessClient client) { + this.client = client; + } + + public void updateNetworkPolicy( + final String networkPolicyName, + final String collectionName, + final String vpceId + ) { + try { + final Document newStatement = createNetworkPolicyStatement(collectionName, vpceId); + final Optional maybeNetworkPolicy = getNetworkPolicy(networkPolicyName); + + if (maybeNetworkPolicy.isPresent()) { + final Document existingPolicy = maybeNetworkPolicy.get().policy(); + final String policyVersion = maybeNetworkPolicy.get().policyVersion(); + final List existingStatements = existingPolicy.asList(); + if (hasAcceptablePolicy(existingStatements, collectionName, vpceId)) { + LOG.info("Policy statement already exists that matches collection and vpce id"); + return; + } + + final List statements = new ArrayList<>(existingStatements); + statements.add(newStatement); + final Document newPolicy = Document.fromList(statements); + updateNetworkPolicy(networkPolicyName, newPolicy, policyVersion); + } else { + final Document newPolicy = Document.fromList(List.of(newStatement)); + createNetworkPolicy(networkPolicyName, newPolicy); + } + } catch (final Exception e) { + LOG.error("Failed to create or update network policy", e); + } + } + + private Optional getNetworkPolicy(final String networkPolicyName) { + // Call the GetSecurityPolicy API + GetSecurityPolicyRequest getRequest = GetSecurityPolicyRequest.builder() + .name(networkPolicyName) + .type(SecurityPolicyType.NETWORK) + .build(); + + GetSecurityPolicyResponse response; + try { + response = client.getSecurityPolicy(getRequest); + } catch (final ResourceNotFoundException e) { + LOG.info("Could not find network policy {}", networkPolicyName); + return Optional.empty(); + } + + if (response.securityPolicyDetail() == null) { + LOG.info("Security policy exists but had no detail."); + return Optional.empty(); + } + + return Optional.of(response.securityPolicyDetail()); + } + + private void createNetworkPolicy(final String networkPolicyName, final Document policy) { + final CreateSecurityPolicyRequest request = CreateSecurityPolicyRequest.builder() + .name(networkPolicyName) + .policy(policy.toString()) + .type(SecurityPolicyType.NETWORK) + .build(); + + client.createSecurityPolicy(request); + } + + private void updateNetworkPolicy(final String networkPolicyName, final Document policy, final String policyVersion) { + final UpdateSecurityPolicyRequest request = UpdateSecurityPolicyRequest.builder() + .name(networkPolicyName) + .policy(policy.toString()) + .type(SecurityPolicyType.NETWORK) + .policyVersion(policyVersion) + .build(); + + client.updateSecurityPolicy(request); + } + + private static Document createNetworkPolicyStatement(final String collectionName, final String vpceId) { + return Document.mapBuilder() + .putString(DESCRIPTION, "Created by Data Prepper") + .putList(RULES, List.of(Document.mapBuilder() + .putString(RESOURCE_TYPE, COLLECTION) + .putList(RESOURCE, List.of(Document.fromString(String.format("%s/%s", COLLECTION, collectionName)))) + .build())) + .putList(SOURCE_VPCES, List.of(Document.fromString(vpceId))) + .build(); + } + + private static boolean hasAcceptablePolicy(final List statements, final String collectionName, final String vpceId) { + for (final Document statement : statements) { + final Map statementFields = statement.asMap(); + if (!statementFields.containsKey(SOURCE_VPCES) || !statementFields.containsKey(RULES)) { + continue; + } + + // Check if the statement has the SourceVPCEs field that matches the given vpceId + boolean hasMatchingVpce = statementFields.get(SOURCE_VPCES).asList().stream() + .map(Document::asString) + .anyMatch(vpce -> vpce.equals(vpceId)); + + // Check if the statement has the Rules field with the ResourceType set to COLLECTION + // that matches (or covers) the given collectionName + boolean hasMatchingCollection = statementFields.get(RULES).asList().stream() + .filter(rule -> rule.asMap().get(RESOURCE_TYPE).asString().equals(COLLECTION)) + .flatMap(rule -> rule.asMap().get(RESOURCE).asList().stream()) + .map(Document::asString) + .anyMatch(collectionPattern -> matchesPattern(collectionPattern, String.format("%s/%s", COLLECTION, collectionName))); + + // If both conditions are met, return true + if (hasMatchingVpce && hasMatchingCollection) { + return true; + } + } + return false; + } + + private static boolean matchesPattern(String pattern, String value) { + // Convert wildcard pattern to regex + String regex = "^" + Pattern.quote(pattern).replace("*", "\\E.*\\Q") + "$"; + return value.matches(regex); + } + +} \ No newline at end of file diff --git a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/ConnectionConfigurationTests.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/ConnectionConfigurationTests.java index d5da3e6d3d..e5503cd9a0 100644 --- a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/ConnectionConfigurationTests.java +++ b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/ConnectionConfigurationTests.java @@ -264,6 +264,36 @@ void testCreateClientWithAWSSigV4AndRegion() throws IOException { assertTrue(connectionConfiguration.isAwsSigv4()); } + @Test + void testServerlessOptions() throws IOException { + final String serverlessNetworkPolicyName = UUID.randomUUID().toString(); + final String serverlessCollectionName = UUID.randomUUID().toString(); + final String serverlessVpceId = UUID.randomUUID().toString(); + + final Map metadata = new HashMap<>(); + final Map awsOptionMetadata = new HashMap<>(); + final Map serverlessOptionsMetadata = new HashMap<>(); + serverlessOptionsMetadata.put("network_policy_name", serverlessNetworkPolicyName); + serverlessOptionsMetadata.put("collection_name", serverlessCollectionName); + serverlessOptionsMetadata.put("vpce_id", serverlessVpceId); + awsOptionMetadata.put("region", UUID.randomUUID().toString()); + awsOptionMetadata.put("serverless", true); + awsOptionMetadata.put("serverless_options", serverlessOptionsMetadata); + awsOptionMetadata.put("sts_role_arn", TEST_ROLE); + metadata.put("hosts", TEST_HOSTS); + metadata.put("username", UUID.randomUUID().toString()); + metadata.put("password", UUID.randomUUID().toString()); + metadata.put("connect_timeout", 1); + metadata.put("socket_timeout", 1); + metadata.put("aws", awsOptionMetadata); + + final PluginSetting pluginSetting = getPluginSettingByConfigurationMetadata(metadata); + final ConnectionConfiguration connectionConfiguration = ConnectionConfiguration.readConnectionConfiguration(pluginSetting); + assertThat(connectionConfiguration.getServerlessNetworkPolicyName(), equalTo(serverlessNetworkPolicyName)); + assertThat(connectionConfiguration.getServerlessCollectionName(), equalTo(serverlessCollectionName)); + assertThat(connectionConfiguration.getServerlessVpceId(), equalTo(serverlessVpceId)); + } + @Test void testCreateClientWithAWSSigV4DefaultRegion() throws IOException { final PluginSetting pluginSetting = generatePluginSetting( diff --git a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/ServerlessNetworkPolicyUpdaterTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/ServerlessNetworkPolicyUpdaterTest.java new file mode 100644 index 0000000000..7fed6b523a --- /dev/null +++ b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/ServerlessNetworkPolicyUpdaterTest.java @@ -0,0 +1,380 @@ +package org.opensearch.dataprepper.plugins.sink.opensearch; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.InjectMocks; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import software.amazon.awssdk.core.document.Document; +import software.amazon.awssdk.services.opensearchserverless.OpenSearchServerlessClient; +import software.amazon.awssdk.services.opensearchserverless.model.CreateSecurityPolicyRequest; +import software.amazon.awssdk.services.opensearchserverless.model.GetSecurityPolicyRequest; +import software.amazon.awssdk.services.opensearchserverless.model.GetSecurityPolicyResponse; +import software.amazon.awssdk.services.opensearchserverless.model.ResourceNotFoundException; +import software.amazon.awssdk.services.opensearchserverless.model.SecurityPolicyDetail; +import software.amazon.awssdk.services.opensearchserverless.model.UpdateSecurityPolicyRequest; + +import java.util.List; +import java.util.UUID; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.plugins.sink.opensearch.ServerlessNetworkPolicyUpdater.COLLECTION; +import static org.opensearch.dataprepper.plugins.sink.opensearch.ServerlessNetworkPolicyUpdater.CREATED_BY_DATA_PREPPER; +import static org.opensearch.dataprepper.plugins.sink.opensearch.ServerlessNetworkPolicyUpdater.DESCRIPTION; +import static org.opensearch.dataprepper.plugins.sink.opensearch.ServerlessNetworkPolicyUpdater.RESOURCE; +import static org.opensearch.dataprepper.plugins.sink.opensearch.ServerlessNetworkPolicyUpdater.RESOURCE_TYPE; +import static org.opensearch.dataprepper.plugins.sink.opensearch.ServerlessNetworkPolicyUpdater.RULES; +import static org.opensearch.dataprepper.plugins.sink.opensearch.ServerlessNetworkPolicyUpdater.SOURCE_VPCES; + +public class ServerlessNetworkPolicyUpdaterTest { + + @Mock + private OpenSearchServerlessClient client; + + @InjectMocks + private ServerlessNetworkPolicyUpdater updater; + + + @BeforeEach + public void setUp() { + MockitoAnnotations.openMocks(this); + } + + @Test + public void testUpdateNetworkPolicyNoExistingPolicy() { + final String networkPolicyName = UUID.randomUUID().toString(); + final String collectionName = UUID.randomUUID().toString(); + final String vpceId = UUID.randomUUID().toString(); + + // Mock client to throw ResourceNotFoundException + when(client.getSecurityPolicy(any(GetSecurityPolicyRequest.class))).thenThrow(ResourceNotFoundException.class); + + updater.updateNetworkPolicy(networkPolicyName, collectionName, vpceId); + + final ArgumentCaptor argumentCaptor = ArgumentCaptor.forClass(CreateSecurityPolicyRequest.class); + + verify(client, never()).updateSecurityPolicy(any(UpdateSecurityPolicyRequest.class)); + verify(client).createSecurityPolicy(argumentCaptor.capture()); + + final Document expectedPolicy = Document.fromList(List.of( + Document.mapBuilder() + .putString(DESCRIPTION, CREATED_BY_DATA_PREPPER) + .putList(RULES, List.of( + Document.mapBuilder() + .putString(RESOURCE_TYPE, COLLECTION) + .putList(RESOURCE, List.of(Document.fromString(COLLECTION + "/" + collectionName))) + .build())) + .putList(SOURCE_VPCES, List.of(Document.fromString(vpceId))) + .build() + )); + + final CreateSecurityPolicyRequest actualRequest = argumentCaptor.getValue(); + final String actualPolicy = actualRequest.policy(); + assertThat(actualPolicy, equalTo(expectedPolicy.toString())); + } + + + @Test + public void testUpdateNetworkPolicyNoPolicyDetail() { + final String networkPolicyName = UUID.randomUUID().toString(); + final String collectionName = UUID.randomUUID().toString(); + final String vpceId = UUID.randomUUID().toString(); + + final GetSecurityPolicyResponse mockedResponse = GetSecurityPolicyResponse.builder().build(); + + // Mocking client behavior + when(client.getSecurityPolicy(any(GetSecurityPolicyRequest.class))).thenReturn(mockedResponse); + + updater.updateNetworkPolicy(networkPolicyName, collectionName, vpceId); + + final ArgumentCaptor argumentCaptor = ArgumentCaptor.forClass(CreateSecurityPolicyRequest.class); + + verify(client, never()).updateSecurityPolicy(any(UpdateSecurityPolicyRequest.class)); + verify(client).createSecurityPolicy(argumentCaptor.capture()); + + final Document expectedPolicy = Document.fromList(List.of( + Document.mapBuilder() + .putString(DESCRIPTION, CREATED_BY_DATA_PREPPER) + .putList(RULES, List.of( + Document.mapBuilder() + .putString(RESOURCE_TYPE, COLLECTION) + .putList(RESOURCE, List.of(Document.fromString(COLLECTION + "/" + collectionName))) + .build())) + .putList(SOURCE_VPCES, List.of(Document.fromString(vpceId))) + .build() + )); + + final CreateSecurityPolicyRequest actualRequest = argumentCaptor.getValue(); + final String actualPolicy = actualRequest.policy(); + assertThat(actualPolicy, equalTo(expectedPolicy.toString())); + } + + @Test + public void testUpdateNetworkPolicyExistingAcceptablePolicyBothConditionsTrueFullMatch() { + final String networkPolicyName = UUID.randomUUID().toString(); + final String collectionName = UUID.randomUUID().toString(); + final String vpceId = UUID.randomUUID().toString(); + + // Mocking an existing policy with acceptable conditions + final Document policy = Document.fromList(List.of( + Document.mapBuilder() + .putList(SOURCE_VPCES, List.of(Document.fromString(vpceId))) + .putList(RULES, List.of( + Document.mapBuilder() + .putString(RESOURCE_TYPE, COLLECTION) + .putList(RESOURCE, List.of(Document.fromString(COLLECTION + "/" + collectionName))) + .build())) + .build() + )); + + final GetSecurityPolicyResponse response = mock(GetSecurityPolicyResponse.class); + when(response.securityPolicyDetail()).thenReturn(SecurityPolicyDetail.builder().policy(policy).build()); + when(client.getSecurityPolicy(any(GetSecurityPolicyRequest.class))).thenReturn(response); + + updater.updateNetworkPolicy(networkPolicyName, collectionName, vpceId); + + verify(client, never()).updateSecurityPolicy(any(UpdateSecurityPolicyRequest.class)); + verify(client, never()).createSecurityPolicy(any(CreateSecurityPolicyRequest.class)); + } + + @Test + public void testUpdateNetworkPolicyExistingAcceptablePolicyBothConditionsTrueWildcardMatch() { + final String networkPolicyName = UUID.randomUUID().toString(); + final String collectionName = UUID.randomUUID().toString(); + final String vpceId = UUID.randomUUID().toString(); + + // Mocking an existing policy with acceptable conditions + final Document policy = Document.fromList(List.of( + Document.mapBuilder() + .putList(SOURCE_VPCES, List.of(Document.fromString(vpceId))) + .putList(RULES, List.of( + Document.mapBuilder() + .putString(RESOURCE_TYPE, COLLECTION) + .putList(RESOURCE, List.of(Document.fromString(COLLECTION + "/*"))) + .build())) + .build() + )); + + final GetSecurityPolicyResponse response = mock(GetSecurityPolicyResponse.class); + when(response.securityPolicyDetail()).thenReturn(SecurityPolicyDetail.builder().policy(policy).build()); + when(client.getSecurityPolicy(any(GetSecurityPolicyRequest.class))).thenReturn(response); + + updater.updateNetworkPolicy(networkPolicyName, collectionName, vpceId); + + verify(client, never()).updateSecurityPolicy(any(UpdateSecurityPolicyRequest.class)); + verify(client, never()).createSecurityPolicy(any(CreateSecurityPolicyRequest.class)); + } + + @Test + public void testUpdateNetworkPolicyExistingUnacceptablePolicyOnlyVpceMatches() { + final String networkPolicyName = UUID.randomUUID().toString(); + final String collectionName = UUID.randomUUID().toString(); + final String vpceId = UUID.randomUUID().toString(); + + // Mocking an existing policy with acceptable conditions + final Document policy = Document.fromList(List.of( + Document.mapBuilder() + .putList(SOURCE_VPCES, List.of(Document.fromString(vpceId))) + .putList(RULES, List.of( + Document.mapBuilder() + .putString(RESOURCE_TYPE, COLLECTION) + .putList(RESOURCE, List.of(Document.fromString(COLLECTION + "/differentCollection"))) + .build())) + .build() + )); + + final GetSecurityPolicyResponse response = mock(GetSecurityPolicyResponse.class); + when(response.securityPolicyDetail()).thenReturn(SecurityPolicyDetail.builder().policy(policy).build()); + when(client.getSecurityPolicy(any(GetSecurityPolicyRequest.class))).thenReturn(response); + + updater.updateNetworkPolicy(networkPolicyName, collectionName, vpceId); + + final ArgumentCaptor argumentCaptor = ArgumentCaptor.forClass(UpdateSecurityPolicyRequest.class); + + verify(client).updateSecurityPolicy(argumentCaptor.capture()); + verify(client, never()).createSecurityPolicy(any(CreateSecurityPolicyRequest.class)); + + final Document expectedPolicy = Document.fromList(List.of( + policy.asList().get(0), + Document.mapBuilder() + .putString(DESCRIPTION, CREATED_BY_DATA_PREPPER) + .putList(RULES, List.of( + Document.mapBuilder() + .putString(RESOURCE_TYPE, COLLECTION) + .putList(RESOURCE, List.of(Document.fromString(COLLECTION + "/" + collectionName))) + .build())) + .putList(SOURCE_VPCES, List.of(Document.fromString(vpceId))) + .build() + )); + + final UpdateSecurityPolicyRequest actualRequest = argumentCaptor.getValue(); + final String actualPolicy = actualRequest.policy(); + + assertThat(actualPolicy, equalTo(expectedPolicy.toString())); + } + + @Test + public void testUpdateNetworkPolicyExistingUnacceptablePolicyOnlyCollectionMatches() { + final String networkPolicyName = UUID.randomUUID().toString(); + final String collectionName = UUID.randomUUID().toString(); + final String vpceId = UUID.randomUUID().toString(); + + // Mocking an existing policy with acceptable conditions + final Document policy = Document.fromList(List.of( + Document.mapBuilder() + .putList(SOURCE_VPCES, List.of(Document.fromString("vpce5678"))) + .putList(RULES, List.of( + Document.mapBuilder() + .putString(RESOURCE_TYPE, COLLECTION) + .putList(RESOURCE, List.of(Document.fromString(COLLECTION + "/" + collectionName))) + .build())) + .build() + )); + + final GetSecurityPolicyResponse response = mock(GetSecurityPolicyResponse.class); + when(response.securityPolicyDetail()).thenReturn(SecurityPolicyDetail.builder().policy(policy).build()); + when(client.getSecurityPolicy(any(GetSecurityPolicyRequest.class))).thenReturn(response); + + updater.updateNetworkPolicy(networkPolicyName, collectionName, vpceId); + + final ArgumentCaptor argumentCaptor = ArgumentCaptor.forClass(UpdateSecurityPolicyRequest.class); + + verify(client).updateSecurityPolicy(argumentCaptor.capture()); + verify(client, never()).createSecurityPolicy(any(CreateSecurityPolicyRequest.class)); + + final Document expectedPolicy = Document.fromList(List.of( + policy.asList().get(0), + Document.mapBuilder() + .putString(DESCRIPTION, CREATED_BY_DATA_PREPPER) + .putList(RULES, List.of( + Document.mapBuilder() + .putString(RESOURCE_TYPE, COLLECTION) + .putList(RESOURCE, List.of(Document.fromString(COLLECTION + "/" + collectionName))) + .build())) + .putList(SOURCE_VPCES, List.of(Document.fromString(vpceId))) + .build() + )); + + final UpdateSecurityPolicyRequest actualRequest = argumentCaptor.getValue(); + final String actualPolicy = actualRequest.policy(); + + assertThat(actualPolicy, equalTo(expectedPolicy.toString())); + } + + @Test + public void testUpdateNetworkPolicyExistingUnacceptablePolicyBothConditionsFalse() { + final String networkPolicyName = UUID.randomUUID().toString(); + final String collectionName = UUID.randomUUID().toString(); + final String vpceId = UUID.randomUUID().toString(); + + // Mocking an existing policy with acceptable conditions + final Document policy = Document.fromList(List.of( + Document.mapBuilder() + .putList(SOURCE_VPCES, List.of(Document.fromString("vpce5678"))) + .putList(RULES, List.of( + Document.mapBuilder() + .putString(RESOURCE_TYPE, COLLECTION) + .putList(RESOURCE, List.of(Document.fromString(COLLECTION + "/differentCollection"))) + .build())) + .build() + )); + + final GetSecurityPolicyResponse response = mock(GetSecurityPolicyResponse.class); + when(response.securityPolicyDetail()).thenReturn(SecurityPolicyDetail.builder().policy(policy).build()); + when(client.getSecurityPolicy(any(GetSecurityPolicyRequest.class))).thenReturn(response); + + updater.updateNetworkPolicy(networkPolicyName, collectionName, vpceId); + + final ArgumentCaptor argumentCaptor = ArgumentCaptor.forClass(UpdateSecurityPolicyRequest.class); + + verify(client).updateSecurityPolicy(argumentCaptor.capture()); + verify(client, never()).createSecurityPolicy(any(CreateSecurityPolicyRequest.class)); + + final Document expectedPolicy = Document.fromList(List.of( + policy.asList().get(0), + Document.mapBuilder() + .putString(DESCRIPTION, CREATED_BY_DATA_PREPPER) + .putList(RULES, List.of( + Document.mapBuilder() + .putString(RESOURCE_TYPE, COLLECTION) + .putList(RESOURCE, List.of(Document.fromString(COLLECTION + "/" + collectionName))) + .build())) + .putList(SOURCE_VPCES, List.of(Document.fromString(vpceId))) + .build() + )); + + final UpdateSecurityPolicyRequest actualRequest = argumentCaptor.getValue(); + final String actualPolicy = actualRequest.policy(); + + assertThat(actualPolicy, equalTo(expectedPolicy.toString())); + } + + @Test + public void testUpdateNetworkPolicyGetExceptionScenario() { + final String networkPolicyName = UUID.randomUUID().toString(); + final String collectionName = UUID.randomUUID().toString(); + final String vpceId = UUID.randomUUID().toString(); + + // Mock client to throw a generic exception + when(client.getSecurityPolicy(any(GetSecurityPolicyRequest.class))).thenThrow(new RuntimeException("Test Exception")); + + updater.updateNetworkPolicy(networkPolicyName, collectionName, vpceId); + + verify(client, never()).updateSecurityPolicy(any(UpdateSecurityPolicyRequest.class)); + verify(client, never()).createSecurityPolicy(any(CreateSecurityPolicyRequest.class)); + } + + @Test + public void testUpdateNetworkPolicyCreateExceptionScenario() { + final String networkPolicyName = UUID.randomUUID().toString(); + final String collectionName = UUID.randomUUID().toString(); + final String vpceId = UUID.randomUUID().toString(); + + when(client.getSecurityPolicy(any(GetSecurityPolicyRequest.class))).thenThrow(ResourceNotFoundException.class); + + when(client.createSecurityPolicy(any(CreateSecurityPolicyRequest.class))) + .thenThrow(new RuntimeException("Test Exception")); + + updater.updateNetworkPolicy(networkPolicyName, collectionName, vpceId); + + verify(client, never()).updateSecurityPolicy(any(UpdateSecurityPolicyRequest.class)); + verify(client).createSecurityPolicy(any(CreateSecurityPolicyRequest.class)); + } + + @Test + public void testUpdateNetworkPolicyUpdateExceptionScenario() { + final String networkPolicyName = UUID.randomUUID().toString(); + final String collectionName = UUID.randomUUID().toString(); + final String vpceId = UUID.randomUUID().toString(); + + // Mocking an existing policy with acceptable conditions + final Document policy = Document.fromList(List.of( + Document.mapBuilder() + .putList(RULES, List.of( + Document.mapBuilder() + .putString(RESOURCE_TYPE, COLLECTION) + .putList(RESOURCE, List.of(Document.fromString(COLLECTION + "/" + collectionName))) + .build())) + .build() + )); + + final GetSecurityPolicyResponse response = mock(GetSecurityPolicyResponse.class); + when(response.securityPolicyDetail()).thenReturn(SecurityPolicyDetail.builder().policy(policy).build()); + when(client.getSecurityPolicy(any(GetSecurityPolicyRequest.class))).thenReturn(response); + + when(client.updateSecurityPolicy(any(UpdateSecurityPolicyRequest.class))) + .thenThrow(new RuntimeException("Test exception")); + + updater.updateNetworkPolicy(networkPolicyName, collectionName, vpceId); + + verify(client).updateSecurityPolicy(any(UpdateSecurityPolicyRequest.class)); + verify(client, never()).createSecurityPolicy(any(CreateSecurityPolicyRequest.class)); + } +} From 4acab8409c523998ba57d13c6f6d05caa1507902 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 23 Oct 2023 12:13:44 -0700 Subject: [PATCH 02/53] Bump @babel/traverse in /release/staging-resources-cdk (#3521) Bumps [@babel/traverse](https://github.com/babel/babel/tree/HEAD/packages/babel-traverse) from 7.22.5 to 7.23.2. - [Release notes](https://github.com/babel/babel/releases) - [Changelog](https://github.com/babel/babel/blob/main/CHANGELOG.md) - [Commits](https://github.com/babel/babel/commits/v7.23.2/packages/babel-traverse) --- updated-dependencies: - dependency-name: "@babel/traverse" dependency-type: indirect ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .../staging-resources-cdk/package-lock.json | 331 ++++++++++++------ 1 file changed, 231 insertions(+), 100 deletions(-) diff --git a/release/staging-resources-cdk/package-lock.json b/release/staging-resources-cdk/package-lock.json index 83b8d7a0bd..7ac1eaed21 100644 --- a/release/staging-resources-cdk/package-lock.json +++ b/release/staging-resources-cdk/package-lock.json @@ -57,17 +57,89 @@ "integrity": "sha512-bsyLQD/vqXQcc9RDmlM1XqiFNO/yewgVFXmkMcQkndJbmE/jgYkzewwYGrBlfL725hGLQipXq19+jwWwdsXQqg==" }, "node_modules/@babel/code-frame": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/code-frame/-/code-frame-7.22.5.tgz", - "integrity": "sha512-Xmwn266vad+6DAqEB2A6V/CcZVp62BbwVmcOJc2RPuwih1kw02TjQvWVWlcKGbBPd+8/0V5DEkOcizRGYsspYQ==", + "version": "7.22.13", + "resolved": "https://registry.npmjs.org/@babel/code-frame/-/code-frame-7.22.13.tgz", + "integrity": "sha512-XktuhWlJ5g+3TJXc5upd9Ks1HutSArik6jf2eAjYFyIOf4ej3RN+184cZbzDvbPnuTJIUhPKKJE3cIsYTiAT3w==", "dev": true, "dependencies": { - "@babel/highlight": "^7.22.5" + "@babel/highlight": "^7.22.13", + "chalk": "^2.4.2" }, "engines": { "node": ">=6.9.0" } }, + "node_modules/@babel/code-frame/node_modules/ansi-styles": { + "version": "3.2.1", + "resolved": "https://registry.npmjs.org/ansi-styles/-/ansi-styles-3.2.1.tgz", + "integrity": "sha512-VT0ZI6kZRdTh8YyJw3SMbYm/u+NqfsAxEpWO0Pf9sq8/e94WxxOpPKx9FR1FlyCtOVDNOQ+8ntlqFxiRc+r5qA==", + "dev": true, + "dependencies": { + "color-convert": "^1.9.0" + }, + "engines": { + "node": ">=4" + } + }, + "node_modules/@babel/code-frame/node_modules/chalk": { + "version": "2.4.2", + "resolved": "https://registry.npmjs.org/chalk/-/chalk-2.4.2.tgz", + "integrity": "sha512-Mti+f9lpJNcwF4tWV8/OrTTtF1gZi+f8FqlyAdouralcFWFQWF2+NgCHShjkCb+IFBLq9buZwE1xckQU4peSuQ==", + "dev": true, + "dependencies": { + "ansi-styles": "^3.2.1", + "escape-string-regexp": "^1.0.5", + "supports-color": "^5.3.0" + }, + "engines": { + "node": ">=4" + } + }, + "node_modules/@babel/code-frame/node_modules/color-convert": { + "version": "1.9.3", + "resolved": "https://registry.npmjs.org/color-convert/-/color-convert-1.9.3.tgz", + "integrity": "sha512-QfAUtd+vFdAtFQcC8CCyYt1fYWxSqAiK2cSD6zDB8N3cpsEBAvRxp9zOGg6G/SHHJYAT88/az/IuDGALsNVbGg==", + "dev": true, + "dependencies": { + "color-name": "1.1.3" + } + }, + "node_modules/@babel/code-frame/node_modules/color-name": { + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/color-name/-/color-name-1.1.3.tgz", + "integrity": "sha512-72fSenhMw2HZMTVHeCA9KCmpEIbzWiQsjN+BHcBbS9vr1mtt+vJjPdksIBNUmKAW8TFUDPJK5SUU3QhE9NEXDw==", + "dev": true + }, + "node_modules/@babel/code-frame/node_modules/escape-string-regexp": { + "version": "1.0.5", + "resolved": "https://registry.npmjs.org/escape-string-regexp/-/escape-string-regexp-1.0.5.tgz", + "integrity": "sha512-vbRorB5FUQWvla16U8R/qgaFIya2qGzwDrNmCZuYKrbdSUMG6I1ZCGQRefkRVhuOkIGVne7BQ35DSfo1qvJqFg==", + "dev": true, + "engines": { + "node": ">=0.8.0" + } + }, + "node_modules/@babel/code-frame/node_modules/has-flag": { + "version": "3.0.0", + "resolved": "https://registry.npmjs.org/has-flag/-/has-flag-3.0.0.tgz", + "integrity": "sha512-sKJf1+ceQBr4SMkvQnBDNDtf4TXpVhVGateu0t918bl30FnbE2m4vNLX+VWe/dpjlb+HugGYzW7uQXH98HPEYw==", + "dev": true, + "engines": { + "node": ">=4" + } + }, + "node_modules/@babel/code-frame/node_modules/supports-color": { + "version": "5.5.0", + "resolved": "https://registry.npmjs.org/supports-color/-/supports-color-5.5.0.tgz", + "integrity": "sha512-QjVjwdXIt408MIiAqCX4oUKsgU2EqAGzs2Ppkm4aQYbjm+ZEWEcW4SfFNTr4uMNZma0ey4f5lgLrkB0aX0QMow==", + "dev": true, + "dependencies": { + "has-flag": "^3.0.0" + }, + "engines": { + "node": ">=4" + } + }, "node_modules/@babel/compat-data": { "version": "7.22.5", "resolved": "https://registry.npmjs.org/@babel/compat-data/-/compat-data-7.22.5.tgz", @@ -117,12 +189,12 @@ } }, "node_modules/@babel/generator": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/generator/-/generator-7.22.5.tgz", - "integrity": "sha512-+lcUbnTRhd0jOewtFSedLyiPsD5tswKkbgcezOqqWFUVNEwoUTlpPOBmvhG7OXWLR4jMdv0czPGH5XbflnD1EA==", + "version": "7.23.0", + "resolved": "https://registry.npmjs.org/@babel/generator/-/generator-7.23.0.tgz", + "integrity": "sha512-lN85QRR+5IbYrMWM6Y4pE/noaQtg4pNiqeNGX60eqOfo6gtEj6uw/JagelB8vVztSd7R6M5n1+PQkDbHbBRU4g==", "dev": true, "dependencies": { - "@babel/types": "^7.22.5", + "@babel/types": "^7.23.0", "@jridgewell/gen-mapping": "^0.3.2", "@jridgewell/trace-mapping": "^0.3.17", "jsesc": "^2.5.1" @@ -174,22 +246,22 @@ } }, "node_modules/@babel/helper-environment-visitor": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/helper-environment-visitor/-/helper-environment-visitor-7.22.5.tgz", - "integrity": "sha512-XGmhECfVA/5sAt+H+xpSg0mfrHq6FzNr9Oxh7PSEBBRUb/mL7Kz3NICXb194rCqAEdxkhPT1a88teizAFyvk8Q==", + "version": "7.22.20", + "resolved": "https://registry.npmjs.org/@babel/helper-environment-visitor/-/helper-environment-visitor-7.22.20.tgz", + "integrity": "sha512-zfedSIzFhat/gFhWfHtgWvlec0nqB9YEIVrpuwjruLlXfUSnA8cJB0miHKwqDnQ7d32aKo2xt88/xZptwxbfhA==", "dev": true, "engines": { "node": ">=6.9.0" } }, "node_modules/@babel/helper-function-name": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/helper-function-name/-/helper-function-name-7.22.5.tgz", - "integrity": "sha512-wtHSq6jMRE3uF2otvfuD3DIvVhOsSNshQl0Qrd7qC9oQJzHvOL4qQXlQn2916+CXGywIjpGuIkoyZRRxHPiNQQ==", + "version": "7.23.0", + "resolved": "https://registry.npmjs.org/@babel/helper-function-name/-/helper-function-name-7.23.0.tgz", + "integrity": "sha512-OErEqsrxjZTJciZ4Oo+eoZqeW9UIiOcuYKRJA4ZAgV9myA+pOXhhmpfNCKjEH/auVfEYVFJ6y1Tc4r0eIApqiw==", "dev": true, "dependencies": { - "@babel/template": "^7.22.5", - "@babel/types": "^7.22.5" + "@babel/template": "^7.22.15", + "@babel/types": "^7.23.0" }, "engines": { "node": ">=6.9.0" @@ -260,9 +332,9 @@ } }, "node_modules/@babel/helper-split-export-declaration": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/helper-split-export-declaration/-/helper-split-export-declaration-7.22.5.tgz", - "integrity": "sha512-thqK5QFghPKWLhAV321lxF95yCg2K3Ob5yw+M3VHWfdia0IkPXUtoLH8x/6Fh486QUvzhb8YOWHChTVen2/PoQ==", + "version": "7.22.6", + "resolved": "https://registry.npmjs.org/@babel/helper-split-export-declaration/-/helper-split-export-declaration-7.22.6.tgz", + "integrity": "sha512-AsUnxuLhRYsisFiaJwvp1QF+I3KjD5FOxut14q/GzovUe6orHLesW2C7d754kRm53h5gqrz6sFl6sxc4BVtE/g==", "dev": true, "dependencies": { "@babel/types": "^7.22.5" @@ -281,9 +353,9 @@ } }, "node_modules/@babel/helper-validator-identifier": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/helper-validator-identifier/-/helper-validator-identifier-7.22.5.tgz", - "integrity": "sha512-aJXu+6lErq8ltp+JhkJUfk1MTGyuA4v7f3pA+BJ5HLfNC6nAQ0Cpi9uOquUj8Hehg0aUiHzWQbOVJGao6ztBAQ==", + "version": "7.22.20", + "resolved": "https://registry.npmjs.org/@babel/helper-validator-identifier/-/helper-validator-identifier-7.22.20.tgz", + "integrity": "sha512-Y4OZ+ytlatR8AI+8KZfKuL5urKp7qey08ha31L8b3BwewJAoJamTzyvxPR/5D+KkdJCGPq/+8TukHBlY10FX9A==", "dev": true, "engines": { "node": ">=6.9.0" @@ -313,13 +385,13 @@ } }, "node_modules/@babel/highlight": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/highlight/-/highlight-7.22.5.tgz", - "integrity": "sha512-BSKlD1hgnedS5XRnGOljZawtag7H1yPfQp0tdNJCHoH6AZ+Pcm9VvkrK59/Yy593Ypg0zMxH2BxD1VPYUQ7UIw==", + "version": "7.22.20", + "resolved": "https://registry.npmjs.org/@babel/highlight/-/highlight-7.22.20.tgz", + "integrity": "sha512-dkdMCN3py0+ksCgYmGG8jKeGA/8Tk+gJwSYYlFGxG5lmhfKNoAy004YpLxpS1W2J8m/EK2Ew+yOs9pVRwO89mg==", "dev": true, "dependencies": { - "@babel/helper-validator-identifier": "^7.22.5", - "chalk": "^2.0.0", + "@babel/helper-validator-identifier": "^7.22.20", + "chalk": "^2.4.2", "js-tokens": "^4.0.0" }, "engines": { @@ -398,9 +470,9 @@ } }, "node_modules/@babel/parser": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/parser/-/parser-7.22.5.tgz", - "integrity": "sha512-DFZMC9LJUG9PLOclRC32G63UXwzqS2koQC8dkx+PLdmt1xSePYpbT/NbsrJy8Q/muXz7o/h/d4A7Fuyixm559Q==", + "version": "7.23.0", + "resolved": "https://registry.npmjs.org/@babel/parser/-/parser-7.23.0.tgz", + "integrity": "sha512-vvPKKdMemU85V9WE/l5wZEmImpCtLqbnTvqDS2U1fJ96KrxoW7KrXhNsNCblQlg8Ck4b85yxdTyelsMUgFUXiw==", "dev": true, "bin": { "parser": "bin/babel-parser.js" @@ -557,33 +629,33 @@ } }, "node_modules/@babel/template": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/template/-/template-7.22.5.tgz", - "integrity": "sha512-X7yV7eiwAxdj9k94NEylvbVHLiVG1nvzCV2EAowhxLTwODV1jl9UzZ48leOC0sH7OnuHrIkllaBgneUykIcZaw==", + "version": "7.22.15", + "resolved": "https://registry.npmjs.org/@babel/template/-/template-7.22.15.tgz", + "integrity": "sha512-QPErUVm4uyJa60rkI73qneDacvdvzxshT3kksGqlGWYdOTIUOwJ7RDUL8sGqslY1uXWSL6xMFKEXDS3ox2uF0w==", "dev": true, "dependencies": { - "@babel/code-frame": "^7.22.5", - "@babel/parser": "^7.22.5", - "@babel/types": "^7.22.5" + "@babel/code-frame": "^7.22.13", + "@babel/parser": "^7.22.15", + "@babel/types": "^7.22.15" }, "engines": { "node": ">=6.9.0" } }, "node_modules/@babel/traverse": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/traverse/-/traverse-7.22.5.tgz", - "integrity": "sha512-7DuIjPgERaNo6r+PZwItpjCZEa5vyw4eJGufeLxrPdBXBoLcCJCIasvK6pK/9DVNrLZTLFhUGqaC6X/PA007TQ==", + "version": "7.23.2", + "resolved": "https://registry.npmjs.org/@babel/traverse/-/traverse-7.23.2.tgz", + "integrity": "sha512-azpe59SQ48qG6nu2CzcMLbxUudtN+dOM9kDbUqGq3HXUJRlo7i8fvPoxQUzYgLZ4cMVmuZgm8vvBpNeRhd6XSw==", "dev": true, "dependencies": { - "@babel/code-frame": "^7.22.5", - "@babel/generator": "^7.22.5", - "@babel/helper-environment-visitor": "^7.22.5", - "@babel/helper-function-name": "^7.22.5", + "@babel/code-frame": "^7.22.13", + "@babel/generator": "^7.23.0", + "@babel/helper-environment-visitor": "^7.22.20", + "@babel/helper-function-name": "^7.23.0", "@babel/helper-hoist-variables": "^7.22.5", - "@babel/helper-split-export-declaration": "^7.22.5", - "@babel/parser": "^7.22.5", - "@babel/types": "^7.22.5", + "@babel/helper-split-export-declaration": "^7.22.6", + "@babel/parser": "^7.23.0", + "@babel/types": "^7.23.0", "debug": "^4.1.0", "globals": "^11.1.0" }, @@ -601,13 +673,13 @@ } }, "node_modules/@babel/types": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/types/-/types-7.22.5.tgz", - "integrity": "sha512-zo3MIHGOkPOfoRXitsgHLjEXmlDaD/5KU1Uzuc9GNiZPhSqVxVRtxuPaSBZDsYZ9qV88AjtMtWW7ww98loJ9KA==", + "version": "7.23.0", + "resolved": "https://registry.npmjs.org/@babel/types/-/types-7.23.0.tgz", + "integrity": "sha512-0oIyUfKoI3mSqMvsxBdclDwxXKXAUA8v/apZbc+iSyARYou1o8ZGDxbUYyLFoW2arqS2jDGqJuZvv1d/io1axg==", "dev": true, "dependencies": { "@babel/helper-string-parser": "^7.22.5", - "@babel/helper-validator-identifier": "^7.22.5", + "@babel/helper-validator-identifier": "^7.22.20", "to-fast-properties": "^2.0.0" }, "engines": { @@ -7837,12 +7909,71 @@ "integrity": "sha512-bsyLQD/vqXQcc9RDmlM1XqiFNO/yewgVFXmkMcQkndJbmE/jgYkzewwYGrBlfL725hGLQipXq19+jwWwdsXQqg==" }, "@babel/code-frame": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/code-frame/-/code-frame-7.22.5.tgz", - "integrity": "sha512-Xmwn266vad+6DAqEB2A6V/CcZVp62BbwVmcOJc2RPuwih1kw02TjQvWVWlcKGbBPd+8/0V5DEkOcizRGYsspYQ==", + "version": "7.22.13", + "resolved": "https://registry.npmjs.org/@babel/code-frame/-/code-frame-7.22.13.tgz", + "integrity": "sha512-XktuhWlJ5g+3TJXc5upd9Ks1HutSArik6jf2eAjYFyIOf4ej3RN+184cZbzDvbPnuTJIUhPKKJE3cIsYTiAT3w==", "dev": true, "requires": { - "@babel/highlight": "^7.22.5" + "@babel/highlight": "^7.22.13", + "chalk": "^2.4.2" + }, + "dependencies": { + "ansi-styles": { + "version": "3.2.1", + "resolved": "https://registry.npmjs.org/ansi-styles/-/ansi-styles-3.2.1.tgz", + "integrity": "sha512-VT0ZI6kZRdTh8YyJw3SMbYm/u+NqfsAxEpWO0Pf9sq8/e94WxxOpPKx9FR1FlyCtOVDNOQ+8ntlqFxiRc+r5qA==", + "dev": true, + "requires": { + "color-convert": "^1.9.0" + } + }, + "chalk": { + "version": "2.4.2", + "resolved": "https://registry.npmjs.org/chalk/-/chalk-2.4.2.tgz", + "integrity": "sha512-Mti+f9lpJNcwF4tWV8/OrTTtF1gZi+f8FqlyAdouralcFWFQWF2+NgCHShjkCb+IFBLq9buZwE1xckQU4peSuQ==", + "dev": true, + "requires": { + "ansi-styles": "^3.2.1", + "escape-string-regexp": "^1.0.5", + "supports-color": "^5.3.0" + } + }, + "color-convert": { + "version": "1.9.3", + "resolved": "https://registry.npmjs.org/color-convert/-/color-convert-1.9.3.tgz", + "integrity": "sha512-QfAUtd+vFdAtFQcC8CCyYt1fYWxSqAiK2cSD6zDB8N3cpsEBAvRxp9zOGg6G/SHHJYAT88/az/IuDGALsNVbGg==", + "dev": true, + "requires": { + "color-name": "1.1.3" + } + }, + "color-name": { + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/color-name/-/color-name-1.1.3.tgz", + "integrity": "sha512-72fSenhMw2HZMTVHeCA9KCmpEIbzWiQsjN+BHcBbS9vr1mtt+vJjPdksIBNUmKAW8TFUDPJK5SUU3QhE9NEXDw==", + "dev": true + }, + "escape-string-regexp": { + "version": "1.0.5", + "resolved": "https://registry.npmjs.org/escape-string-regexp/-/escape-string-regexp-1.0.5.tgz", + "integrity": "sha512-vbRorB5FUQWvla16U8R/qgaFIya2qGzwDrNmCZuYKrbdSUMG6I1ZCGQRefkRVhuOkIGVne7BQ35DSfo1qvJqFg==", + "dev": true + }, + "has-flag": { + "version": "3.0.0", + "resolved": "https://registry.npmjs.org/has-flag/-/has-flag-3.0.0.tgz", + "integrity": "sha512-sKJf1+ceQBr4SMkvQnBDNDtf4TXpVhVGateu0t918bl30FnbE2m4vNLX+VWe/dpjlb+HugGYzW7uQXH98HPEYw==", + "dev": true + }, + "supports-color": { + "version": "5.5.0", + "resolved": "https://registry.npmjs.org/supports-color/-/supports-color-5.5.0.tgz", + "integrity": "sha512-QjVjwdXIt408MIiAqCX4oUKsgU2EqAGzs2Ppkm4aQYbjm+ZEWEcW4SfFNTr4uMNZma0ey4f5lgLrkB0aX0QMow==", + "dev": true, + "requires": { + "has-flag": "^3.0.0" + } + } } }, "@babel/compat-data": { @@ -7883,12 +8014,12 @@ } }, "@babel/generator": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/generator/-/generator-7.22.5.tgz", - "integrity": "sha512-+lcUbnTRhd0jOewtFSedLyiPsD5tswKkbgcezOqqWFUVNEwoUTlpPOBmvhG7OXWLR4jMdv0czPGH5XbflnD1EA==", + "version": "7.23.0", + "resolved": "https://registry.npmjs.org/@babel/generator/-/generator-7.23.0.tgz", + "integrity": "sha512-lN85QRR+5IbYrMWM6Y4pE/noaQtg4pNiqeNGX60eqOfo6gtEj6uw/JagelB8vVztSd7R6M5n1+PQkDbHbBRU4g==", "dev": true, "requires": { - "@babel/types": "^7.22.5", + "@babel/types": "^7.23.0", "@jridgewell/gen-mapping": "^0.3.2", "@jridgewell/trace-mapping": "^0.3.17", "jsesc": "^2.5.1" @@ -7929,19 +8060,19 @@ } }, "@babel/helper-environment-visitor": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/helper-environment-visitor/-/helper-environment-visitor-7.22.5.tgz", - "integrity": "sha512-XGmhECfVA/5sAt+H+xpSg0mfrHq6FzNr9Oxh7PSEBBRUb/mL7Kz3NICXb194rCqAEdxkhPT1a88teizAFyvk8Q==", + "version": "7.22.20", + "resolved": "https://registry.npmjs.org/@babel/helper-environment-visitor/-/helper-environment-visitor-7.22.20.tgz", + "integrity": "sha512-zfedSIzFhat/gFhWfHtgWvlec0nqB9YEIVrpuwjruLlXfUSnA8cJB0miHKwqDnQ7d32aKo2xt88/xZptwxbfhA==", "dev": true }, "@babel/helper-function-name": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/helper-function-name/-/helper-function-name-7.22.5.tgz", - "integrity": "sha512-wtHSq6jMRE3uF2otvfuD3DIvVhOsSNshQl0Qrd7qC9oQJzHvOL4qQXlQn2916+CXGywIjpGuIkoyZRRxHPiNQQ==", + "version": "7.23.0", + "resolved": "https://registry.npmjs.org/@babel/helper-function-name/-/helper-function-name-7.23.0.tgz", + "integrity": "sha512-OErEqsrxjZTJciZ4Oo+eoZqeW9UIiOcuYKRJA4ZAgV9myA+pOXhhmpfNCKjEH/auVfEYVFJ6y1Tc4r0eIApqiw==", "dev": true, "requires": { - "@babel/template": "^7.22.5", - "@babel/types": "^7.22.5" + "@babel/template": "^7.22.15", + "@babel/types": "^7.23.0" } }, "@babel/helper-hoist-variables": { @@ -7994,9 +8125,9 @@ } }, "@babel/helper-split-export-declaration": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/helper-split-export-declaration/-/helper-split-export-declaration-7.22.5.tgz", - "integrity": "sha512-thqK5QFghPKWLhAV321lxF95yCg2K3Ob5yw+M3VHWfdia0IkPXUtoLH8x/6Fh486QUvzhb8YOWHChTVen2/PoQ==", + "version": "7.22.6", + "resolved": "https://registry.npmjs.org/@babel/helper-split-export-declaration/-/helper-split-export-declaration-7.22.6.tgz", + "integrity": "sha512-AsUnxuLhRYsisFiaJwvp1QF+I3KjD5FOxut14q/GzovUe6orHLesW2C7d754kRm53h5gqrz6sFl6sxc4BVtE/g==", "dev": true, "requires": { "@babel/types": "^7.22.5" @@ -8009,9 +8140,9 @@ "dev": true }, "@babel/helper-validator-identifier": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/helper-validator-identifier/-/helper-validator-identifier-7.22.5.tgz", - "integrity": "sha512-aJXu+6lErq8ltp+JhkJUfk1MTGyuA4v7f3pA+BJ5HLfNC6nAQ0Cpi9uOquUj8Hehg0aUiHzWQbOVJGao6ztBAQ==", + "version": "7.22.20", + "resolved": "https://registry.npmjs.org/@babel/helper-validator-identifier/-/helper-validator-identifier-7.22.20.tgz", + "integrity": "sha512-Y4OZ+ytlatR8AI+8KZfKuL5urKp7qey08ha31L8b3BwewJAoJamTzyvxPR/5D+KkdJCGPq/+8TukHBlY10FX9A==", "dev": true }, "@babel/helper-validator-option": { @@ -8032,13 +8163,13 @@ } }, "@babel/highlight": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/highlight/-/highlight-7.22.5.tgz", - "integrity": "sha512-BSKlD1hgnedS5XRnGOljZawtag7H1yPfQp0tdNJCHoH6AZ+Pcm9VvkrK59/Yy593Ypg0zMxH2BxD1VPYUQ7UIw==", + "version": "7.22.20", + "resolved": "https://registry.npmjs.org/@babel/highlight/-/highlight-7.22.20.tgz", + "integrity": "sha512-dkdMCN3py0+ksCgYmGG8jKeGA/8Tk+gJwSYYlFGxG5lmhfKNoAy004YpLxpS1W2J8m/EK2Ew+yOs9pVRwO89mg==", "dev": true, "requires": { - "@babel/helper-validator-identifier": "^7.22.5", - "chalk": "^2.0.0", + "@babel/helper-validator-identifier": "^7.22.20", + "chalk": "^2.4.2", "js-tokens": "^4.0.0" }, "dependencies": { @@ -8101,9 +8232,9 @@ } }, "@babel/parser": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/parser/-/parser-7.22.5.tgz", - "integrity": "sha512-DFZMC9LJUG9PLOclRC32G63UXwzqS2koQC8dkx+PLdmt1xSePYpbT/NbsrJy8Q/muXz7o/h/d4A7Fuyixm559Q==", + "version": "7.23.0", + "resolved": "https://registry.npmjs.org/@babel/parser/-/parser-7.23.0.tgz", + "integrity": "sha512-vvPKKdMemU85V9WE/l5wZEmImpCtLqbnTvqDS2U1fJ96KrxoW7KrXhNsNCblQlg8Ck4b85yxdTyelsMUgFUXiw==", "dev": true }, "@babel/plugin-syntax-async-generators": { @@ -8215,30 +8346,30 @@ } }, "@babel/template": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/template/-/template-7.22.5.tgz", - "integrity": "sha512-X7yV7eiwAxdj9k94NEylvbVHLiVG1nvzCV2EAowhxLTwODV1jl9UzZ48leOC0sH7OnuHrIkllaBgneUykIcZaw==", + "version": "7.22.15", + "resolved": "https://registry.npmjs.org/@babel/template/-/template-7.22.15.tgz", + "integrity": "sha512-QPErUVm4uyJa60rkI73qneDacvdvzxshT3kksGqlGWYdOTIUOwJ7RDUL8sGqslY1uXWSL6xMFKEXDS3ox2uF0w==", "dev": true, "requires": { - "@babel/code-frame": "^7.22.5", - "@babel/parser": "^7.22.5", - "@babel/types": "^7.22.5" + "@babel/code-frame": "^7.22.13", + "@babel/parser": "^7.22.15", + "@babel/types": "^7.22.15" } }, "@babel/traverse": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/traverse/-/traverse-7.22.5.tgz", - "integrity": "sha512-7DuIjPgERaNo6r+PZwItpjCZEa5vyw4eJGufeLxrPdBXBoLcCJCIasvK6pK/9DVNrLZTLFhUGqaC6X/PA007TQ==", + "version": "7.23.2", + "resolved": "https://registry.npmjs.org/@babel/traverse/-/traverse-7.23.2.tgz", + "integrity": "sha512-azpe59SQ48qG6nu2CzcMLbxUudtN+dOM9kDbUqGq3HXUJRlo7i8fvPoxQUzYgLZ4cMVmuZgm8vvBpNeRhd6XSw==", "dev": true, "requires": { - "@babel/code-frame": "^7.22.5", - "@babel/generator": "^7.22.5", - "@babel/helper-environment-visitor": "^7.22.5", - "@babel/helper-function-name": "^7.22.5", + "@babel/code-frame": "^7.22.13", + "@babel/generator": "^7.23.0", + "@babel/helper-environment-visitor": "^7.22.20", + "@babel/helper-function-name": "^7.23.0", "@babel/helper-hoist-variables": "^7.22.5", - "@babel/helper-split-export-declaration": "^7.22.5", - "@babel/parser": "^7.22.5", - "@babel/types": "^7.22.5", + "@babel/helper-split-export-declaration": "^7.22.6", + "@babel/parser": "^7.23.0", + "@babel/types": "^7.23.0", "debug": "^4.1.0", "globals": "^11.1.0" }, @@ -8252,13 +8383,13 @@ } }, "@babel/types": { - "version": "7.22.5", - "resolved": "https://registry.npmjs.org/@babel/types/-/types-7.22.5.tgz", - "integrity": "sha512-zo3MIHGOkPOfoRXitsgHLjEXmlDaD/5KU1Uzuc9GNiZPhSqVxVRtxuPaSBZDsYZ9qV88AjtMtWW7ww98loJ9KA==", + "version": "7.23.0", + "resolved": "https://registry.npmjs.org/@babel/types/-/types-7.23.0.tgz", + "integrity": "sha512-0oIyUfKoI3mSqMvsxBdclDwxXKXAUA8v/apZbc+iSyARYou1o8ZGDxbUYyLFoW2arqS2jDGqJuZvv1d/io1axg==", "dev": true, "requires": { "@babel/helper-string-parser": "^7.22.5", - "@babel/helper-validator-identifier": "^7.22.5", + "@babel/helper-validator-identifier": "^7.22.20", "to-fast-properties": "^2.0.0" } }, From 4c3827ccff38d7f3838d34354308466e540ffce9 Mon Sep 17 00:00:00 2001 From: David Venable Date: Mon, 23 Oct 2023 14:50:56 -0700 Subject: [PATCH 03/53] Run the end-to-end tests on Java 21 in the GitHub Actions. (#3523) Signed-off-by: David Venable --- .../data-prepper-log-analytics-basic-grok-e2e-tests.yml | 2 +- .../data-prepper-peer-forwarder-local-node-e2e-tests.yml | 2 +- .../workflows/data-prepper-peer-forwarder-static-e2e-tests.yml | 2 +- ...prepper-trace-analytics-raw-span-compatibility-e2e-tests.yml | 2 +- .../data-prepper-trace-analytics-raw-span-e2e-tests.yml | 2 +- ...repper-trace-analytics-raw-span-peer-forwarder-e2e-tests.yml | 2 +- .../data-prepper-trace-analytics-service-map-e2e-tests.yml | 2 +- 7 files changed, 7 insertions(+), 7 deletions(-) diff --git a/.github/workflows/data-prepper-log-analytics-basic-grok-e2e-tests.yml b/.github/workflows/data-prepper-log-analytics-basic-grok-e2e-tests.yml index 71405b8ee4..5449ed7f6c 100644 --- a/.github/workflows/data-prepper-log-analytics-basic-grok-e2e-tests.yml +++ b/.github/workflows/data-prepper-log-analytics-basic-grok-e2e-tests.yml @@ -13,7 +13,7 @@ jobs: build: strategy: matrix: - java: [11, 17] + java: [11, 17, 21] test: ['basicLogEndToEndTest', 'parallelGrokStringSubstituteTest'] fail-fast: false diff --git a/.github/workflows/data-prepper-peer-forwarder-local-node-e2e-tests.yml b/.github/workflows/data-prepper-peer-forwarder-local-node-e2e-tests.yml index a0f7ece568..de36390665 100644 --- a/.github/workflows/data-prepper-peer-forwarder-local-node-e2e-tests.yml +++ b/.github/workflows/data-prepper-peer-forwarder-local-node-e2e-tests.yml @@ -13,7 +13,7 @@ jobs: build: strategy: matrix: - java: [11, 17] + java: [11, 17, 21] fail-fast: false runs-on: ubuntu-latest diff --git a/.github/workflows/data-prepper-peer-forwarder-static-e2e-tests.yml b/.github/workflows/data-prepper-peer-forwarder-static-e2e-tests.yml index 06623cba3a..10efb39f01 100644 --- a/.github/workflows/data-prepper-peer-forwarder-static-e2e-tests.yml +++ b/.github/workflows/data-prepper-peer-forwarder-static-e2e-tests.yml @@ -13,7 +13,7 @@ jobs: build: strategy: matrix: - java: [11, 17] + java: [11, 17, 21] test: ['staticAggregateEndToEndTest', 'staticLogMetricsEndToEndTest'] fail-fast: false diff --git a/.github/workflows/data-prepper-trace-analytics-raw-span-compatibility-e2e-tests.yml b/.github/workflows/data-prepper-trace-analytics-raw-span-compatibility-e2e-tests.yml index 4a2f8bd345..e85d4b17ef 100644 --- a/.github/workflows/data-prepper-trace-analytics-raw-span-compatibility-e2e-tests.yml +++ b/.github/workflows/data-prepper-trace-analytics-raw-span-compatibility-e2e-tests.yml @@ -13,7 +13,7 @@ jobs: build: strategy: matrix: - java: [11, 17] + java: [11, 17, 21] fail-fast: false runs-on: ubuntu-latest diff --git a/.github/workflows/data-prepper-trace-analytics-raw-span-e2e-tests.yml b/.github/workflows/data-prepper-trace-analytics-raw-span-e2e-tests.yml index 45e063c660..d7d705c93e 100644 --- a/.github/workflows/data-prepper-trace-analytics-raw-span-e2e-tests.yml +++ b/.github/workflows/data-prepper-trace-analytics-raw-span-e2e-tests.yml @@ -13,7 +13,7 @@ jobs: build: strategy: matrix: - java: [11, 17] + java: [11, 17, 21] otelVersion: ['0.9.0-alpha', '0.16.0-alpha'] fail-fast: false diff --git a/.github/workflows/data-prepper-trace-analytics-raw-span-peer-forwarder-e2e-tests.yml b/.github/workflows/data-prepper-trace-analytics-raw-span-peer-forwarder-e2e-tests.yml index 0dce35be61..44291cdbf2 100644 --- a/.github/workflows/data-prepper-trace-analytics-raw-span-peer-forwarder-e2e-tests.yml +++ b/.github/workflows/data-prepper-trace-analytics-raw-span-peer-forwarder-e2e-tests.yml @@ -13,7 +13,7 @@ jobs: build: strategy: matrix: - java: [11, 17] + java: [11, 17, 21] fail-fast: false runs-on: ubuntu-latest diff --git a/.github/workflows/data-prepper-trace-analytics-service-map-e2e-tests.yml b/.github/workflows/data-prepper-trace-analytics-service-map-e2e-tests.yml index 14e2618546..0c5aa7eb16 100644 --- a/.github/workflows/data-prepper-trace-analytics-service-map-e2e-tests.yml +++ b/.github/workflows/data-prepper-trace-analytics-service-map-e2e-tests.yml @@ -13,7 +13,7 @@ jobs: build: strategy: matrix: - java: [11, 17] + java: [11, 17, 21] fail-fast: false runs-on: ubuntu-latest From 658c2031e3cc78b87d185df167379e00d76d3342 Mon Sep 17 00:00:00 2001 From: David Venable Date: Mon, 23 Oct 2023 14:51:22 -0700 Subject: [PATCH 04/53] Perform a full upgrade on the base Docker image when building the Data Prepper Docker image to get latest patches. (#3497) Signed-off-by: David Venable --- release/docker/Dockerfile | 1 + 1 file changed, 1 insertion(+) diff --git a/release/docker/Dockerfile b/release/docker/Dockerfile index 4701bcd37b..c4c138c0a6 100644 --- a/release/docker/Dockerfile +++ b/release/docker/Dockerfile @@ -11,6 +11,7 @@ ENV ENV_PIPELINE_FILEPATH=$PIPELINE_FILEPATH # Update all packages RUN apt -y update RUN apt -y install bash bc +RUN apt -y full-upgrade RUN mkdir -p /var/log/data-prepper ADD $ARCHIVE_FILE /usr/share From e79add68caf12e6cb736cbb56144349aaed15d0b Mon Sep 17 00:00:00 2001 From: David Venable Date: Mon, 23 Oct 2023 14:51:48 -0700 Subject: [PATCH 05/53] Test against newer versions of OpenSearch and use odd versions in the 2.x series to avoid testing against too many different versions. Updated to the latest 1.3 version. (#3512) Signed-off-by: David Venable --- .../workflows/opensearch-sink-opensearch-integration-tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/opensearch-sink-opensearch-integration-tests.yml b/.github/workflows/opensearch-sink-opensearch-integration-tests.yml index 50a8f00a3b..1d3106c136 100644 --- a/.github/workflows/opensearch-sink-opensearch-integration-tests.yml +++ b/.github/workflows/opensearch-sink-opensearch-integration-tests.yml @@ -20,7 +20,7 @@ jobs: strategy: matrix: java: [11] - opensearch: [1.0.1, 1.1.0, 1.2.4, 1.3.9, 2.0.1, 2.1.0, 2.2.1, 2.3.0, 2.4.0, 2.5.0, 2.6.0] + opensearch: [1.0.1, 1.1.0, 1.2.4, 1.3.13, 2.0.1, 2.1.0, 2.3.0, 2.5.0, 2.7.0, 2.9.0, 2.11.0] fail-fast: false runs-on: ubuntu-latest From 2f8372bd0536941d89c62e2dbc8e46de86df8469 Mon Sep 17 00:00:00 2001 From: David Venable Date: Mon, 23 Oct 2023 14:52:26 -0700 Subject: [PATCH 06/53] Removed the deprecated annotation from Record and RecordMetadata as these are currently still very necessary. Resolves #3536. (#3540) Signed-off-by: David Venable --- .../java/org/opensearch/dataprepper/model/record/Record.java | 1 - .../org/opensearch/dataprepper/model/record/RecordMetadata.java | 1 - 2 files changed, 2 deletions(-) diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/record/Record.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/record/Record.java index a5bb6e8704..65ae2d6ce5 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/record/Record.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/record/Record.java @@ -12,7 +12,6 @@ * TODO: The current implementation focuses on proving the bare bones for which this class only need to * TODO: support sample test cases. */ -@Deprecated public class Record { private final T data; private final RecordMetadata metadata; diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/record/RecordMetadata.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/record/RecordMetadata.java index 823abd30e1..8efbe0960e 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/record/RecordMetadata.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/record/RecordMetadata.java @@ -13,7 +13,6 @@ * The RecordMetadata class provides a wrapper around the ImmutableMap making metadata management easier for the * user to access. */ -@Deprecated public class RecordMetadata { private static final RecordMetadata DEFAULT_METADATA = new RecordMetadata(); From fe750f5580c5a9b60ee0253c54b5fe9425938015 Mon Sep 17 00:00:00 2001 From: Aiden Dai <68811299+daixba@users.noreply.github.com> Date: Tue, 24 Oct 2023 07:08:21 +0800 Subject: [PATCH 07/53] Add bug fixes and improvements to DDB source (#3534) Signed-off-by: Aiden Dai --- .../dynamodb-source/build.gradle | 1 + .../source/dynamodb/DynamoDBService.java | 23 ++--- .../source/dynamodb/DynamoDBSource.java | 3 +- .../converter/ExportRecordConverter.java | 32 ++++--- .../converter/MetadataKeyAttributes.java | 12 +-- .../dynamodb/converter/RecordConverter.java | 86 ++++++++++--------- .../converter/StreamRecordConverter.java | 74 ++++++++++++---- .../export/DataFileLoaderFactory.java | 18 ++-- .../dynamodb/export/DataFileScheduler.java | 4 +- .../dynamodb/stream/ShardConsumerFactory.java | 14 ++- .../source/dynamodb/stream/ShardManager.java | 20 +++-- .../converter/ExportRecordConverterTest.java | 52 ++++++++--- .../converter/StreamRecordConverterTest.java | 60 ++++++++++--- .../export/DataFileLoaderFactoryTest.java | 1 - 14 files changed, 264 insertions(+), 136 deletions(-) diff --git a/data-prepper-plugins/dynamodb-source/build.gradle b/data-prepper-plugins/dynamodb-source/build.gradle index 71b8dc2afc..d3a57de0e6 100644 --- a/data-prepper-plugins/dynamodb-source/build.gradle +++ b/data-prepper-plugins/dynamodb-source/build.gradle @@ -26,6 +26,7 @@ dependencies { implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-ion' implementation project(path: ':data-prepper-plugins:aws-plugin-api') + implementation project(path: ':data-prepper-plugins:buffer-common') testImplementation platform('org.junit:junit-bom:5.9.1') diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java index 6b5430d997..f09b5b8ded 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java @@ -41,6 +41,7 @@ import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsClient; import software.amazon.awssdk.services.s3.S3Client; +import java.time.Duration; import java.time.Instant; import java.util.List; import java.util.Map; @@ -69,6 +70,9 @@ public class DynamoDBService { private final PluginMetrics pluginMetrics; + static final Duration BUFFER_TIMEOUT = Duration.ofSeconds(60); + static final int DEFAULT_BUFFER_BATCH_SIZE = 1_000; + public DynamoDBService(EnhancedSourceCoordinator coordinator, ClientFactory clientFactory, DynamoDBSourceConfig sourceConfig, PluginMetrics pluginMetrics) { this.coordinator = coordinator; @@ -156,14 +160,13 @@ public void init() { Instant startTime = Instant.now(); if (tableInfo.getMetadata().isExportRequired()) { -// exportTime = Instant.now(); createExportPartition(tableInfo.getTableArn(), startTime, tableInfo.getMetadata().getExportBucket(), tableInfo.getMetadata().getExportPrefix()); } if (tableInfo.getMetadata().isStreamRequired()) { List shardIds; // start position by default is TRIM_HORIZON if not provided. - if (tableInfo.getMetadata().isExportRequired() || String.valueOf(StreamStartPosition.LATEST).equals(tableInfo.getMetadata().getStreamStartPosition())) { + if (tableInfo.getMetadata().isExportRequired() || tableInfo.getMetadata().getStreamStartPosition() == StreamStartPosition.LATEST) { // For a continued data extraction process that involves both export and stream // The export must be completed and loaded before stream can start. // Moreover, there should not be any gaps between the export time and the time start reading the stream @@ -274,15 +277,13 @@ private TableInfo getTableInfo(TableConfig tableConfig) { throw new InvalidPluginConfigurationException(errorMessage); } // Validate view type of DynamoDB stream - if (describeTableResult.table().streamSpecification() != null) { - String viewType = describeTableResult.table().streamSpecification().streamViewTypeAsString(); - LOG.debug("The stream view type for table " + tableName + " is " + viewType); - List supportedType = List.of("NEW_IMAGE", "NEW_AND_OLD_IMAGES"); - if (!supportedType.contains(viewType)) { - String errorMessage = "Stream " + tableConfig.getTableArn() + " is enabled with " + viewType + ". Supported types are " + supportedType; - LOG.error(errorMessage); - throw new InvalidPluginConfigurationException(errorMessage); - } + String viewType = describeTableResult.table().streamSpecification().streamViewTypeAsString(); + LOG.debug("The stream view type for table " + tableName + " is " + viewType); + List supportedType = List.of("NEW_IMAGE", "NEW_AND_OLD_IMAGES"); + if (!supportedType.contains(viewType)) { + String errorMessage = "Stream " + tableConfig.getTableArn() + " is enabled with " + viewType + ". Supported types are " + supportedType; + LOG.error(errorMessage); + throw new InvalidPluginConfigurationException(errorMessage); } streamStartPosition = tableConfig.getStreamConfig().getStartPosition(); } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSource.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSource.java index ef35c82825..ffb1809ac6 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSource.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSource.java @@ -10,7 +10,6 @@ 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.configuration.PluginSetting; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.plugin.PluginFactory; import org.opensearch.dataprepper.model.record.Record; @@ -46,7 +45,7 @@ public class DynamoDBSource implements Source>, UsesEnhancedSource @DataPrepperPluginConstructor - public DynamoDBSource(PluginMetrics pluginMetrics, final DynamoDBSourceConfig sourceConfig, final PluginFactory pluginFactory, final PluginSetting pluginSetting, final AwsCredentialsSupplier awsCredentialsSupplier) { + public DynamoDBSource(PluginMetrics pluginMetrics, final DynamoDBSourceConfig sourceConfig, final PluginFactory pluginFactory, final AwsCredentialsSupplier awsCredentialsSupplier) { LOG.info("Create DynamoDB Source"); this.pluginMetrics = pluginMetrics; this.sourceConfig = sourceConfig; diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverter.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverter.java index 70a6cbcf31..9f770dff69 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverter.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverter.java @@ -8,8 +8,8 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.dataformat.ion.IonObjectMapper; import io.micrometer.core.instrument.Counter; +import org.opensearch.dataprepper.buffer.common.BufferAccumulator; import org.opensearch.dataprepper.metrics.PluginMetrics; -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.plugins.source.dynamodb.model.TableInfo; @@ -18,7 +18,6 @@ import java.util.List; import java.util.Map; -import java.util.stream.Collectors; public class ExportRecordConverter extends RecordConverter { @@ -37,8 +36,8 @@ public class ExportRecordConverter extends RecordConverter { private final Counter exportRecordSuccessCounter; private final Counter exportRecordErrorCounter; - public ExportRecordConverter(Buffer> buffer, TableInfo tableInfo, PluginMetrics pluginMetrics) { - super(buffer, tableInfo); + public ExportRecordConverter(final BufferAccumulator> bufferAccumulator, TableInfo tableInfo, PluginMetrics pluginMetrics) { + super(bufferAccumulator, tableInfo); this.pluginMetrics = pluginMetrics; this.exportRecordSuccessCounter = pluginMetrics.counter(EXPORT_RECORDS_PROCESSED_COUNT); this.exportRecordErrorCounter = pluginMetrics.counter(EXPORT_RECORDS_PROCESSING_ERROR_COUNT); @@ -60,21 +59,26 @@ String getEventType() { } public void writeToBuffer(List lines) { - List> data = lines.stream() - .map(this::convertToMap) - .map(d -> (Map) d.get(ITEM_KEY)) - .collect(Collectors.toList()); - List> events = data.stream().map(this::convertToEvent).collect(Collectors.toList()); + int eventCount = 0; + for (String line : lines) { + Map data = (Map) convertToMap(line).get(ITEM_KEY); + try { + addToBuffer(data); + eventCount++; + } catch (Exception e) { + // will this cause too many logs? + LOG.error("Failed to add event to buffer due to {}", e.getMessage()); + } + } try { - writeEventsToBuffer(events); - exportRecordSuccessCounter.increment(events.size()); + flushBuffer(); + exportRecordSuccessCounter.increment(eventCount); } catch (Exception e) { - LOG.error("Failed to write {} events to buffer due to {}", events.size(), e.getMessage()); - exportRecordErrorCounter.increment(events.size()); + LOG.error("Failed to write {} events to buffer due to {}", eventCount, e.getMessage()); + exportRecordErrorCounter.increment(eventCount); } } - } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java index 0286627ba6..204ec82e02 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java @@ -6,15 +6,15 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.converter; public class MetadataKeyAttributes { - static final String PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE = "primary_key"; + static final String PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE = "primary_key"; - static final String PARTITION_KEY_METADATA_ATTRIBUTE = "partition_key"; + static final String PARTITION_KEY_METADATA_ATTRIBUTE = "partition_key"; - static final String SORT_KEY_METADATA_ATTRIBUTE = "sort_key"; + static final String SORT_KEY_METADATA_ATTRIBUTE = "sort_key"; - static final String EVENT_TIMESTAMP_METADATA_ATTRIBUTE = "ts"; + static final String EVENT_TIMESTAMP_METADATA_ATTRIBUTE = "ts"; - static final String STREAM_EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE = "op"; + static final String EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE = "op"; - static final String EVENT_TABLE_NAME_METADATA_ATTRIBUTE = "table_name"; + static final String EVENT_TABLE_NAME_METADATA_ATTRIBUTE = "table_name"; } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java index cedf7fb0f1..d5c5a059a8 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java @@ -5,7 +5,7 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.converter; -import org.opensearch.dataprepper.model.buffer.Buffer; +import org.opensearch.dataprepper.buffer.common.BufferAccumulator; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.event.EventMetadata; import org.opensearch.dataprepper.model.event.JacksonEvent; @@ -13,16 +13,14 @@ import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo; -import java.time.Instant; -import java.util.List; import java.util.Map; -import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_TABLE_NAME_METADATA_ATTRIBUTE; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_TIMESTAMP_METADATA_ATTRIBUTE; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.PARTITION_KEY_METADATA_ATTRIBUTE; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.SORT_KEY_METADATA_ATTRIBUTE; -import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.STREAM_EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE; /** * Base Record Processor definition. @@ -33,68 +31,72 @@ public abstract class RecordConverter { private static final String DEFAULT_ACTION = OpenSearchBulkActions.INDEX.toString(); - private static final int DEFAULT_WRITE_TIMEOUT_MILLIS = 60_000; - private final Buffer> buffer; + private final BufferAccumulator> bufferAccumulator; private final TableInfo tableInfo; - public RecordConverter(Buffer> buffer, TableInfo tableInfo) { - this.buffer = buffer; + public RecordConverter(final BufferAccumulator> bufferAccumulator, TableInfo tableInfo) { + this.bufferAccumulator = bufferAccumulator; this.tableInfo = tableInfo; } - abstract String getEventType(); /** - * Default method to conduct the document ID value, - * Using partition key plus sort key (if any) + * Extract the value based on attribute map + * + * @param data A map of attribute name and value + * @param attributeName Attribute name + * @return the related attribute value, return null if the attribute name doesn't exist. */ - String getId(Map data) { - String partitionKey = String.valueOf(data.get(tableInfo.getMetadata().getPartitionKeyAttributeName())); - if (tableInfo.getMetadata().getSortKeyAttributeName() == null) { - return partitionKey; + private String getAttributeValue(final Map data, String attributeName) { + if (data.containsKey(attributeName)) { + return String.valueOf(data.get(attributeName)); } - String sortKey = String.valueOf(data.get(tableInfo.getMetadata().getSortKeyAttributeName())); - return partitionKey + "_" + sortKey; + return null; } - String getPartitionKey(final Map data) { - return String.valueOf(data.get(tableInfo.getMetadata().getPartitionKeyAttributeName())); + void flushBuffer() throws Exception { + bufferAccumulator.flush(); } - String getSortKey(final Map data) { - return String.valueOf(data.get(tableInfo.getMetadata().getSortKeyAttributeName())); - } - - void writeEventsToBuffer(List> events) throws Exception { - buffer.writeAll(events, DEFAULT_WRITE_TIMEOUT_MILLIS); - } - - public Record convertToEvent(Map data, Instant eventCreationTime, String streamEventName) { - Event event; - event = JacksonEvent.builder() + /** + * Add event record to buffer + * + * @param data A map to hold event data, note that it may be empty. + * @param keys A map to hold the keys (partition key and sort key) + * @param eventCreationTimeMillis Creation timestamp of the event + * @param eventName Event name + * @throws Exception Exception if failed to write to buffer. + */ + public void addToBuffer(Map data, Map keys, long eventCreationTimeMillis, String eventName) throws Exception { + Event event = JacksonEvent.builder() .withEventType(getEventType()) .withData(data) .build(); EventMetadata eventMetadata = event.getMetadata(); eventMetadata.setAttribute(EVENT_TABLE_NAME_METADATA_ATTRIBUTE, tableInfo.getTableName()); - if (eventCreationTime != null) { - eventMetadata.setAttribute(EVENT_TIMESTAMP_METADATA_ATTRIBUTE, eventCreationTime.toEpochMilli()); + eventMetadata.setAttribute(EVENT_TIMESTAMP_METADATA_ATTRIBUTE, eventCreationTimeMillis); + eventMetadata.setAttribute(EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE, mapStreamEventNameToBulkAction(eventName)); + String partitionKey = getAttributeValue(keys, tableInfo.getMetadata().getPartitionKeyAttributeName()); + eventMetadata.setAttribute(PARTITION_KEY_METADATA_ATTRIBUTE, partitionKey); + + String sortKey = getAttributeValue(keys, tableInfo.getMetadata().getSortKeyAttributeName()); + if (sortKey != null) { + eventMetadata.setAttribute(SORT_KEY_METADATA_ATTRIBUTE, sortKey); + eventMetadata.setAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE, partitionKey + "_" + sortKey); + } else { + eventMetadata.setAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE, partitionKey); } - - eventMetadata.setAttribute(STREAM_EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE, mapStreamEventNameToBulkAction(streamEventName)); - eventMetadata.setAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE, getId(data)); - eventMetadata.setAttribute(PARTITION_KEY_METADATA_ATTRIBUTE, getPartitionKey(data)); - eventMetadata.setAttribute(SORT_KEY_METADATA_ATTRIBUTE, getSortKey(data)); - - return new Record<>(event); + bufferAccumulator.add(new Record<>(event)); } - public Record convertToEvent(Map data) { - return convertToEvent(data, null, null); + public void addToBuffer(Map data) throws Exception { + // Export data doesn't have an event timestamp + // Default to current timestamp when the event is added to buffer + addToBuffer(data, data, System.currentTimeMillis(), null); } private String mapStreamEventNameToBulkAction(final String streamEventName) { diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java index bdfe6dcbe1..4b0bba0a4a 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java @@ -9,18 +9,19 @@ import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import io.micrometer.core.instrument.Counter; +import org.opensearch.dataprepper.buffer.common.BufferAccumulator; import org.opensearch.dataprepper.metrics.PluginMetrics; -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.plugins.source.dynamodb.model.TableInfo; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.awssdk.enhanced.dynamodb.document.EnhancedDocument; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.Record; +import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.stream.Collectors; public class StreamRecordConverter extends RecordConverter { private static final Logger LOG = LoggerFactory.getLogger(StreamRecordConverter.class); @@ -39,8 +40,8 @@ public class StreamRecordConverter extends RecordConverter { private final Counter changeEventSuccessCounter; private final Counter changeEventErrorCounter; - public StreamRecordConverter(Buffer> buffer, TableInfo tableInfo, PluginMetrics pluginMetrics) { - super(buffer, tableInfo); + public StreamRecordConverter(final BufferAccumulator> bufferAccumulator, TableInfo tableInfo, PluginMetrics pluginMetrics) { + super(bufferAccumulator, tableInfo); this.pluginMetrics = pluginMetrics; this.changeEventSuccessCounter = pluginMetrics.counter(CHANGE_EVENTS_PROCESSED_COUNT); this.changeEventErrorCounter = pluginMetrics.counter(CHANGE_EVENTS_PROCESSING_ERROR_COUNT); @@ -51,30 +52,67 @@ String getEventType() { return "STREAM"; } - public void writeToBuffer(List records) { - // TODO: What if convert failed. - List> events = records.stream() - .map(record -> convertToEvent( - toMap(EnhancedDocument.fromAttributeValueMap(record.dynamodb().newImage()).toJson()), - record.dynamodb().approximateCreationDateTime(), - record.eventNameAsString())) - .collect(Collectors.toList()); + + public void writeToBuffer(List records) { + + int eventCount = 0; + for (Record record : records) { + // NewImage may be empty + Map data = convertData(record.dynamodb().newImage()); + // Always get keys from dynamodb().keys() + Map keys = convertKeys(record.dynamodb().keys()); + + try { + addToBuffer(data, keys, record.dynamodb().approximateCreationDateTime().toEpochMilli(), record.eventNameAsString()); + eventCount++; + } catch (Exception e) { + // will this cause too many logs? + LOG.error("Failed to add event to buffer due to {}", e.getMessage()); + changeEventErrorCounter.increment(); + } + } try { - writeEventsToBuffer(events); - changeEventSuccessCounter.increment(events.size()); + flushBuffer(); + changeEventSuccessCounter.increment(eventCount); } catch (Exception e) { - LOG.error("Failed to write {} events to buffer due to {}", events.size(), e.getMessage()); - changeEventErrorCounter.increment(events.size()); + LOG.error("Failed to write {} events to buffer due to {}", eventCount, e.getMessage()); + changeEventErrorCounter.increment(eventCount); } } - private Map toMap(String jsonData) { + /** + * Convert the DynamoDB attribute map to a normal map for data + */ + private Map convertData(Map data) { try { + String jsonData = EnhancedDocument.fromAttributeValueMap(data).toJson(); return MAPPER.readValue(jsonData, MAP_TYPE_REFERENCE); } catch (JsonProcessingException e) { return null; } } + + /** + * Convert the DynamoDB attribute map to a normal map for keys + * This method may not be necessary, can use convertData() alternatively + */ + private Map convertKeys(Map keys) { + Map result = new HashMap<>(); + // The attribute type for key can only be N, B or S + keys.forEach(((attributeName, attributeValue) -> { + if (attributeValue.type() == AttributeValue.Type.N) { + // N for number + result.put(attributeName, attributeValue.n()); + } else if (attributeValue.type() == AttributeValue.Type.B) { + // B for Binary + result.put(attributeName, attributeValue.b().toString()); + } else { + result.put(attributeName, attributeValue.s()); + } + })); + return result; + + } } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactory.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactory.java index 75dccf8218..9e2378e3df 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactory.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactory.java @@ -5,6 +5,7 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.export; +import org.opensearch.dataprepper.buffer.common.BufferAccumulator; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.buffer.Buffer; import org.opensearch.dataprepper.model.event.Event; @@ -15,34 +16,37 @@ import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo; import software.amazon.awssdk.services.s3.S3Client; +import java.time.Duration; + /** * Factory class for DataFileLoader thread. */ public class DataFileLoaderFactory { + static final Duration BUFFER_TIMEOUT = Duration.ofSeconds(60); + static final int DEFAULT_BUFFER_BATCH_SIZE = 1_000; private final EnhancedSourceCoordinator coordinator; - private final S3ObjectReader fileReader; - + private final S3ObjectReader objectReader; private final PluginMetrics pluginMetrics; - private final Buffer> buffer; - public DataFileLoaderFactory(EnhancedSourceCoordinator coordinator, S3Client s3Client, PluginMetrics pluginMetrics, Buffer> buffer) { + public DataFileLoaderFactory(EnhancedSourceCoordinator coordinator, S3Client s3Client, PluginMetrics pluginMetrics, final Buffer> buffer) { this.coordinator = coordinator; this.pluginMetrics = pluginMetrics; this.buffer = buffer; - fileReader = new S3ObjectReader(s3Client); + objectReader = new S3ObjectReader(s3Client); } public Runnable createDataFileLoader(DataFilePartition dataFilePartition, TableInfo tableInfo) { - ExportRecordConverter recordProcessor = new ExportRecordConverter(buffer, tableInfo, pluginMetrics); + final BufferAccumulator> bufferAccumulator = BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT); + ExportRecordConverter recordProcessor = new ExportRecordConverter(bufferAccumulator, tableInfo, pluginMetrics); DataFileCheckpointer checkpointer = new DataFileCheckpointer(coordinator, dataFilePartition); // Start a data loader thread. DataFileLoader loader = DataFileLoader.builder() - .s3ObjectReader(fileReader) + .s3ObjectReader(objectReader) .bucketName(dataFilePartition.getBucket()) .key(dataFilePartition.getKey()) .recordConverter(recordProcessor) diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java index 6c56323cca..8c1979e320 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java @@ -35,7 +35,7 @@ public class DataFileScheduler implements Runnable { /** * Maximum concurrent data loader per node */ - private static final int MAX_JOB_COUNT = 4; + private static final int MAX_JOB_COUNT = 3; /** * Default interval to acquire a lease from coordination store @@ -85,7 +85,7 @@ private void processDataFilePartition(DataFilePartition dataFilePartition) { @Override public void run() { - LOG.info("Start running Data File Scheduler"); + LOG.debug("Start running Data File Scheduler"); while (!Thread.currentThread().isInterrupted()) { if (numOfWorkers.get() < MAX_JOB_COUNT) { diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java index 1c35f882ad..ba311456b1 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java @@ -5,6 +5,7 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.stream; +import org.opensearch.dataprepper.buffer.common.BufferAccumulator; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.buffer.Buffer; import org.opensearch.dataprepper.model.event.Event; @@ -20,6 +21,7 @@ import org.slf4j.LoggerFactory; import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsClient; +import java.time.Duration; import java.time.Instant; import java.util.Optional; @@ -31,17 +33,21 @@ public class ShardConsumerFactory { private static final int STREAM_TO_TABLE_OFFSET = "stream/".length(); + static final Duration BUFFER_TIMEOUT = Duration.ofSeconds(60); + static final int DEFAULT_BUFFER_BATCH_SIZE = 1_000; + private final DynamoDbStreamsClient streamsClient; private final EnhancedSourceCoordinator enhancedSourceCoordinator; private final PluginMetrics pluginMetrics; private final ShardManager shardManager; - private final Buffer> buffer; + public ShardConsumerFactory(final EnhancedSourceCoordinator enhancedSourceCoordinator, final DynamoDbStreamsClient streamsClient, final PluginMetrics pluginMetrics, - final ShardManager shardManager, final Buffer> buffer) { + final ShardManager shardManager, + final Buffer> buffer) { this.streamsClient = streamsClient; this.enhancedSourceCoordinator = enhancedSourceCoordinator; this.pluginMetrics = pluginMetrics; @@ -77,7 +83,9 @@ public Runnable createConsumer(StreamPartition streamPartition) { StreamCheckpointer checkpointer = new StreamCheckpointer(enhancedSourceCoordinator, streamPartition); String tableArn = getTableArn(streamPartition.getStreamArn()); TableInfo tableInfo = getTableInfo(tableArn); - StreamRecordConverter recordConverter = new StreamRecordConverter(buffer, tableInfo, pluginMetrics); + final BufferAccumulator> bufferAccumulator = BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT); + + StreamRecordConverter recordConverter = new StreamRecordConverter(bufferAccumulator, tableInfo, pluginMetrics); ShardConsumer shardConsumer = ShardConsumer.builder(streamsClient) .recordConverter(recordConverter) diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardManager.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardManager.java index 1e342178de..32ae8b573b 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardManager.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardManager.java @@ -38,12 +38,18 @@ public ShardManager(DynamoDbStreamsClient streamsClient) { this.streamsClient = streamsClient; } - private List listShards(String streamArn) { + /** + * List all shards using describe stream API. + * + * @param streamArn Stream Arn + * @param lastEvaluatedShardId Start shard id for listing, useful when trying to get child shards. If not provided, all shards will be returned. + * @return A list of {@link Shard} + */ + private List listShards(String streamArn, String lastEvaluatedShardId) { LOG.info("Start getting all shards from {}", streamArn); long startTime = System.currentTimeMillis(); // Get all the shard IDs from the stream. List shards = new ArrayList<>(); - String lastEvaluatedShardId = null; do { DescribeStreamRequest req = DescribeStreamRequest.builder() .streamArn(streamArn) @@ -70,12 +76,12 @@ private List listShards(String streamArn) { * Get a list of Child Shard Ids based on a parent shard id provided. * * @param streamArn Stream Arn - * @param shardId Parent Shard Id - * @return A list of child shard Ids. + * @param shardId Parent Shard id + * @return A list of child shard ids. */ public List getChildShardIds(String streamArn, String shardId) { LOG.debug("Getting child ids for " + shardId); - List shards = listShards(streamArn); + List shards = listShards(streamArn, shardId); return shards.stream() .filter(s -> shardId.equals(s.parentShardId())) .map(s -> s.shardId()) @@ -90,7 +96,7 @@ public List getChildShardIds(String streamArn, String shardId) { * @return A list of shard Ids */ public List getActiveShards(String streamArn) { - List shards = listShards(streamArn); + List shards = listShards(streamArn, null); return shards.stream() .filter(s -> s.sequenceNumberRange().endingSequenceNumber() == null) .map(s -> s.shardId()) @@ -155,7 +161,7 @@ public String getShardIterator(String streamArn, String shardId, String sequence * @return A list of root shard Ids */ public List getRootShardIds(String streamArn) { - List shards = listShards(streamArn); + List shards = listShards(streamArn, null); List childIds = shards.stream().map(shard -> shard.shardId()).collect(Collectors.toList()); List rootIds = shards.stream() diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java index 838bb9f0ab..e79d231cb3 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java @@ -12,29 +12,37 @@ import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.buffer.common.BufferAccumulator; import org.opensearch.dataprepper.metrics.PluginMetrics; -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.plugins.source.dynamodb.model.TableInfo; import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableMetadata; import java.util.ArrayList; -import java.util.Collection; import java.util.List; import java.util.Random; 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.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyDouble; -import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.BDDMockito.given; import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; -import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.ExportRecordConverter.EXPORT_RECORDS_PROCESSING_ERROR_COUNT; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.ExportRecordConverter.EXPORT_RECORDS_PROCESSED_COUNT; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.ExportRecordConverter.EXPORT_RECORDS_PROCESSING_ERROR_COUNT; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_TIMESTAMP_METADATA_ATTRIBUTE; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.PARTITION_KEY_METADATA_ATTRIBUTE; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.SORT_KEY_METADATA_ATTRIBUTE; @ExtendWith(MockitoExtension.class) class ExportRecordConverterTest { @@ -43,7 +51,7 @@ class ExportRecordConverterTest { private PluginMetrics pluginMetrics; @Mock - private Buffer> buffer; + private BufferAccumulator> bufferAccumulator; private TableInfo tableInfo; @@ -97,16 +105,40 @@ void test_writeToBuffer() throws Exception { int numberOfRecords = random.nextInt(10); List data = generateData(numberOfRecords); - ExportRecordConverter recordConverter = new ExportRecordConverter(buffer, tableInfo, pluginMetrics); + ExportRecordConverter recordConverter = new ExportRecordConverter(bufferAccumulator, tableInfo, pluginMetrics); - final ArgumentCaptor>> writeRequestArgumentCaptor = ArgumentCaptor.forClass(Collection.class); - doNothing().when(buffer).writeAll(writeRequestArgumentCaptor.capture(), anyInt()); recordConverter.writeToBuffer(data); - - assertThat(writeRequestArgumentCaptor.getValue().size(), equalTo(numberOfRecords)); + verify(bufferAccumulator, times(numberOfRecords)).add(any(Record.class)); verify(exportRecordSuccess).increment(anyDouble()); verifyNoInteractions(exportRecordErrors); } + + @Test + void test_writeSingleRecordToBuffer() throws Exception { + final String pk = UUID.randomUUID().toString(); + final String sk = UUID.randomUUID().toString(); + String line = " $ion_1_0 {Item:{PK:\"" + pk + "\",SK:\"" + sk + "\"}}"; + + final ArgumentCaptor recordArgumentCaptor = ArgumentCaptor.forClass(Record.class); + + ExportRecordConverter recordConverter = new ExportRecordConverter(bufferAccumulator, tableInfo, pluginMetrics); + doNothing().when(bufferAccumulator).add(recordArgumentCaptor.capture()); +// doNothing().when(bufferAccumulator).flush(); + + recordConverter.writeToBuffer(List.of(line)); + verify(bufferAccumulator).add(any(Record.class)); + verify(bufferAccumulator).flush(); + assertThat(recordArgumentCaptor.getValue().getData(), notNullValue()); + JacksonEvent event = (JacksonEvent) recordArgumentCaptor.getValue().getData(); + + assertThat(event.getMetadata(), notNullValue()); + + assertThat(event.getMetadata().getAttribute(PARTITION_KEY_METADATA_ATTRIBUTE), equalTo(pk)); + assertThat(event.getMetadata().getAttribute(SORT_KEY_METADATA_ATTRIBUTE), equalTo(sk)); + assertThat(event.getMetadata().getAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE), equalTo(pk + "_" + sk)); + assertThat(event.getMetadata().getAttribute(EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE), equalTo(OpenSearchBulkActions.INDEX.toString())); + assertThat(event.getMetadata().getAttribute(EVENT_TIMESTAMP_METADATA_ATTRIBUTE), notNullValue()); + } } \ No newline at end of file diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java index 1b9b161ecc..23ac32f6df 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java @@ -12,9 +12,11 @@ import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.buffer.common.BufferAccumulator; import org.opensearch.dataprepper.metrics.PluginMetrics; -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.plugins.source.dynamodb.model.TableInfo; import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableMetadata; @@ -24,7 +26,6 @@ import java.time.Instant; import java.util.ArrayList; -import java.util.Collection; import java.util.List; import java.util.Map; import java.util.Random; @@ -32,14 +33,21 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.notNullValue; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyDouble; -import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.BDDMockito.given; import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; -import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.StreamRecordConverter.CHANGE_EVENTS_PROCESSING_ERROR_COUNT; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_TIMESTAMP_METADATA_ATTRIBUTE; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.PARTITION_KEY_METADATA_ATTRIBUTE; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.SORT_KEY_METADATA_ATTRIBUTE; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.StreamRecordConverter.CHANGE_EVENTS_PROCESSED_COUNT; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.StreamRecordConverter.CHANGE_EVENTS_PROCESSING_ERROR_COUNT; @ExtendWith(MockitoExtension.class) class StreamRecordConverterTest { @@ -47,7 +55,7 @@ class StreamRecordConverterTest { private PluginMetrics pluginMetrics; @Mock - private Buffer> buffer; + private BufferAccumulator> bufferAccumulator; private TableInfo tableInfo; @@ -90,21 +98,46 @@ void test_writeToBuffer() throws Exception { int numberOfRecords = random.nextInt(10); - List data = buildRecords(numberOfRecords); - - StreamRecordConverter recordConverter = new StreamRecordConverter(buffer, tableInfo, pluginMetrics); + List records = buildRecords(numberOfRecords); - final ArgumentCaptor>> writeRequestArgumentCaptor = ArgumentCaptor.forClass(Collection.class); - doNothing().when(buffer).writeAll(writeRequestArgumentCaptor.capture(), anyInt()); + StreamRecordConverter recordConverter = new StreamRecordConverter(bufferAccumulator, tableInfo, pluginMetrics); - - recordConverter.writeToBuffer(data); - assertThat(writeRequestArgumentCaptor.getValue().size(), equalTo(numberOfRecords)); + recordConverter.writeToBuffer(records); + verify(bufferAccumulator, times(numberOfRecords)).add(any(Record.class)); + verify(bufferAccumulator).flush(); verify(changeEventSuccessCounter).increment(anyDouble()); verifyNoInteractions(changeEventErrorCounter); + } + + @Test + void test_writeSingleRecordToBuffer() throws Exception { + + List records = buildRecords(1); + final ArgumentCaptor recordArgumentCaptor = ArgumentCaptor.forClass(Record.class); + software.amazon.awssdk.services.dynamodb.model.Record record = records.get(0); + StreamRecordConverter recordConverter = new StreamRecordConverter(bufferAccumulator, tableInfo, pluginMetrics); + doNothing().when(bufferAccumulator).add(recordArgumentCaptor.capture()); + + recordConverter.writeToBuffer(records); + verify(bufferAccumulator).add(any(Record.class)); + verify(bufferAccumulator).flush(); + verify(changeEventSuccessCounter).increment(anyDouble()); + assertThat(recordArgumentCaptor.getValue().getData(), notNullValue()); + JacksonEvent event = (JacksonEvent) recordArgumentCaptor.getValue().getData(); + + assertThat(event.getMetadata(), notNullValue()); + String partitionKey = record.dynamodb().keys().get(partitionKeyAttrName).s(); + String sortKey = record.dynamodb().keys().get(sortKeyAttrName).s(); + assertThat(event.getMetadata().getAttribute(PARTITION_KEY_METADATA_ATTRIBUTE), equalTo(partitionKey)); + assertThat(event.getMetadata().getAttribute(SORT_KEY_METADATA_ATTRIBUTE), equalTo(sortKey)); + assertThat(event.getMetadata().getAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE), equalTo(partitionKey + "_" + sortKey)); + assertThat(event.getMetadata().getAttribute(EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE), equalTo(OpenSearchBulkActions.CREATE.toString())); + assertThat(event.getMetadata().getAttribute(EVENT_TIMESTAMP_METADATA_ATTRIBUTE), equalTo(record.dynamodb().approximateCreationDateTime().toEpochMilli())); + + verifyNoInteractions(changeEventErrorCounter); } private List buildRecords(int count) { @@ -116,6 +149,7 @@ private List buildRecords StreamRecord streamRecord = StreamRecord.builder() .newImage(data) + .keys(data) .sequenceNumber(UUID.randomUUID().toString()) .approximateCreationDateTime(Instant.now()) .build(); diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactoryTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactoryTest.java index d1f7426ad7..061cb91a93 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactoryTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactoryTest.java @@ -45,7 +45,6 @@ class DataFileLoaderFactoryTest { @Mock private Buffer> buffer; - private TableInfo tableInfo; private final String tableName = UUID.randomUUID().toString(); From a9e419af24810cad29a950205376bc5fc2b89391 Mon Sep 17 00:00:00 2001 From: mallikagogoi7 Date: Wed, 25 Oct 2023 05:40:29 +0530 Subject: [PATCH 08/53] HttpSink plugin functionality for #874. (#3065) HttpSink plugin functionality for #874. Signed-off-by: mallikagogoi7 --- .../plugins/accumulator/Buffer.java | 6 +- .../plugins/accumulator/InMemoryBuffer.java | 18 +- .../plugins/accumulator/LocalFileBuffer.java | 22 +-- .../accumulator/InMemoryBufferTest.java | 17 -- .../accumulator/LocalFileBufferTest.java | 14 -- .../plugins/sink/ThresholdValidatorTest.java | 79 -------- data-prepper-plugins/http-sink/README.md | 187 ++++++++++++++++++ data-prepper-plugins/http-sink/build.gradle | 39 +++- .../plugins/sink/http/HttpSinkServiceIT.java | 154 +++++++++++++++ .../AwsRequestSigningApacheInterceptor.java | 21 +- .../http/FailedHttpResponseInterceptor.java | 13 +- .../plugins/sink/http/HTTPSink.java | 33 ++-- .../sink/http/OAuthAccessTokenManager.java | 69 +++++++ .../HttpClientSSLConnectionManager.java | 26 ++- .../configuration/AuthenticationOptions.java | 24 +++ .../AwsAuthenticationOptions.java | 10 + .../configuration/BasicAuthCredentials.java | 25 +++ .../configuration/BearerTokenOptions.java | 82 ++++++++ .../configuration/CustomHeaderOptions.java | 52 ----- .../configuration/HttpSinkConfiguration.java | 59 +++++- .../configuration/UrlConfigurationOption.java | 58 ------ .../plugins/sink/http/dlq/DlqPushHandler.java | 8 +- .../plugins/sink/http/dlq/FailedDlqData.java | 54 ++--- .../BearerTokenAuthHttpSinkHandler.java | 26 ++- .../sink/http/service/HttpSinkAwsService.java | 5 +- .../sink/http/service/HttpSinkService.java | 134 +++++++++---- .../FailedHttpResponseInterceptorTest.java | 34 ++++ .../plugins/sink/http/HttpSinkTest.java | 122 ++++++++++++ .../CertificateProviderFactoryTest.java | 76 +++++++ .../HttpClientSSLConnectionManagerTest.java | 50 +++++ .../CustomHeaderOptionsTest.java | 43 ---- .../HttpSinkConfigurationTest.java | 86 +++++++- .../configuration/ThresholdOptionsTest.java | 3 +- .../UrlConfigurationOptionTest.java | 44 ----- .../sink/http/dlq/DlqPushHandlerTest.java | 96 +++++++++ .../handler/BasicAuthHttpSinkHandlerTest.java | 52 +++++ .../BearerTokenAuthHttpSinkHandlerTest.java | 44 +++++ .../http/service/HttpSinkAwsServiceTest.java | 56 ++++++ .../http/service/HttpSinkServiceTest.java | 43 ++-- .../sink/http/service/WebhookServiceTest.java | 13 +- .../sink/http/util/HttpSinkUtilTest.java | 34 ++++ .../src/test/resources/test_cert.crt | 14 ++ .../src/test/resources/test_decrypted_key.key | 15 ++ 43 files changed, 1584 insertions(+), 476 deletions(-) create mode 100644 data-prepper-plugins/http-sink/README.md create mode 100644 data-prepper-plugins/http-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/http/HttpSinkServiceIT.java rename data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/{ => http}/AwsRequestSigningApacheInterceptor.java (91%) create mode 100644 data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/OAuthAccessTokenManager.java create mode 100644 data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/AuthenticationOptions.java create mode 100644 data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/BasicAuthCredentials.java create mode 100644 data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/BearerTokenOptions.java delete mode 100644 data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/CustomHeaderOptions.java delete mode 100644 data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/UrlConfigurationOption.java create mode 100644 data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/FailedHttpResponseInterceptorTest.java create mode 100644 data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/HttpSinkTest.java create mode 100644 data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/certificate/CertificateProviderFactoryTest.java create mode 100644 data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/certificate/HttpClientSSLConnectionManagerTest.java delete mode 100644 data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/configuration/CustomHeaderOptionsTest.java delete mode 100644 data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/configuration/UrlConfigurationOptionTest.java create mode 100644 data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/dlq/DlqPushHandlerTest.java create mode 100644 data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/handler/BasicAuthHttpSinkHandlerTest.java create mode 100644 data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/handler/BearerTokenAuthHttpSinkHandlerTest.java create mode 100644 data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkAwsServiceTest.java create mode 100644 data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/util/HttpSinkUtilTest.java create mode 100644 data-prepper-plugins/http-sink/src/test/resources/test_cert.crt create mode 100644 data-prepper-plugins/http-sink/src/test/resources/test_decrypted_key.key diff --git a/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/Buffer.java b/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/Buffer.java index c6f6018a4f..29dd0c204d 100644 --- a/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/Buffer.java +++ b/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/Buffer.java @@ -6,6 +6,7 @@ package org.opensearch.dataprepper.plugins.accumulator; import java.io.IOException; +import java.io.OutputStream; /** * A buffer can hold data before flushing it any Sink. @@ -21,5 +22,8 @@ public interface Buffer { long getDuration(); byte[] getSinkBufferData() throws IOException; - void writeEvent(byte[] bytes) throws IOException; + + OutputStream getOutputStream(); + + void setEventCount(int eventCount); } diff --git a/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBuffer.java b/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBuffer.java index e583e54260..252dac88a9 100644 --- a/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBuffer.java +++ b/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBuffer.java @@ -8,6 +8,7 @@ import org.apache.commons.lang3.time.StopWatch; import java.io.ByteArrayOutputStream; import java.io.IOException; +import java.io.OutputStream; import java.util.concurrent.TimeUnit; /** @@ -49,16 +50,13 @@ public byte[] getSinkBufferData() throws IOException { return byteArrayOutputStream.toByteArray(); } - /** - * write byte array to output stream. - * - * @param bytes byte array. - * @throws IOException while writing to output stream fails. - */ @Override - public void writeEvent(byte[] bytes) throws IOException { - byteArrayOutputStream.write(bytes); - byteArrayOutputStream.write(System.lineSeparator().getBytes()); - eventCount++; + public OutputStream getOutputStream() { + return byteArrayOutputStream; + } + + @Override + public void setEventCount(int eventCount) { + this.eventCount = eventCount; } } \ No newline at end of file diff --git a/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBuffer.java b/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBuffer.java index 9f9b4a3aac..f419fde25b 100644 --- a/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBuffer.java +++ b/data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBuffer.java @@ -70,18 +70,6 @@ public byte[] getSinkBufferData() throws IOException { return fileData; } - /** - * write byte array to output stream. - * @param bytes byte array. - * @throws IOException while writing to output stream fails. - */ - @Override - public void writeEvent(byte[] bytes) throws IOException { - outputStream.write(bytes); - outputStream.write(System.lineSeparator().getBytes()); - eventCount++; - } - /** * Flushing the buffered data into the output stream. */ @@ -106,4 +94,14 @@ protected void removeTemporaryFile() { } } } + + @Override + public OutputStream getOutputStream() { + return outputStream; + } + + @Override + public void setEventCount(int eventCount) { + this.eventCount = eventCount; + } } \ No newline at end of file diff --git a/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBufferTest.java b/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBufferTest.java index ad07cc4011..88d5205dca 100644 --- a/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBufferTest.java +++ b/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBufferTest.java @@ -10,11 +10,6 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.junit.jupiter.MockitoExtension; -import java.io.IOException; - -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.greaterThanOrEqualTo; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; @ExtendWith(MockitoExtension.class) @@ -24,18 +19,6 @@ class InMemoryBufferTest { private InMemoryBuffer inMemoryBuffer; - @Test - void test_with_write_event_into_buffer() throws IOException { - inMemoryBuffer = new InMemoryBuffer(); - - while (inMemoryBuffer.getEventCount() < MAX_EVENTS) { - inMemoryBuffer.writeEvent(generateByteArray()); - } - assertThat(inMemoryBuffer.getSize(), greaterThanOrEqualTo(54110L)); - assertThat(inMemoryBuffer.getEventCount(), equalTo(MAX_EVENTS)); - assertThat(inMemoryBuffer.getDuration(), greaterThanOrEqualTo(0L)); - - } @Test void test_getSinkData_success() { diff --git a/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBufferTest.java b/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBufferTest.java index 53c556e75c..6090e9b57c 100644 --- a/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBufferTest.java +++ b/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBufferTest.java @@ -16,7 +16,6 @@ import java.util.UUID; import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.greaterThan; import static org.hamcrest.Matchers.equalTo; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; @@ -38,19 +37,6 @@ void setUp() throws IOException { localFileBuffer = new LocalFileBuffer(tempFile); } - @Test - void test_with_write_events_into_buffer() throws IOException { - while (localFileBuffer.getEventCount() < 55) { - localFileBuffer.writeEvent(generateByteArray()); - } - assertThat(localFileBuffer.getSize(), greaterThan(1l)); - assertThat(localFileBuffer.getEventCount(), equalTo(55)); - assertThat(localFileBuffer.getDuration(), equalTo(0L)); - localFileBuffer.flushAndCloseStream(); - localFileBuffer.removeTemporaryFile(); - assertFalse(tempFile.exists(), "The temp file has not been deleted."); - } - @Test void test_without_write_events_into_buffer() { assertThat(localFileBuffer.getSize(), equalTo(0L)); diff --git a/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/sink/ThresholdValidatorTest.java b/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/sink/ThresholdValidatorTest.java index d21a664fb5..20667e1564 100644 --- a/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/sink/ThresholdValidatorTest.java +++ b/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/sink/ThresholdValidatorTest.java @@ -14,7 +14,6 @@ import java.io.IOException; import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertTrue; class ThresholdValidatorTest { @@ -23,20 +22,6 @@ class ThresholdValidatorTest { @BeforeEach void setUp() throws IOException { inMemoryBuffer = new InMemoryBufferFactory().getBuffer(); - - while (inMemoryBuffer.getEventCount() < 100) { - inMemoryBuffer.writeEvent(generateByteArray()); - } - } - - @Test - void test_exceedThreshold_true_dueTo_maxEvents_is_less_than_buffered_event_count() { - final int maxEvents = 95; - final ByteCount maxBytes = ByteCount.parse("50kb"); - final long maxCollectionDuration = 15; - boolean isThresholdExceed = ThresholdValidator.checkThresholdExceed(inMemoryBuffer, maxEvents, - maxBytes, maxCollectionDuration); - assertTrue(isThresholdExceed, "Threshold not exceeded"); } @Test @@ -49,16 +34,6 @@ void test_exceedThreshold_false_dueTo_maxEvents_is_greater_than_buffered_event_c assertFalse(isThresholdExceed, "Threshold exceeded"); } - @Test - void test_exceedThreshold_ture_dueTo_maxBytes_is_less_than_buffered_byte_count() { - final int maxEvents = 500; - final ByteCount maxBytes = ByteCount.parse("1b"); - final long maxCollectionDuration = 15; - boolean isThresholdExceed = ThresholdValidator.checkThresholdExceed(inMemoryBuffer, maxEvents, maxBytes, - maxCollectionDuration); - assertTrue(isThresholdExceed, "Threshold not exceeded"); - } - @Test void test_exceedThreshold_false_dueTo_maxBytes_is_greater_than_buffered_byte_count() { final int maxEvents = 500; @@ -69,58 +44,4 @@ void test_exceedThreshold_false_dueTo_maxBytes_is_greater_than_buffered_byte_cou assertFalse(isThresholdExceed, "Threshold exceeded"); } - @Test - void test_exceedThreshold_ture_dueTo_maxCollectionDuration_is_less_than_buffered_event_collection_duration() - throws IOException, InterruptedException { - final int maxEvents = 500; - final ByteCount maxBytes = ByteCount.parse("500mb"); - final long maxCollectionDuration = 10; - - inMemoryBuffer = new InMemoryBufferFactory().getBuffer(); - boolean isThresholdExceed = Boolean.FALSE; - synchronized (this) { - while (inMemoryBuffer.getEventCount() < 100) { - inMemoryBuffer.writeEvent(generateByteArray()); - isThresholdExceed = ThresholdValidator.checkThresholdExceed(inMemoryBuffer, maxEvents, - maxBytes, maxCollectionDuration); - if (isThresholdExceed) { - break; - } - wait(5000); - } - } - assertTrue(isThresholdExceed, "Threshold not exceeded"); - } - - @Test - void test_exceedThreshold_ture_dueTo_maxCollectionDuration_is_greater_than_buffered_event_collection_duration() - throws IOException, InterruptedException { - final int maxEvents = 500; - final ByteCount maxBytes = ByteCount.parse("500mb"); - final long maxCollectionDuration = 240; - - inMemoryBuffer = new InMemoryBufferFactory().getBuffer(); - - boolean isThresholdExceed = Boolean.FALSE; - synchronized (this) { - while (inMemoryBuffer.getEventCount() < 100) { - inMemoryBuffer.writeEvent(generateByteArray()); - isThresholdExceed = ThresholdValidator.checkThresholdExceed(inMemoryBuffer, - maxEvents, maxBytes, maxCollectionDuration); - if (isThresholdExceed) { - break; - } - wait(50); - } - } - assertFalse(isThresholdExceed, "Threshold exceeded"); - } - - private byte[] generateByteArray() { - byte[] bytes = new byte[10000]; - for (int i = 0; i < 10000; i++) { - bytes[i] = (byte) i; - } - return bytes; - } } diff --git a/data-prepper-plugins/http-sink/README.md b/data-prepper-plugins/http-sink/README.md new file mode 100644 index 0000000000..3e79121ce7 --- /dev/null +++ b/data-prepper-plugins/http-sink/README.md @@ -0,0 +1,187 @@ +# Http Sink + +This is the Data Prepper Http sink plugin that sends records to http/https endpoints. You can use the sink to send data to arbitrary HTTP Endpoints. + + +## Usages + +The Http sink should be configured as part of Data Prepper pipeline yaml file. + +### Response status + +* `200`: the request data has been successfully pushed to http endpoint. +* `500`: internal server error while process the request data. +* `400`: bad request error +* `404`: the http endpoint is not reachable +* `501`: the server does not recognize the request method and is incapable of supporting it for any resource + +### HTTP Basic authentication +``` +pipeline: + ... + sink: + - http: + authentication: + http_basic: + username: my-user + password: my_s3cr3t +``` + +### HTTP Bearer token authentication +``` +pipeline: + ... + sink: + - http: + authentication: + bearer_token: + client_id: 0oaafr4j79grYGC5d7 + client_secret: fFel-3FutCXAOndezEsOVlght6D6DR4OIt7G5D1_oJ6YtgU17JdyXmGf0M + token_url: https://localhost/oauth2/default/v1/token + grant_type: client_credentials + scope: httpSink +``` + +## Configuration + +- `url` The http/https endpoint url. + +- `proxy`(optional): A String of the address of a forward HTTP proxy. The format is like ":\". Examples: "example.com:8100", "http://example.com:8100", "112.112.112.112:8100". Note: port number cannot be omitted. + +- `codec` : This plugin is integrated with sink codec + +- `http_method` (Optional) : HttpMethod to be used. Default is POST. + +- `auth_type` (Optional): Authentication type configuration. By default, this runs an unauthenticated server. + +- `username`(optional): A string of username required for basic authentication + +- `password`(optional): A string of password required for basic authentication + +- `client_id`: It is the client id is the public identifier of your authorization server. + +- `client_secret` : It is a secret known only to the application and the authorization server. + +- `token_url`: The End point URL of the OAuth server.(Eg: /oauth2/default/v1/token) + +- `grant_type` (Optional) : This grant type refers to the way an application gets an access token. Example: client_credentials/refresh_token + +- `scope` (Optional) : This scope limit an application's access to a user's account. + +- `aws_sigv4`: A boolean flag to sign the HTTP request with AWS credentials. Default to `false`. For aws_sigv4, we don't need any auth_type or ssl + +- `aws` (Optional) : AWS configurations. See [AWS Configuration](#aws_configuration) for details. SigV4 is enabled by default when this option is used. If this option is present, `aws_` options are not expected to be present. If any of `aws_` options are present along with this, error is thrown. + +- `custom_header` (Optional) : A Map for custom headers such as AWS Sagemaker etc + +- `dlq_file`(optional): A String of absolute file path for DLQ failed output records. Defaults to null. + If not provided, failed records will be written into the default data-prepper log file (`logs/Data-Prepper.log`). If the `dlq` option is present along with this, an error is thrown. + +- `dlq` (optional): DLQ configurations. See [DLQ](https://github.com/opensearch-project/data-prepper/tree/main/data-prepper-plugins/failures-common/src/main/java/org/opensearch/dataprepper/plugins/dlq/README.md) for details. If the `dlq_file` option is present along with this, an error is thrown. + +- `buffer_type`(optional) : Buffer type can be in_memory/local_file. Default is in_memory + +- `max_retries`(optional): A number indicating the maximum number of times Http Sink should try to push the data to the Http arbitrary endpoint before considering it as failure. Defaults to `Integer.MAX_VALUE`. + +- `request_timout`(optional): A duration that represents the request timeout. Example: 1000ms, 5s etc +### HTTP Sink full pipeline +``` + sink: + - http: + url: http/https arbitrary endpoint url + proxy: proxy url + codec: + json: + http_method: "POST" + auth_type: "unauthenticated" + authentication: + http_basic: + username: "username" + password: "password" + bearer_token: + client_id: 0oaafr4j79segd7 + client_secret: fFel-3FutCXAOndezEsOVlghoJ6w0wNoaYtgU17JdyXmGf0M + token_url: token url + grant_type: client_credentials + scope: + insecure: false + insecure_skip_verify: false + ssl_certificate_file: "/full/path/to/certfile.crt" + ssl_key_file: "/full/path/to/keyfile.key" + buffer_type: "in_memory" + use_acm_cert_for_ssl: false + acm_certificate_arn: + custom_header: + header: ["value"] + aws_sigv4: false + dlq_file : + dlq: + s3: + bucket: + key_path_prefix: + webhook_url: + aws: + region: "us-east-2" + sts_role_arn: "arn:aws:iam::1234567890:role/data-prepper-s3source-execution-role" + service_name: lambda + threshold: + event_count: 5 + event_collect_timeout: PT2M + max_retries: 5 + request_timout: 20s +``` + +### SSL + +* insecure_skip_verify(Optional) => A `boolean` that enables mTLS/SSL. Default is ```false```. + * If set to false then the user has two options: + * Use default trust. This can allow for reaching many endpoints. The user does not need to provide any .crt/.key files. + * Allow the user to specify a .crt file for a certificate (no .key is required because this is the client). By the user providing the .crt file, the user is stating he trusts that certificate. We will still verify the signature match. + * If set to true, then skip any verification of the certificate. The user does not need to provide a .crt or .key file. +* insecure (Optional) => A `boolean` that allows http/https endpoints. Default is ```false```. + * If set to false, then only https:// URLs are permitted. Throw an InvalidPluginConfigurationException if the URL is configured with an http:// scheme in the URL. + * If set to true, then the user can provide both http:// https:// as the scheme. +* ssl_certificate_file(Optional) => A `String` that represents the SSL certificate chain file path or AWS S3 path. S3 path example `s3:///`. Required if `ssl` is set to `true` and `use_acm_certificate_for_ssl` is set to `false`. +* ssl_key_file(Optional) => A `String` that represents the SSL key file path or AWS S3 path. S3 path example `s3:///`. Only decrypted key file is supported. Required if `ssl` is set to `true` and `use_acm_certificate_for_ssl` is set to `false`. +* use_acm_certificate_for_ssl(Optional) : A `boolean` that enables mTLS/SSL using certificate and private key from AWS Certificate Manager (ACM). Default is `false`. +* acm_certificate_arn(Optional) : A `String` that represents the ACM certificate ARN. ACM certificate take preference over S3 or local file system certificate. Required if `use_acm_certificate_for_ssl` is set to `true`. +* acm_private_key_password(Optional): A `String` that represents the ACM private key password which that will be used to decrypt the private key. If it's not provided, a random password will be generated. +* acm_certificate_timeout_millis(Optional) : An `int` that represents the timeout in milliseconds for ACM to get certificates. Default value is `120000`. + +### AWS Configuration + +* `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 STS role to assume for requests to AWS. 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_header_overrides` (Optional): A map of header overrides to make when assuming the IAM role for the sink plugin. +* `sts_external_id` (Optional): An optional external ID to use when assuming an IAM role. +* `service_name` The AWS service name to who endpoint we are connecting to. Default: execute-api, Example: lambda, apigateway +### Threshold +* event_count => The event_count size should be between 0 and 10000000 +* maximum_size => The size of byte capacity, Default is 50mb +* event_collect_timeout => The event_collect timeout is between 1 and 3600 seconds + +## Metrics + +### Counter + +- `httpSinkRecordsSuccessCounter`: measures total number of records successfully pushed to http end points (200 response status code) by HTTP sink plugin. +- `httpSinkRecordsFailedCounter`: measures total number of records failed to pushed to http end points (500/400/404/501 response status code) by HTTP sink plugin. + +### End-to-End acknowledgements + +If the events received by the Http Sink have end-to-end acknowledgements enabled (which is tracked using the presence of EventHandle in the event received for processing), then upon successful posting to OpenSearch or upon successful write to DLQ, a positive acknowledgement is sent to the acknowledgementSetManager, otherwise a negative acknowledgement is sent. + +## Developer Guide + +This plugin is compatible with Java 8. 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) + +The integration tests for this plugin do not run as part of the Data Prepper build. + +The following command runs the integration tests: + +``` +./gradlew :data-prepper-plugins:http-sink:integrationTest -Dtests.http.sink.http.endpoint= +``` diff --git a/data-prepper-plugins/http-sink/build.gradle b/data-prepper-plugins/http-sink/build.gradle index b6b39ce120..1c466fc67f 100644 --- a/data-prepper-plugins/http-sink/build.gradle +++ b/data-prepper-plugins/http-sink/build.gradle @@ -7,6 +7,7 @@ dependencies { implementation 'io.micrometer:micrometer-core' implementation 'com.fasterxml.jackson.core:jackson-core' implementation 'com.fasterxml.jackson.core:jackson-databind' + implementation 'org.apache.commons:commons-compress:1.21' implementation libs.commons.compress implementation 'joda-time:joda-time:2.11.1' implementation project(':data-prepper-plugins:common') @@ -21,10 +22,46 @@ dependencies { implementation 'software.amazon.awssdk:auth' implementation libs.commons.lang3 implementation project(':data-prepper-plugins:failures-common') - implementation group: 'org.apache.httpcomponents.client5', name: 'httpclient5', version: '5.2' + implementation 'org.apache.httpcomponents.client5:httpclient5:5.2.1' + implementation 'com.github.scribejava:scribejava-core:8.3.3' + implementation project(path: ':data-prepper-core') + implementation project(':data-prepper-plugins:parse-json-processor') + implementation 'com.amazonaws:aws-java-sdk-sts:1.12.395' testImplementation project(':data-prepper-test-common') } test { useJUnitPlatform() +} + +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.http.sink.http.endpoint', System.getProperty('tests.http.sink.http.endpoint') + systemProperty 'tests.http.sink.region', System.getProperty('tests.http.sink.region') + systemProperty 'tests.http.sink.bucket', System.getProperty('tests.http.sink.bucket') + + filter { + includeTestsMatching '*IT' + } } \ No newline at end of file diff --git a/data-prepper-plugins/http-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/http/HttpSinkServiceIT.java b/data-prepper-plugins/http-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/http/HttpSinkServiceIT.java new file mode 100644 index 0000000000..47a9a18f85 --- /dev/null +++ b/data-prepper-plugins/http-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/http/HttpSinkServiceIT.java @@ -0,0 +1,154 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.http; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; +import com.fasterxml.jackson.dataformat.yaml.YAMLGenerator; +import io.micrometer.core.instrument.Counter; +import org.apache.hc.client5.http.impl.classic.HttpClientBuilder; +import org.apache.hc.client5.http.impl.classic.HttpClients; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.codec.OutputCodec; +import org.opensearch.dataprepper.model.configuration.PipelineDescription; +import org.opensearch.dataprepper.model.configuration.PluginModel; +import org.opensearch.dataprepper.model.configuration.PluginSetting; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.event.EventHandle; +import org.opensearch.dataprepper.model.event.JacksonEvent; +import org.opensearch.dataprepper.model.log.JacksonLog; +import org.opensearch.dataprepper.model.plugin.PluginFactory; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.plugins.accumulator.BufferFactory; +import org.opensearch.dataprepper.plugins.accumulator.InMemoryBufferFactory; +import org.opensearch.dataprepper.plugins.codec.json.NdjsonOutputCodec; +import org.opensearch.dataprepper.plugins.codec.json.NdjsonOutputConfig; +import org.opensearch.dataprepper.plugins.sink.http.configuration.HttpSinkConfiguration; +import org.opensearch.dataprepper.plugins.sink.http.configuration.ThresholdOptions; +import org.opensearch.dataprepper.plugins.sink.http.dlq.DlqPushHandler; +import org.opensearch.dataprepper.plugins.sink.http.service.HttpSinkService; +import org.opensearch.dataprepper.test.helper.ReflectivelySetField; + +import java.text.MessageFormat; +import java.time.Duration; +import java.util.Collection; +import java.util.LinkedList; +import java.util.UUID; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.plugins.sink.http.service.HttpSinkService.HTTP_SINK_RECORDS_SUCCESS_COUNTER; + +public class HttpSinkServiceIT { + + private String urlString; + + OutputCodec codec; + + private ObjectMapper objectMapper = new ObjectMapper(new YAMLFactory().enable(YAMLGenerator.Feature.USE_PLATFORM_LINE_BREAKS)); + + private String config = + " url: {0}\n" + + " http_method: POST\n" + + " auth_type: {1}\n" + + " codec:\n" + + " json:\n" + + " insecure_skip_verify: true\n" + + " threshold:\n" + + " event_count: 1"; + + private HttpSinkConfiguration httpSinkConfiguration; + + private BufferFactory bufferFactory; + + private DlqPushHandler dlqPushHandler; + + private PluginMetrics pluginMetrics; + + private PluginSetting pluginSetting; + + @BeforeEach + void setUp() throws JsonProcessingException{ + this.urlString = System.getProperty("tests.http.sink.http.endpoint"); + String[] values = { urlString,"unauthenticated"}; + final String configYaml = MessageFormat.format(config, values); + this.httpSinkConfiguration = objectMapper.readValue(configYaml, HttpSinkConfiguration.class); + } + + @Mock + private PipelineDescription pipelineDescription; + + private PluginFactory pluginFactory; + + @Mock + private Counter httpSinkRecordsSuccessCounter; + + @Mock + NdjsonOutputConfig ndjsonOutputConfig; + + public HttpSinkService createHttpSinkServiceUnderTest() throws NoSuchFieldException, IllegalAccessException { + this.pipelineDescription = mock(PipelineDescription.class); + this.pluginFactory = mock(PluginFactory.class); + this.httpSinkRecordsSuccessCounter = mock(Counter.class); + this.pluginMetrics = mock(PluginMetrics.class); + this.pluginSetting = mock(PluginSetting.class); + + when(pluginMetrics.counter(HTTP_SINK_RECORDS_SUCCESS_COUNTER)).thenReturn(httpSinkRecordsSuccessCounter); + when(pipelineDescription.getPipelineName()).thenReturn("http-plugin"); + ReflectivelySetField.setField(ThresholdOptions.class,httpSinkConfiguration.getThresholdOptions(),"eventCollectTimeOut", Duration.ofNanos(1)); + + final PluginModel codecConfiguration = httpSinkConfiguration.getCodec(); + final PluginSetting codecPluginSettings = new PluginSetting(codecConfiguration.getPluginName(), + codecConfiguration.getPluginSettings()); + this.ndjsonOutputConfig = mock(NdjsonOutputConfig.class); + codec = new NdjsonOutputCodec(ndjsonOutputConfig); + this.bufferFactory = new InMemoryBufferFactory(); + this.dlqPushHandler = new DlqPushHandler(httpSinkConfiguration.getDlqFile(), pluginFactory, + "bucket", + "arn", "region", + "keypath"); + + HttpClientBuilder httpClientBuilder = HttpClients.custom(); + + return new HttpSinkService( + httpSinkConfiguration, + bufferFactory, + dlqPushHandler, + codecPluginSettings, + null, + httpClientBuilder, + pluginMetrics, + pluginSetting, + codec, + null); + } + + private Collection> setEventQueue(final int records) { + final Collection> jsonObjects = new LinkedList<>(); + for (int i = 0; i < records; i++) + jsonObjects.add(createRecord()); + return jsonObjects; + } + + private static Record createRecord() { + final JacksonEvent event = JacksonLog.builder().withData("{\"name\":\""+ UUID.randomUUID() +"\"}").build(); + event.setEventHandle(mock(EventHandle.class)); + return new Record<>(event); + } + + @Test + public void http_endpoint_test_with_single_record() throws NoSuchFieldException, IllegalAccessException { + final HttpSinkService httpSinkService = createHttpSinkServiceUnderTest(); + final Collection> records = setEventQueue(1); + httpSinkService.output(records); + verify(httpSinkRecordsSuccessCounter).increment(records.size()); + } + +} diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/AwsRequestSigningApacheInterceptor.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/AwsRequestSigningApacheInterceptor.java similarity index 91% rename from data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/AwsRequestSigningApacheInterceptor.java rename to data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/AwsRequestSigningApacheInterceptor.java index b70740708d..84b13a17c9 100644 --- a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/AwsRequestSigningApacheInterceptor.java +++ b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/AwsRequestSigningApacheInterceptor.java @@ -10,18 +10,17 @@ * CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions * and limitations under the License. */ -package org.opensearch.dataprepper.plugins.sink; +package org.opensearch.dataprepper.plugins.sink.http; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.auth.DefaultAWSCredentialsProviderChain; import org.apache.hc.core5.http.HttpRequest; import org.apache.hc.core5.http.ClassicHttpRequest; import org.apache.hc.core5.http.EntityDetails; -import org.apache.hc.core5.http.HttpEntity; import org.apache.hc.core5.http.Header; import org.apache.hc.core5.http.NameValuePair; -import org.apache.hc.core5.http.ContentType; import org.apache.hc.core5.http.HttpHost; import org.apache.hc.core5.http.HttpRequestInterceptor; -import org.apache.hc.core5.http.io.entity.BasicHttpEntity; import org.apache.hc.core5.http.message.BasicHeader; import org.apache.hc.core5.http.protocol.HttpContext; import org.apache.hc.core5.net.URIBuilder; @@ -146,6 +145,8 @@ public void process(final HttpRequest request, final EntityDetails entity, final requestBuilder.rawQueryParameters(nvpToMapParams(uriBuilder.getQueryParams())); requestBuilder.headers(headerArrayToMap(request.getHeaders())); + AWSCredentials credentials = new DefaultAWSCredentialsProviderChain().getCredentials(); + ExecutionAttributes attributes = new ExecutionAttributes(); attributes.putAttribute(AwsSignerExecutionAttribute.AWS_CREDENTIALS, awsCredentialsProvider.resolveCredentials()); attributes.putAttribute(AwsSignerExecutionAttribute.SERVICE_SIGNING_NAME, service); @@ -156,16 +157,6 @@ public void process(final HttpRequest request, final EntityDetails entity, final // Now copy everything back request.setHeaders(mapToHeaderArray(signedRequest.headers())); - if (request instanceof ClassicHttpRequest) { - ClassicHttpRequest classicHttpRequest = - (ClassicHttpRequest) request; - if (classicHttpRequest.getEntity() != null) { - HttpEntity basicHttpEntity = new BasicHttpEntity(signedRequest.contentStreamProvider() - .orElseThrow(() -> new IllegalStateException("There must be content")) - .newStream(), ContentType.APPLICATION_JSON); - classicHttpRequest.setEntity(basicHttpEntity); - } - } } private URI buildUri(final HttpContext context, URIBuilder uriBuilder) throws IOException { @@ -239,4 +230,4 @@ private static Header[] mapToHeaderArray(final Map> mapHead } return headers; } -} +} \ No newline at end of file diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/FailedHttpResponseInterceptor.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/FailedHttpResponseInterceptor.java index 8a8d3b523c..f2961b4948 100644 --- a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/FailedHttpResponseInterceptor.java +++ b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/FailedHttpResponseInterceptor.java @@ -13,13 +13,7 @@ public class FailedHttpResponseInterceptor implements HttpResponseInterceptor { - public static final int ERROR_CODE_500 = 500; - - public static final int ERROR_CODE_400 = 400; - - public static final int ERROR_CODE_404 = 404; - - public static final int ERROR_CODE_501 = 501; + public static final int STATUS_CODE_200 = 200; private final String url; @@ -29,10 +23,7 @@ public FailedHttpResponseInterceptor(final String url){ @Override public void process(HttpResponse response, EntityDetails entity, HttpContext context) throws IOException { - if (response.getCode() == ERROR_CODE_500 || - response.getCode() == ERROR_CODE_400 || - response.getCode() == ERROR_CODE_404 || - response.getCode() == ERROR_CODE_501) { + if (response.getCode() != STATUS_CODE_200) { throw new IOException(String.format("url: %s , status code: %s", url,response.getCode())); } } diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/HTTPSink.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/HTTPSink.java index cc34b7b75d..3f43df2755 100644 --- a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/HTTPSink.java +++ b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/HTTPSink.java @@ -12,6 +12,7 @@ import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; +import org.opensearch.dataprepper.model.codec.OutputCodec; import org.opensearch.dataprepper.model.configuration.PipelineDescription; import org.opensearch.dataprepper.model.configuration.PluginModel; import org.opensearch.dataprepper.model.configuration.PluginSetting; @@ -20,7 +21,9 @@ import org.opensearch.dataprepper.model.plugin.PluginFactory; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.sink.AbstractSink; +import org.opensearch.dataprepper.model.sink.OutputCodecContext; import org.opensearch.dataprepper.model.sink.Sink; +import org.opensearch.dataprepper.model.sink.SinkContext; import org.opensearch.dataprepper.plugins.accumulator.BufferFactory; import org.opensearch.dataprepper.plugins.accumulator.BufferTypeOptions; import org.opensearch.dataprepper.plugins.accumulator.InMemoryBufferFactory; @@ -36,6 +39,7 @@ import org.slf4j.LoggerFactory; import java.util.Collection; +import java.util.Collections; import java.util.Objects; @DataPrepperPlugin(name = "http", pluginType = Sink.class, pluginConfigurationType = HttpSinkConfiguration.class) @@ -56,13 +60,19 @@ public class HTTPSink extends AbstractSink> { private DlqPushHandler dlqPushHandler; + private final OutputCodec codec; + + private final SinkContext sinkContext; + @DataPrepperPluginConstructor public HTTPSink(final PluginSetting pluginSetting, final HttpSinkConfiguration httpSinkConfiguration, final PluginFactory pluginFactory, final PipelineDescription pipelineDescription, + final SinkContext sinkContext, final AwsCredentialsSupplier awsCredentialsSupplier) { super(pluginSetting); + this.sinkContext = sinkContext != null ? sinkContext : new SinkContext(null, Collections.emptyList(), Collections.emptyList(), Collections.emptyList()); final PluginModel codecConfiguration = httpSinkConfiguration.getCodec(); final PluginSetting codecPluginSettings = new PluginSetting(codecConfiguration.getPluginName(), codecConfiguration.getPluginSettings()); @@ -74,19 +84,17 @@ public HTTPSink(final PluginSetting pluginSetting, this.bufferFactory = new InMemoryBufferFactory(); } - if(httpSinkConfiguration.getDlqFile() != null) - this.dlqPushHandler = new DlqPushHandler(httpSinkConfiguration.getDlqFile(), pluginFactory, - null, null, null, null); - - else if(Objects.nonNull(httpSinkConfiguration.getDlq())) - this.dlqPushHandler = new DlqPushHandler(httpSinkConfiguration.getDlqFile(), pluginFactory, - httpSinkConfiguration.getDlq().getPluginSettings().get(BUCKET).toString(), httpSinkConfiguration.getAwsAuthenticationOptions() - .getAwsStsRoleArn(), httpSinkConfiguration.getAwsAuthenticationOptions().getAwsRegion().toString(), - httpSinkConfiguration.getDlq().getPluginSettings().get(KEY_PATH).toString()); - + this.dlqPushHandler = new DlqPushHandler(httpSinkConfiguration.getDlqFile(), pluginFactory, + String.valueOf(httpSinkConfiguration.getDlqPluginSetting().get(BUCKET)), + httpSinkConfiguration.getDlqStsRoleARN() + ,httpSinkConfiguration.getDlqStsRegion(), + String.valueOf(httpSinkConfiguration.getDlqPluginSetting().get(KEY_PATH))); final HttpRequestRetryStrategy httpRequestRetryStrategy = new DefaultHttpRequestRetryStrategy(httpSinkConfiguration.getMaxUploadRetries(), TimeValue.of(httpSinkConfiguration.getHttpRetryInterval())); + if((!httpSinkConfiguration.isInsecure()) && (httpSinkConfiguration.isHttpUrl())){ + throw new InvalidPluginConfigurationException ("Cannot configure http url with insecure as false"); + } final HttpClientBuilder httpClientBuilder = HttpClients.custom() .setRetryStrategy(httpRequestRetryStrategy); @@ -98,6 +106,7 @@ else if(Objects.nonNull(httpSinkConfiguration.getDlq())) if(httpSinkConfiguration.isAwsSigv4() && httpSinkConfiguration.isValidAWSUrl()){ HttpSinkAwsService.attachSigV4(httpSinkConfiguration, httpClientBuilder, awsCredentialsSupplier); } + this.codec = pluginFactory.loadPlugin(OutputCodec.class, codecPluginSettings); this.httpSinkService = new HttpSinkService( httpSinkConfiguration, bufferFactory, @@ -106,7 +115,9 @@ else if(Objects.nonNull(httpSinkConfiguration.getDlq())) webhookService, httpClientBuilder, pluginMetrics, - pluginSetting); + pluginSetting, + codec, + OutputCodecContext.fromSinkContext(sinkContext)); } @Override diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/OAuthAccessTokenManager.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/OAuthAccessTokenManager.java new file mode 100644 index 0000000000..9a7d1b2bae --- /dev/null +++ b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/OAuthAccessTokenManager.java @@ -0,0 +1,69 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.http; + +import com.github.scribejava.core.builder.api.DefaultApi20; +import com.github.scribejava.core.model.OAuth2AccessToken; +import com.github.scribejava.core.oauth.OAuth20Service; +import com.github.scribejava.core.builder.ServiceBuilder; +import org.opensearch.dataprepper.plugins.sink.http.configuration.BearerTokenOptions; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.time.Instant; +import java.time.ZoneOffset; + +public class OAuthAccessTokenManager { + + private static final Logger LOG = LoggerFactory.getLogger(OAuthAccessTokenManager.class); + + public String getAccessToken(final BearerTokenOptions bearerTokenOptions) { + OAuth20Service service = getOAuth20ServiceObj(bearerTokenOptions); + OAuth2AccessToken accessTokenObj = null; + try { + if(bearerTokenOptions.getRefreshToken() != null) { + accessTokenObj = new OAuth2AccessToken(bearerTokenOptions.getAccessToken(), bearerTokenOptions.getRefreshToken()); + accessTokenObj = service.refreshAccessToken(accessTokenObj.getRefreshToken()); + + }else { + accessTokenObj = service.getAccessTokenClientCredentialsGrant(); + } + bearerTokenOptions.setRefreshToken(accessTokenObj.getRefreshToken()); + bearerTokenOptions.setAccessToken(accessTokenObj.getAccessToken()); + bearerTokenOptions.setTokenExpired(accessTokenObj.getExpiresIn()); + }catch (Exception e) { + LOG.info("Exception : "+ e.getMessage() ); + } + return bearerTokenOptions.getAccessToken(); + } + + + public boolean isTokenExpired(final Integer tokenExpired){ + final Instant systemCurrentTimeStamp = Instant.now().atOffset(ZoneOffset.UTC).toInstant(); + Instant accessTokenExpTimeStamp = systemCurrentTimeStamp.plusSeconds(tokenExpired); + if(systemCurrentTimeStamp.compareTo(accessTokenExpTimeStamp)>=0) { + return true; + } + return false; + } + + private OAuth20Service getOAuth20ServiceObj(BearerTokenOptions bearerTokenOptions){ + return new ServiceBuilder(bearerTokenOptions.getClientId()) + .apiSecret(bearerTokenOptions.getClientSecret()) + .defaultScope(bearerTokenOptions.getScope()) + .build(new DefaultApi20() { + @Override + public String getAccessTokenEndpoint() { + return bearerTokenOptions.getTokenUrl(); + } + + @Override + protected String getAuthorizationBaseUrl() { + return bearerTokenOptions.getTokenUrl(); + } + }); + } + +} diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/certificate/HttpClientSSLConnectionManager.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/certificate/HttpClientSSLConnectionManager.java index 6edefa0f11..c33a4d3f99 100644 --- a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/certificate/HttpClientSSLConnectionManager.java +++ b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/certificate/HttpClientSSLConnectionManager.java @@ -7,6 +7,7 @@ import org.apache.hc.client5.http.config.TlsConfig; import org.apache.hc.client5.http.impl.io.PoolingHttpClientConnectionManagerBuilder; import org.apache.hc.client5.http.io.HttpClientConnectionManager; +import org.apache.hc.client5.http.ssl.NoopHostnameVerifier; import org.apache.hc.client5.http.ssl.SSLConnectionSocketFactory; import org.apache.hc.client5.http.ssl.SSLConnectionSocketFactoryBuilder; import org.apache.hc.client5.http.ssl.TrustAllStrategy; @@ -15,14 +16,18 @@ import org.apache.hc.core5.ssl.SSLContexts; import org.apache.hc.core5.ssl.TrustStrategy; import org.apache.hc.core5.util.Timeout; -import org.opensearch.dataprepper.plugins.certificate.CertificateProvider; import org.opensearch.dataprepper.plugins.sink.http.configuration.HttpSinkConfiguration; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import javax.net.ssl.SSLContext; import java.io.ByteArrayInputStream; import java.io.InputStream; import java.nio.charset.StandardCharsets; +import java.security.KeyManagementException; import java.security.KeyStore; +import java.security.KeyStoreException; +import java.security.NoSuchAlgorithmException; import java.security.cert.Certificate; import java.security.cert.CertificateFactory; @@ -32,6 +37,8 @@ */ public class HttpClientSSLConnectionManager { + private static final Logger LOG = LoggerFactory.getLogger(HttpClientSSLConnectionManager.class); + /** * This method creates HttpClientConnectionManager for SSL certs authentication * @param sinkConfiguration HttpSinkConfiguration @@ -40,10 +47,8 @@ public class HttpClientSSLConnectionManager { */ public HttpClientConnectionManager createHttpClientConnectionManager(final HttpSinkConfiguration sinkConfiguration, final CertificateProviderFactory providerFactory){ - final CertificateProvider certificateProvider = providerFactory.getCertificateProvider(); - final org.opensearch.dataprepper.plugins.certificate.model.Certificate certificate = certificateProvider.getCertificate(); final SSLContext sslContext = sinkConfiguration.getSslCertificateFile() != null ? - getCAStrategy(new ByteArrayInputStream(certificate.getCertificate().getBytes(StandardCharsets.UTF_8))) : getTrustAllStrategy(); + getCAStrategy(new ByteArrayInputStream(providerFactory.getCertificateProvider().getCertificate().getCertificate().getBytes(StandardCharsets.UTF_8))) : getTrustAllStrategy(); SSLConnectionSocketFactory sslSocketFactory = SSLConnectionSocketFactoryBuilder.create() .setSslContext(sslContext) .build(); @@ -80,4 +85,17 @@ private SSLContext getTrustAllStrategy() { throw new RuntimeException(ex.getMessage(), ex); } } + + public HttpClientConnectionManager createHttpClientConnectionManagerWithoutValidation() throws NoSuchAlgorithmException, KeyStoreException, KeyManagementException { + { + return PoolingHttpClientConnectionManagerBuilder.create() + .setSSLSocketFactory(SSLConnectionSocketFactoryBuilder.create() + .setSslContext(SSLContextBuilder.create() + .loadTrustMaterial(TrustAllStrategy.INSTANCE) + .build()) + .setHostnameVerifier(NoopHostnameVerifier.INSTANCE) + .build()) + .build(); + } + } } diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/AuthenticationOptions.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/AuthenticationOptions.java new file mode 100644 index 0000000000..1622967696 --- /dev/null +++ b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/AuthenticationOptions.java @@ -0,0 +1,24 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.http.configuration; + +import com.fasterxml.jackson.annotation.JsonProperty; + +public class AuthenticationOptions { + + @JsonProperty("http_basic") + private BasicAuthCredentials httpBasic; + + @JsonProperty("bearer_token") + private BearerTokenOptions bearerTokenOptions; + + public BasicAuthCredentials getHttpBasic() { + return httpBasic; + } + + public BearerTokenOptions getBearerTokenOptions() { + return bearerTokenOptions; + } +} diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/AwsAuthenticationOptions.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/AwsAuthenticationOptions.java index 7a7be57d34..703c635fa2 100644 --- a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/AwsAuthenticationOptions.java +++ b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/AwsAuthenticationOptions.java @@ -11,6 +11,8 @@ import java.util.Map; public class AwsAuthenticationOptions { + + private static final String DEFAULT_SERVICE_NAME = "execute-api"; @JsonProperty("region") @Size(min = 1, message = "Region cannot be empty string") private String awsRegion; @@ -27,6 +29,10 @@ public class AwsAuthenticationOptions { @Size(max = 5, message = "sts_header_overrides supports a maximum of 5 headers to override") private Map awsStsHeaderOverrides; + @JsonProperty("service_name") + @Size(min = 1, message = "Service Name cannot be empty") + private String serviceName = DEFAULT_SERVICE_NAME; + public Region getAwsRegion() { return awsRegion != null ? Region.of(awsRegion) : null; } @@ -42,4 +48,8 @@ public String getAwsStsExternalId() { public Map getAwsStsHeaderOverrides() { return awsStsHeaderOverrides; } + + public String getServiceName() { + return serviceName; + } } \ No newline at end of file diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/BasicAuthCredentials.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/BasicAuthCredentials.java new file mode 100644 index 0000000000..9f5defd1ef --- /dev/null +++ b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/BasicAuthCredentials.java @@ -0,0 +1,25 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.http.configuration; + +import com.fasterxml.jackson.annotation.JsonProperty; + +public class BasicAuthCredentials { + + + @JsonProperty("username") + private String username; + + @JsonProperty("password") + private String password; + + public String getUsername() { + return username; + } + + public String getPassword() { + return password; + } +} diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/BearerTokenOptions.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/BearerTokenOptions.java new file mode 100644 index 0000000000..89ec7dd834 --- /dev/null +++ b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/BearerTokenOptions.java @@ -0,0 +1,82 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.http.configuration; + +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.constraints.NotNull; + +public class BearerTokenOptions { + + @JsonProperty("client_id") + @NotNull(message = "client id type is mandatory for refresh token") + private String clientId; + + @JsonProperty("client_secret") + @NotNull(message = "client secret type is mandatory for refresh token") + private String clientSecret; + + @JsonProperty("token_url") + @NotNull(message = "token url type is mandatory for refresh token") + private String tokenUrl; + + @JsonProperty("grant_type") + @NotNull(message = "grant type is mandatory for refresh token") + private String grantType; + + @JsonProperty("scope") + @NotNull(message = "scope is mandatory for refresh token") + private String scope; + + private String refreshToken; + + private String accessToken; + + private Integer tokenExpired; + + public String getScope() { + return scope; + } + + public String getGrantType() { + return grantType; + } + + public String getRefreshToken() { + return refreshToken; + } + + public String getAccessToken() { + return accessToken; + } + + public Integer getTokenExpired() { + return tokenExpired; + } + + public String getClientId() { + return clientId; + } + + public String getClientSecret() { + return clientSecret; + } + + public void setRefreshToken(String refreshToken) { + this.refreshToken = refreshToken; + } + + public void setAccessToken(String accessToken) { + this.accessToken = accessToken; + } + + public void setTokenExpired(Integer tokenExpired) { + this.tokenExpired = tokenExpired; + } + + public String getTokenUrl() { + return tokenUrl; + } + +} diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/CustomHeaderOptions.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/CustomHeaderOptions.java deleted file mode 100644 index f6b9f7bbd8..0000000000 --- a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/CustomHeaderOptions.java +++ /dev/null @@ -1,52 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.dataprepper.plugins.sink.http.configuration; - -import com.fasterxml.jackson.annotation.JsonProperty; - -public class CustomHeaderOptions { - - @JsonProperty("X-Amzn-SageMaker-Custom-Attributes") - private String customAttributes; - - @JsonProperty("X-Amzn-SageMaker-Target-Model") - private String targetModel; - - @JsonProperty("X-Amzn-SageMaker-Target-Variant") - private String targetVariant; - - @JsonProperty("X-Amzn-SageMaker-Target-Container-Hostname") - private String targetContainerHostname; - - @JsonProperty("X-Amzn-SageMaker-Inference-Id") - private String inferenceId; - - @JsonProperty("X-Amzn-SageMaker-Enable-Explanations") - private String enableExplanations; - - public String getCustomAttributes() { - return customAttributes; - } - - public String getTargetModel() { - return targetModel; - } - - public String getTargetVariant() { - return targetVariant; - } - - public String getTargetContainerHostname() { - return targetContainerHostname; - } - - public String getInferenceId() { - return inferenceId; - } - - public String getEnableExplanations() { - return enableExplanations; - } -} diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/HttpSinkConfiguration.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/HttpSinkConfiguration.java index 23d87ccfc1..e1e9ac9c57 100644 --- a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/HttpSinkConfiguration.java +++ b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/HttpSinkConfiguration.java @@ -16,6 +16,7 @@ import java.time.Duration; import java.util.List; import java.util.Map; +import java.util.Objects; public class HttpSinkConfiguration { @@ -23,7 +24,7 @@ public class HttpSinkConfiguration { private static final int DEFAULT_WORKERS = 1; - static final boolean DEFAULT_SSL = false; + static final boolean DEFAULT_INSECURE = false; private static final String S3_PREFIX = "s3://"; @@ -31,6 +32,11 @@ public class HttpSinkConfiguration { static final String SSL_KEY_FILE = "sslKeyFile"; static final String SSL = "ssl"; static final String AWS_REGION = "awsRegion"; + + + public static final String STS_REGION = "sts_region"; + + public static final String STS_ROLE_ARN = "sts_role_arn"; static final boolean DEFAULT_USE_ACM_CERT_FOR_SSL = false; static final int DEFAULT_ACM_CERT_ISSUE_TIME_OUT_MILLIS = 120000; public static final String SSL_IS_ENABLED = "%s is enabled"; @@ -39,6 +45,8 @@ public class HttpSinkConfiguration { private static final String HTTPS = "https"; + private static final String HTTP = "http"; + private static final String AWS_HOST_AMAZONAWS_COM = "amazonaws.com"; private static final String AWS_HOST_API_AWS = "api.aws"; @@ -65,7 +73,7 @@ public class HttpSinkConfiguration { private AuthTypeOptions authType = AuthTypeOptions.UNAUTHENTICATED; @JsonProperty("authentication") - private PluginModel authentication; + private AuthenticationOptions authentication; @JsonProperty("ssl_certificate_file") private String sslCertificateFile; @@ -114,8 +122,14 @@ public class HttpSinkConfiguration { @JsonProperty("acm_cert_issue_time_out_millis") private long acmCertIssueTimeOutMillis = DEFAULT_ACM_CERT_ISSUE_TIME_OUT_MILLIS; - @JsonProperty("ssl") - private boolean ssl = DEFAULT_SSL; + @JsonProperty("insecure") + private boolean insecure = DEFAULT_INSECURE; + + @JsonProperty("insecure_skip_verify") + private boolean insecureSkipVerify = DEFAULT_INSECURE; + + @JsonProperty("request_timout") + private Duration requestTimout; @JsonProperty("http_retry_interval") private Duration httpRetryInterval = DEFAULT_HTTP_RETRY_INTERVAL; @@ -127,8 +141,8 @@ public String getUrl() { return url; } - public boolean isSsl() { - return ssl; + public boolean isInsecureSkipVerify() { + return insecureSkipVerify; } public Duration getHttpRetryInterval() { @@ -160,7 +174,7 @@ public void validateAndInitializeCertAndKeyFileInS3() { if (useAcmCertForSSL) { validateSSLArgument(String.format(SSL_IS_ENABLED, useAcmCertForSSL), acmCertificateArn, acmCertificateArn); validateSSLArgument(String.format(SSL_IS_ENABLED, useAcmCertForSSL), awsAuthenticationOptions.getAwsRegion().toString(), AWS_REGION); - } else if(ssl) { + } else if(!insecureSkipVerify) { validateSSLCertificateFiles(); certAndKeyFileInS3 = isSSLCertificateLocatedInS3(); if (certAndKeyFileInS3) { @@ -205,7 +219,7 @@ public AuthTypeOptions getAuthType() { return authType; } - public PluginModel getAuthentication() { + public AuthenticationOptions getAuthentication() { return authentication; } @@ -260,4 +274,33 @@ public boolean isValidAWSUrl() { } return false; } + + public String getDlqStsRoleARN(){ + return Objects.nonNull(getDlqPluginSetting().get(STS_ROLE_ARN)) ? + String.valueOf(getDlqPluginSetting().get(STS_ROLE_ARN)) : + awsAuthenticationOptions.getAwsStsRoleArn(); + } + + public String getDlqStsRegion(){ + return Objects.nonNull(getDlqPluginSetting().get(STS_REGION)) ? + String.valueOf(getDlqPluginSetting().get(STS_REGION)) : + awsAuthenticationOptions.getAwsRegion().toString(); + } + + public Map getDlqPluginSetting(){ + return dlq != null ? dlq.getPluginSettings() : Map.of(); + } + + public boolean isInsecure() { + return insecure; + } + + public Duration getRequestTimout() { + return requestTimout; + } + + public boolean isHttpUrl() { + URL parsedUrl = HttpSinkUtil.getURLByUrlString(url); + return parsedUrl.getProtocol().equals(HTTP); + } } diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/UrlConfigurationOption.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/UrlConfigurationOption.java deleted file mode 100644 index 4ddd2dcf06..0000000000 --- a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/configuration/UrlConfigurationOption.java +++ /dev/null @@ -1,58 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.dataprepper.plugins.sink.http.configuration; - -import com.fasterxml.jackson.annotation.JsonProperty; -import jakarta.validation.constraints.NotNull; -import org.opensearch.dataprepper.model.configuration.PluginModel; - -public class UrlConfigurationOption { - - private static final int DEFAULT_WORKERS = 1; - - @NotNull - @JsonProperty("url") - private String url; - - @JsonProperty("workers") - private Integer workers = DEFAULT_WORKERS; - - @JsonProperty("proxy") - private String proxy; - - @JsonProperty("codec") - private PluginModel codec; - - @JsonProperty("http_method") - private HTTPMethodOptions httpMethod; - - @JsonProperty("auth_type") - private AuthTypeOptions authType; - - public String getUrl() { - return url; - } - - public Integer getWorkers() { - return workers; - } - - public String getProxy() { - return proxy; - } - - public PluginModel getCodec() { - return codec; - } - - public HTTPMethodOptions getHttpMethod() { - return httpMethod; - } - - public AuthTypeOptions getAuthType() { - return authType; - } - -} diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/dlq/DlqPushHandler.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/dlq/DlqPushHandler.java index 05f29d78c5..9c00561e8e 100644 --- a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/dlq/DlqPushHandler.java +++ b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/dlq/DlqPushHandler.java @@ -4,6 +4,8 @@ */ package org.opensearch.dataprepper.plugins.sink.http.dlq; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.ObjectWriter; import io.micrometer.core.instrument.util.StringUtils; import org.opensearch.dataprepper.metrics.MetricNames; import org.opensearch.dataprepper.model.configuration.PluginSetting; @@ -52,6 +54,8 @@ public class DlqPushHandler { private DlqProvider dlqProvider; + private ObjectWriter objectWriter; + public DlqPushHandler(final String dlqFile, final PluginFactory pluginFactory, final String bucket, @@ -60,6 +64,7 @@ public DlqPushHandler(final String dlqFile, final String dlqPathPrefix) { if(dlqFile != null) { this.dlqFile = dlqFile; + this.objectWriter = new ObjectMapper().writer(); }else{ this.dlqProvider = getDlqProvider(pluginFactory,bucket,stsRoleArn,awsRegion,dlqPathPrefix); } @@ -76,7 +81,8 @@ public void perform(final PluginSetting pluginSetting, private void writeToFile(Object failedData) { try(BufferedWriter dlqFileWriter = Files.newBufferedWriter(Paths.get(dlqFile), StandardOpenOption.CREATE, StandardOpenOption.APPEND)) { - dlqFileWriter.write(failedData.toString()); + dlqFileWriter.write(objectWriter.writeValueAsString(failedData)+"\n"); + } catch (IOException e) { LOG.error("Exception while writing failed data to DLQ file Exception: ",e); } diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/dlq/FailedDlqData.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/dlq/FailedDlqData.java index 18875fcafa..024710ea8c 100644 --- a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/dlq/FailedDlqData.java +++ b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/dlq/FailedDlqData.java @@ -4,52 +4,56 @@ */ package org.opensearch.dataprepper.plugins.sink.http.dlq; -import com.fasterxml.jackson.annotation.JsonIgnore; -import org.opensearch.dataprepper.plugins.sink.http.HttpEndPointResponse; - public class FailedDlqData { - private final HttpEndPointResponse endPointResponse; - @JsonIgnore - private final String bufferData; + private String url; + + private int status; + + private String message; public FailedDlqData(final Builder builder) { - this.endPointResponse = builder.endPointResponse; - this.bufferData = builder.bufferData; + this.status = builder.status; + this.message = builder.message; + this.url = builder.url; + } + + public String getUrl() { + return url; } - public HttpEndPointResponse getEndPointResponse() { - return endPointResponse; + public int getStatus() { + return status; } - public String getBufferData() { - return bufferData; + public String getMessage() { + return message; } public static Builder builder() { return new Builder(); } - @Override - public String toString() { - return "{" + - "endPointResponse=" + endPointResponse + - ", bufferData='" + bufferData + '\'' + - '}'; - } public static class Builder { - private HttpEndPointResponse endPointResponse; + private String url; + + private int status; + + private String message; - private String bufferData; + public Builder withUrl(String url) { + this.url = url; + return this; + } - public Builder withEndPointResponses(HttpEndPointResponse endPointResponses) { - this.endPointResponse = endPointResponses; + public Builder withStatus(int status) { + this.status = status; return this; } - public Builder withBufferData(String bufferData) { - this.bufferData = bufferData; + public Builder withMessage(String message) { + this.message = message; return this; } diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/handler/BearerTokenAuthHttpSinkHandler.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/handler/BearerTokenAuthHttpSinkHandler.java index 0ed8567fbc..a48ffc7936 100644 --- a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/handler/BearerTokenAuthHttpSinkHandler.java +++ b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/handler/BearerTokenAuthHttpSinkHandler.java @@ -4,7 +4,13 @@ */ package org.opensearch.dataprepper.plugins.sink.http.handler; +import com.fasterxml.jackson.databind.ObjectMapper; import org.apache.hc.client5.http.io.HttpClientConnectionManager; +import org.opensearch.dataprepper.plugins.sink.http.OAuthAccessTokenManager; +import org.opensearch.dataprepper.plugins.sink.http.configuration.BearerTokenOptions; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + import org.opensearch.dataprepper.plugins.sink.http.FailedHttpResponseInterceptor; /** @@ -12,21 +18,31 @@ */ public class BearerTokenAuthHttpSinkHandler implements MultiAuthHttpSinkHandler { + private static final Logger LOG = LoggerFactory.getLogger(BearerTokenAuthHttpSinkHandler.class); + public static final String AUTHORIZATION = "Authorization"; private final HttpClientConnectionManager httpClientConnectionManager; - private final String bearerTokenString; + private final BearerTokenOptions bearerTokenOptions; + + private final ObjectMapper objectMapper; + + private OAuthAccessTokenManager oAuthRefreshTokenManager; - public BearerTokenAuthHttpSinkHandler(final String bearerTokenString, - final HttpClientConnectionManager httpClientConnectionManager){ - this.bearerTokenString = bearerTokenString; + public BearerTokenAuthHttpSinkHandler(final BearerTokenOptions bearerTokenOptions, + final HttpClientConnectionManager httpClientConnectionManager, + final OAuthAccessTokenManager oAuthRefreshTokenManager){ + this.bearerTokenOptions = bearerTokenOptions; this.httpClientConnectionManager = httpClientConnectionManager; + this.objectMapper = new ObjectMapper(); + this.oAuthRefreshTokenManager = oAuthRefreshTokenManager; } @Override public HttpAuthOptions authenticate(final HttpAuthOptions.Builder httpAuthOptionsBuilder) { - httpAuthOptionsBuilder.getClassicHttpRequestBuilder().addHeader(AUTHORIZATION,bearerTokenString); + httpAuthOptionsBuilder.getClassicHttpRequestBuilder() + .addHeader(AUTHORIZATION, oAuthRefreshTokenManager.getAccessToken(bearerTokenOptions)); httpAuthOptionsBuilder.setHttpClientBuilder(httpAuthOptionsBuilder.build().getHttpClientBuilder() .setConnectionManager(httpClientConnectionManager) .addResponseInterceptorLast(new FailedHttpResponseInterceptor(httpAuthOptionsBuilder.getUrl()))); diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkAwsService.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkAwsService.java index bcbe0ca5b5..3b838b9c74 100644 --- a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkAwsService.java +++ b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkAwsService.java @@ -7,7 +7,7 @@ import org.apache.hc.client5.http.impl.classic.HttpClientBuilder; import org.opensearch.dataprepper.aws.api.AwsCredentialsOptions; import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; -import org.opensearch.dataprepper.plugins.sink.AwsRequestSigningApacheInterceptor; +import org.opensearch.dataprepper.plugins.sink.http.AwsRequestSigningApacheInterceptor; import org.opensearch.dataprepper.plugins.sink.http.configuration.HttpSinkConfiguration; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -18,14 +18,13 @@ public class HttpSinkAwsService { private static final Logger LOG = LoggerFactory.getLogger(HttpSinkAwsService.class); public static final String AWS_SIGV4 = "aws_sigv4"; - private static final String AOS_SERVICE_NAME = "http-endpoint"; public static void attachSigV4(final HttpSinkConfiguration httpSinkConfiguration, final HttpClientBuilder httpClientBuilder, final AwsCredentialsSupplier awsCredentialsSupplier) { LOG.info("{} is set, will sign requests using AWSRequestSigningApacheInterceptor", AWS_SIGV4); final Aws4Signer aws4Signer = Aws4Signer.create(); final AwsCredentialsOptions awsCredentialsOptions = createAwsCredentialsOptions(httpSinkConfiguration); final AwsCredentialsProvider credentialsProvider = awsCredentialsSupplier.getProvider(awsCredentialsOptions); - httpClientBuilder.addRequestInterceptorLast(new AwsRequestSigningApacheInterceptor(AOS_SERVICE_NAME, aws4Signer, + httpClientBuilder.addRequestInterceptorLast(new AwsRequestSigningApacheInterceptor(httpSinkConfiguration.getAwsAuthenticationOptions().getServiceName(), aws4Signer, credentialsProvider, httpSinkConfiguration.getAwsAuthenticationOptions().getAwsRegion())); } diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkService.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkService.java index 8b76ab32d5..480886d52b 100644 --- a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkService.java +++ b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkService.java @@ -5,24 +5,29 @@ package org.opensearch.dataprepper.plugins.sink.http.service; import io.micrometer.core.instrument.Counter; +import org.apache.hc.client5.http.config.RequestConfig; import org.apache.hc.client5.http.impl.classic.HttpClientBuilder; import org.apache.hc.client5.http.io.HttpClientConnectionManager; import org.apache.hc.client5.http.protocol.HttpClientContext; +import org.apache.hc.core5.http.ContentType; import org.apache.hc.core5.http.HttpStatus; import org.apache.hc.core5.http.io.support.ClassicRequestBuilder; +import org.apache.hc.core5.util.Timeout; import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.codec.OutputCodec; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.event.EventHandle; import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.sink.OutputCodecContext; import org.opensearch.dataprepper.model.types.ByteCount; import org.opensearch.dataprepper.plugins.accumulator.Buffer; import org.opensearch.dataprepper.plugins.accumulator.BufferFactory; +import org.opensearch.dataprepper.plugins.sink.http.HttpEndPointResponse; +import org.opensearch.dataprepper.plugins.sink.http.OAuthAccessTokenManager; import org.opensearch.dataprepper.plugins.sink.ThresholdValidator; -import org.opensearch.dataprepper.plugins.sink.http.FailedHttpResponseInterceptor; -import org.opensearch.dataprepper.plugins.sink.http.HttpEndPointResponse; import org.opensearch.dataprepper.plugins.sink.http.certificate.CertificateProviderFactory; import org.opensearch.dataprepper.plugins.sink.http.certificate.HttpClientSSLConnectionManager; import org.opensearch.dataprepper.plugins.sink.http.configuration.AuthTypeOptions; @@ -33,13 +38,18 @@ import org.opensearch.dataprepper.plugins.sink.http.handler.BasicAuthHttpSinkHandler; import org.opensearch.dataprepper.plugins.sink.http.handler.BearerTokenAuthHttpSinkHandler; import org.opensearch.dataprepper.plugins.sink.http.handler.HttpAuthOptions; +import org.opensearch.dataprepper.plugins.sink.http.FailedHttpResponseInterceptor; import org.opensearch.dataprepper.plugins.sink.http.handler.MultiAuthHttpSinkHandler; import org.opensearch.dataprepper.plugins.sink.http.util.HttpSinkUtil; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.IOException; +import java.io.OutputStream; +import java.security.KeyManagementException; +import java.security.KeyStoreException; +import java.security.NoSuchAlgorithmException; import java.util.Collection; import java.util.HashMap; import java.util.LinkedList; @@ -56,15 +66,8 @@ public class HttpSinkService { private static final Logger LOG = LoggerFactory.getLogger(HttpSinkService.class); - public static final String USERNAME = "username"; - - public static final String PASSWORD = "password"; - - public static final String TOKEN = "token"; - - public static final String BEARER = "Bearer "; - public static final String HTTP_SINK_RECORDS_SUCCESS_COUNTER = "httpSinkRecordsSuccessPushToEndPoint"; + public static final String HTTP_SINK_RECORDS_FAILED_COUNTER = "httpSinkRecordsFailedToPushEndPoint"; private final Collection bufferedEventHandles; @@ -73,7 +76,7 @@ public class HttpSinkService { private final BufferFactory bufferFactory; - private final Map httpAuthOptions; + private final Map httpAuthOptions; private DlqPushHandler dlqPushHandler; @@ -93,6 +96,8 @@ public class HttpSinkService { private final Counter httpSinkRecordsFailedCounter; + private final OAuthAccessTokenManager oAuthAccessTokenManager; + private CertificateProviderFactory certificateProviderFactory; private WebhookService webhookService; @@ -103,6 +108,12 @@ public class HttpSinkService { private final PluginSetting httpPluginSetting; + private MultiAuthHttpSinkHandler multiAuthHttpSinkHandler; + + private final OutputCodec codec; + + private final OutputCodecContext codecContext; + public HttpSinkService(final HttpSinkConfiguration httpSinkConfiguration, final BufferFactory bufferFactory, final DlqPushHandler dlqPushHandler, @@ -110,7 +121,10 @@ public HttpSinkService(final HttpSinkConfiguration httpSinkConfiguration, final WebhookService webhookService, final HttpClientBuilder httpClientBuilder, final PluginMetrics pluginMetrics, - final PluginSetting httpPluginSetting){ + final PluginSetting httpPluginSetting, + final OutputCodec codec, + final OutputCodecContext codecContext) { + this.httpSinkConfiguration = httpSinkConfiguration; this.bufferFactory = bufferFactory; this.dlqPushHandler = dlqPushHandler; @@ -123,15 +137,25 @@ public HttpSinkService(final HttpSinkConfiguration httpSinkConfiguration, this.maxBytes = httpSinkConfiguration.getThresholdOptions().getMaximumSize(); this.maxCollectionDuration = httpSinkConfiguration.getThresholdOptions().getEventCollectTimeOut().getSeconds(); this.httpPluginSetting = httpPluginSetting; - if (httpSinkConfiguration.isSsl() || httpSinkConfiguration.useAcmCertForSSL()) { + this.oAuthAccessTokenManager = new OAuthAccessTokenManager(); + + if ((!httpSinkConfiguration.isInsecureSkipVerify()) || (httpSinkConfiguration.useAcmCertForSSL())) { this.certificateProviderFactory = new CertificateProviderFactory(httpSinkConfiguration); - httpSinkConfiguration.validateAndInitializeCertAndKeyFileInS3(); this.httpClientConnectionManager = new HttpClientSSLConnectionManager() .createHttpClientConnectionManager(httpSinkConfiguration, certificateProviderFactory); } + else{ + try { + this.httpClientConnectionManager = new HttpClientSSLConnectionManager().createHttpClientConnectionManagerWithoutValidation(); + }catch(NoSuchAlgorithmException | KeyStoreException | KeyManagementException ex){ + LOG.error("Exception while insecure_skip_verify is true ",ex); + } + } this.httpAuthOptions = buildAuthHttpSinkObjectsByConfig(httpSinkConfiguration); this.httpSinkRecordsSuccessCounter = pluginMetrics.counter(HTTP_SINK_RECORDS_SUCCESS_COUNTER); this.httpSinkRecordsFailedCounter = pluginMetrics.counter(HTTP_SINK_RECORDS_FAILED_COUNTER); + this.codec= codec; + this.codecContext = codecContext; } /** @@ -144,28 +168,35 @@ public void output(Collection> records) { this.currentBuffer = bufferFactory.getBuffer(); } try { + OutputStream outputStream = currentBuffer.getOutputStream(); records.forEach(record -> { - final Event event = record.getData(); try { - currentBuffer.writeEvent(event.toJsonString().getBytes()); - } catch (IOException e) { - throw new RuntimeException(e); - } - if (event.getEventHandle() != null) { - this.bufferedEventHandles.add(event.getEventHandle()); - } - if (ThresholdValidator.checkThresholdExceed(currentBuffer, maxEvents, maxBytes, maxCollectionDuration)) { - final HttpEndPointResponse failedHttpEndPointResponses = pushToEndPoint(getCurrentBufferData(currentBuffer)); - if (failedHttpEndPointResponses != null) { - logFailedData(failedHttpEndPointResponses, getCurrentBufferData(currentBuffer)); - } else { - LOG.info("data pushed to the end point successfully"); + final Event event = record.getData(); + if(currentBuffer.getEventCount() == 0) { + codec.start(outputStream,event , codecContext); } - currentBuffer = bufferFactory.getBuffer(); - releaseEventHandles(Boolean.TRUE); + codec.writeEvent(event, outputStream); + int count = currentBuffer.getEventCount() +1; + currentBuffer.setEventCount(count); + if(event.getEventHandle() != null) { + bufferedEventHandles.add(event.getEventHandle()); + } + if (ThresholdValidator.checkThresholdExceed(currentBuffer, maxEvents, maxBytes, maxCollectionDuration)) { + codec.complete(outputStream); + final HttpEndPointResponse failedHttpEndPointResponses = pushToEndPoint(getCurrentBufferData(currentBuffer)); + if (failedHttpEndPointResponses != null) { + logFailedData(failedHttpEndPointResponses); + releaseEventHandles(Boolean.FALSE); + } else { + LOG.info("data pushed to the end point successfully"); + releaseEventHandles(Boolean.TRUE); + } + currentBuffer = bufferFactory.getBuffer(); + }} + catch (IOException e) { + throw new RuntimeException(e); }}); - }finally { reentrantLock.unlock(); } @@ -182,11 +213,14 @@ private byte[] getCurrentBufferData(final Buffer currentBuffer) { /** * * This method logs Failed Data to DLQ and Webhook * @param endPointResponses HttpEndPointResponses. - * @param currentBufferData Current bufferData. */ - private void logFailedData(final HttpEndPointResponse endPointResponses, final byte[] currentBufferData) { + private void logFailedData(final HttpEndPointResponse endPointResponses) { FailedDlqData failedDlqData = - FailedDlqData.builder().withBufferData(new String(currentBufferData)).withEndPointResponses(endPointResponses).build(); + FailedDlqData.builder() + .withUrl(endPointResponses.getUrl()) + .withMessage(endPointResponses.getErrMessage()) + .withStatus(endPointResponses.getStatusCode()).build(); + LOG.info("Failed to push the data. Failed DLQ Data: {}",failedDlqData); logFailureForDlqObjects(failedDlqData); @@ -210,8 +244,10 @@ private HttpEndPointResponse pushToEndPoint(final byte[] currentBufferData) { HttpEndPointResponse httpEndPointResponses = null; final ClassicRequestBuilder classicHttpRequestBuilder = httpAuthOptions.get(httpSinkConfiguration.getUrl()).getClassicHttpRequestBuilder(); - classicHttpRequestBuilder.setEntity(new String(currentBufferData)); + classicHttpRequestBuilder.setEntity(currentBufferData, ContentType.APPLICATION_JSON); try { + if(AuthTypeOptions.BEARER_TOKEN.equals(httpSinkConfiguration.getAuthType())) + accessTokenIfExpired(httpSinkConfiguration.getAuthentication().getBearerTokenOptions().getTokenExpired(),httpSinkConfiguration.getUrl()); httpAuthOptions.get(httpSinkConfiguration.getUrl()).getHttpClientBuilder().build() .execute(classicHttpRequestBuilder.build(), HttpClientContext.create()); LOG.info("No of Records successfully pushed to endpoint {}", httpSinkConfiguration.getUrl() +" " + currentBuffer.getEventCount()); @@ -248,16 +284,17 @@ private void logFailureForWebHook(final FailedDlqData failedDlqData){ */ private HttpAuthOptions getAuthHandlerByConfig(final AuthTypeOptions authType, final HttpAuthOptions.Builder authOptions){ - MultiAuthHttpSinkHandler multiAuthHttpSinkHandler = null; switch(authType) { case HTTP_BASIC: - String username = httpSinkConfiguration.getAuthentication().getPluginSettings().get(USERNAME).toString(); - String password = httpSinkConfiguration.getAuthentication().getPluginSettings().get(PASSWORD).toString(); - multiAuthHttpSinkHandler = new BasicAuthHttpSinkHandler(username,password,httpClientConnectionManager); + multiAuthHttpSinkHandler = new BasicAuthHttpSinkHandler( + httpSinkConfiguration.getAuthentication().getHttpBasic().getUsername(), + httpSinkConfiguration.getAuthentication().getHttpBasic().getPassword(), + httpClientConnectionManager); break; case BEARER_TOKEN: - String token = httpSinkConfiguration.getAuthentication().getPluginSettings().get(TOKEN).toString(); - multiAuthHttpSinkHandler = new BearerTokenAuthHttpSinkHandler(BEARER + token,httpClientConnectionManager); + multiAuthHttpSinkHandler = new BearerTokenAuthHttpSinkHandler( + httpSinkConfiguration.getAuthentication().getBearerTokenOptions(), + httpClientConnectionManager, oAuthAccessTokenManager); break; case UNAUTHENTICATED: default: @@ -285,10 +322,17 @@ private Map buildAuthHttpSinkObjectsByConfig(final HttpS if(Objects.nonNull(httpSinkConfiguration.getCustomHeaderOptions())) addCustomHeaders(classicRequestBuilder,httpSinkConfiguration.getCustomHeaderOptions()); + if(httpSinkConfiguration.isAwsSigv4() && httpSinkConfiguration.isValidAWSUrl()){ + classicRequestBuilder.addHeader("x-amz-content-sha256","required"); + } + if(Objects.nonNull(proxyUrlString)) { httpClientBuilder.setProxy(HttpSinkUtil.getHttpHostByURL(HttpSinkUtil.getURLByUrlString(proxyUrlString))); LOG.info("sending data via proxy {}",proxyUrlString); } + if(httpSinkConfiguration.getRequestTimout() != null) { + httpClientBuilder.setDefaultRequestConfig(RequestConfig.custom().setConnectionRequestTimeout(Timeout.ofMilliseconds(httpSinkConfiguration.getRequestTimout().toMillis())).build()); + } final HttpAuthOptions.Builder authOptions = new HttpAuthOptions.Builder() .setUrl(httpSinkConfiguration.getUrl()) @@ -316,7 +360,7 @@ private void addCustomHeaders(final ClassicRequestBuilder classicRequestBuilder, */ private ClassicRequestBuilder buildRequestByHTTPMethodType(final HTTPMethodOptions httpMethodOptions) { final ClassicRequestBuilder classicRequestBuilder; - switch(httpMethodOptions){ + switch (httpMethodOptions) { case PUT: classicRequestBuilder = ClassicRequestBuilder.put(); break; @@ -328,5 +372,11 @@ private ClassicRequestBuilder buildRequestByHTTPMethodType(final HTTPMethodOptio return classicRequestBuilder; } + private void accessTokenIfExpired(final Integer tokenExpired,final String url){ + if(oAuthAccessTokenManager.isTokenExpired(tokenExpired)) { + httpAuthOptions.get(url).getClassicHttpRequestBuilder() + .setHeader(BearerTokenAuthHttpSinkHandler.AUTHORIZATION, oAuthAccessTokenManager.getAccessToken(httpSinkConfiguration.getAuthentication().getBearerTokenOptions())); + } + } } \ No newline at end of file diff --git a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/FailedHttpResponseInterceptorTest.java b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/FailedHttpResponseInterceptorTest.java new file mode 100644 index 0000000000..a02e3685a2 --- /dev/null +++ b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/FailedHttpResponseInterceptorTest.java @@ -0,0 +1,34 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.http; + +import org.apache.hc.core5.http.EntityDetails; +import org.apache.hc.core5.http.HttpResponse; +import org.apache.hc.core5.http.protocol.HttpContext; +import org.junit.jupiter.api.Test; +import java.io.IOException; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class FailedHttpResponseInterceptorTest { + + private FailedHttpResponseInterceptor failedHttpResponseInterceptor; + + private HttpResponse httpResponse; + + private EntityDetails entityDetails; + + private HttpContext httpContext; + + @Test + public void test_process(){ + httpResponse = mock(HttpResponse.class); + failedHttpResponseInterceptor = new FailedHttpResponseInterceptor("http://localhost:8080"); + when(httpResponse.getCode()).thenReturn(501); + assertThrows(IOException.class, () -> failedHttpResponseInterceptor.process(httpResponse, entityDetails, httpContext)); + } +} diff --git a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/HttpSinkTest.java b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/HttpSinkTest.java new file mode 100644 index 0000000000..430ae9faf8 --- /dev/null +++ b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/HttpSinkTest.java @@ -0,0 +1,122 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.http; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; +import org.opensearch.dataprepper.model.codec.OutputCodec; +import org.opensearch.dataprepper.model.configuration.PipelineDescription; +import org.opensearch.dataprepper.model.configuration.PluginModel; +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.model.sink.SinkContext; +import org.opensearch.dataprepper.plugins.sink.http.configuration.AuthTypeOptions; +import org.opensearch.dataprepper.plugins.accumulator.BufferTypeOptions; +import org.opensearch.dataprepper.plugins.sink.http.configuration.AwsAuthenticationOptions; +import org.opensearch.dataprepper.plugins.sink.http.configuration.HttpSinkConfiguration; +import org.opensearch.dataprepper.plugins.sink.http.configuration.ThresholdOptions; +import org.opensearch.dataprepper.plugins.sink.http.configuration.HTTPMethodOptions; +import org.opensearch.dataprepper.plugins.sink.http.handler.HttpAuthOptions; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class HttpSinkTest { + + HTTPSink httpSink; + + private PluginSetting pluginSetting; + private PluginFactory pluginFactory; + + private HttpSinkConfiguration httpSinkConfiguration; + + private PipelineDescription pipelineDescription; + + private AwsCredentialsSupplier awsCredentialsSupplier; + + private SinkContext sinkContext; + + private ThresholdOptions thresholdOptions; + + private AwsAuthenticationOptions awsAuthenticationOptions; + + private OutputCodec codec; + + private HttpAuthOptions httpAuthOptions; + + + @BeforeEach + void setUp() { + pluginSetting = mock(PluginSetting.class); + pluginFactory = mock(PluginFactory.class); + httpSinkConfiguration = mock(HttpSinkConfiguration.class); + pipelineDescription = mock(PipelineDescription.class); + awsCredentialsSupplier = mock(AwsCredentialsSupplier.class); + thresholdOptions = mock(ThresholdOptions.class); + sinkContext = mock(SinkContext.class); + codec = mock(OutputCodec.class); + httpAuthOptions = mock(HttpAuthOptions.class); + awsAuthenticationOptions = mock(AwsAuthenticationOptions.class); + when(pluginSetting.getPipelineName()).thenReturn("log-pipeline"); + PluginModel codecConfiguration = new PluginModel("http", new HashMap<>()); + when(httpSinkConfiguration.getCodec()).thenReturn(codecConfiguration); + when(httpSinkConfiguration.getBufferType()).thenReturn(BufferTypeOptions.LOCALFILE); + when(httpAuthOptions.getUrl()).thenReturn("http://localhost:8080"); + when(httpSinkConfiguration.getHttpMethod()).thenReturn(HTTPMethodOptions.POST); + when(httpSinkConfiguration.getAuthType()).thenReturn(AuthTypeOptions.UNAUTHENTICATED); + Map dlqSetting = new HashMap<>(); + dlqSetting.put("bucket", "dlq.test"); + dlqSetting.put("key_path_prefix", "\\dlq"); + PluginModel dlq = new PluginModel("s3",dlqSetting); + when(httpSinkConfiguration.getAwsAuthenticationOptions()).thenReturn(awsAuthenticationOptions); + when(httpSinkConfiguration.getDlqStsRoleARN()).thenReturn("arn:aws:iam::1234567890:role/app-test"); + when(httpSinkConfiguration.getDlqStsRegion()).thenReturn("ap-south-1"); + when(httpSinkConfiguration.getDlq()).thenReturn(dlq); + when(httpSinkConfiguration.getThresholdOptions()).thenReturn(thresholdOptions); + when(thresholdOptions.getEventCount()).thenReturn(10); + when(httpSinkConfiguration.getDlqFile()).thenReturn("\\dlq"); + when(sinkContext.getIncludeKeys()).thenReturn(new ArrayList<>()); + when(sinkContext.getExcludeKeys()).thenReturn(new ArrayList<>()); + } + + private HTTPSink createObjectUnderTest() { + return new HTTPSink(pluginSetting, httpSinkConfiguration, pluginFactory, pipelineDescription, sinkContext, + awsCredentialsSupplier); + } + @Test + void test_http_sink_plugin_isReady_positive() { + httpSink = createObjectUnderTest(); + Assertions.assertNotNull(httpSink); + httpSink.doInitialize(); + assertTrue(httpSink.isReady(), "http sink is initialized and ready to work"); + } + + @Test + void test_http_Sink_plugin_isReady_negative() { + httpSink = createObjectUnderTest(); + Assertions.assertNotNull(httpSink); + assertFalse(httpSink.isReady(), "httpSink sink is not initialized and not ready to work"); + } + + @Test + void test_doOutput_with_empty_records() { + httpSink = createObjectUnderTest(); + Assertions.assertNotNull(httpSink); + httpSink.doInitialize(); + Collection> records = new ArrayList<>(); + httpSink.doOutput(records); + } +} diff --git a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/certificate/CertificateProviderFactoryTest.java b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/certificate/CertificateProviderFactoryTest.java new file mode 100644 index 0000000000..d70e5fe120 --- /dev/null +++ b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/certificate/CertificateProviderFactoryTest.java @@ -0,0 +1,76 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.http.certificate; + +import org.hamcrest.core.IsInstanceOf; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.plugins.certificate.CertificateProvider; +import org.opensearch.dataprepper.plugins.certificate.acm.ACMCertificateProvider; +import org.opensearch.dataprepper.plugins.certificate.file.FileCertificateProvider; +import org.opensearch.dataprepper.plugins.certificate.s3.S3CertificateProvider; +import org.opensearch.dataprepper.plugins.sink.http.configuration.AwsAuthenticationOptions; +import org.opensearch.dataprepper.plugins.sink.http.configuration.HttpSinkConfiguration; +import software.amazon.awssdk.regions.Region; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +class CertificateProviderFactoryTest { + private final String TEST_SSL_CERTIFICATE_FILE = getClass().getClassLoader().getResource("test_cert.crt").getFile(); + private final String TEST_SSL_KEY_FILE = getClass().getClassLoader().getResource("test_decrypted_key.key").getFile(); + + private HttpSinkConfiguration httpSinkConfiguration; + + private AwsAuthenticationOptions awsAuthenticationOptions; + private CertificateProviderFactory certificateProviderFactory; + + @BeforeEach + void setUp() { + httpSinkConfiguration = mock(HttpSinkConfiguration.class); + awsAuthenticationOptions = mock(AwsAuthenticationOptions.class); + } + + @Test + void getCertificateProviderFileCertificateProviderSuccess() { + when(httpSinkConfiguration.isInsecureSkipVerify()).thenReturn(true); + when(httpSinkConfiguration.getSslCertificateFile()).thenReturn(TEST_SSL_CERTIFICATE_FILE); + when(httpSinkConfiguration.getSslKeyFile()).thenReturn(TEST_SSL_KEY_FILE); + + certificateProviderFactory = new CertificateProviderFactory(httpSinkConfiguration); + final CertificateProvider certificateProvider = certificateProviderFactory.getCertificateProvider(); + + assertThat(certificateProvider, IsInstanceOf.instanceOf(FileCertificateProvider.class)); + } + + @Test + void getCertificateProviderS3ProviderSuccess() { + when(httpSinkConfiguration.isSslCertAndKeyFileInS3()).thenReturn(true); + when(awsAuthenticationOptions.getAwsRegion()).thenReturn(Region.of("us-east-1")); + when(httpSinkConfiguration.getAwsAuthenticationOptions()).thenReturn(awsAuthenticationOptions); + when(httpSinkConfiguration.getSslCertificateFile()).thenReturn("s3://data/certificate/test_cert.crt"); + when(httpSinkConfiguration.getSslKeyFile()).thenReturn("s3://data/certificate/test_decrypted_key.key"); + + certificateProviderFactory = new CertificateProviderFactory(httpSinkConfiguration); + final CertificateProvider certificateProvider = certificateProviderFactory.getCertificateProvider(); + + assertThat(certificateProvider, IsInstanceOf.instanceOf(S3CertificateProvider.class)); + } + + @Test + void getCertificateProviderAcmProviderSuccess() { + when(httpSinkConfiguration.useAcmCertForSSL()).thenReturn(true); + when(awsAuthenticationOptions.getAwsRegion()).thenReturn(Region.of("us-east-1")); + when(httpSinkConfiguration.getAwsAuthenticationOptions()).thenReturn(awsAuthenticationOptions); + when(httpSinkConfiguration.getAcmCertificateArn()).thenReturn("arn:aws:acm:us-east-1:account:certificate/1234-567-856456"); + + certificateProviderFactory = new CertificateProviderFactory(httpSinkConfiguration); + final CertificateProvider certificateProvider = certificateProviderFactory.getCertificateProvider(); + + assertThat(certificateProvider, IsInstanceOf.instanceOf(ACMCertificateProvider.class)); + } +} \ No newline at end of file diff --git a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/certificate/HttpClientSSLConnectionManagerTest.java b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/certificate/HttpClientSSLConnectionManagerTest.java new file mode 100644 index 0000000000..96bc679d8a --- /dev/null +++ b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/certificate/HttpClientSSLConnectionManagerTest.java @@ -0,0 +1,50 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.http.certificate; + +import org.apache.hc.client5.http.io.HttpClientConnectionManager; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.plugins.certificate.CertificateProvider; +import org.opensearch.dataprepper.plugins.certificate.file.FileCertificateProvider; +import org.opensearch.dataprepper.plugins.sink.http.configuration.HttpSinkConfiguration; + +import java.io.IOException; + +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class HttpClientSSLConnectionManagerTest { + + private final String TEST_SSL_CERTIFICATE_FILE = getClass().getClassLoader().getResource("test_cert.crt").getFile(); + private final String TEST_SSL_KEY_FILE = getClass().getClassLoader().getResource("test_decrypted_key.key").getFile(); + + HttpClientSSLConnectionManager httpClientSSLConnectionManager; + + private CertificateProviderFactory certificateProviderFactory; + + private HttpSinkConfiguration httpSinkConfiguration; + + @BeforeEach + void setup() throws IOException { + this.httpSinkConfiguration = mock(HttpSinkConfiguration.class); + this.certificateProviderFactory = mock(CertificateProviderFactory.class); + } + + @Test + public void create_httpClientConnectionManager_with_ssl_file_test() { + when(httpSinkConfiguration.getSslCertificateFile()).thenReturn(TEST_SSL_CERTIFICATE_FILE); + when(httpSinkConfiguration.getSslKeyFile()).thenReturn(TEST_SSL_KEY_FILE); + CertificateProvider provider = new FileCertificateProvider(httpSinkConfiguration.getSslCertificateFile(), httpSinkConfiguration.getSslKeyFile()); + when(certificateProviderFactory.getCertificateProvider()).thenReturn(provider); + + CertificateProviderFactory providerFactory = new CertificateProviderFactory(httpSinkConfiguration); + httpClientSSLConnectionManager = new HttpClientSSLConnectionManager(); + HttpClientConnectionManager httpClientConnectionManager = httpClientSSLConnectionManager + .createHttpClientConnectionManager(httpSinkConfiguration, providerFactory); + assertNotNull(httpClientConnectionManager); + } +} diff --git a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/configuration/CustomHeaderOptionsTest.java b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/configuration/CustomHeaderOptionsTest.java deleted file mode 100644 index 1812c0abfe..0000000000 --- a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/configuration/CustomHeaderOptionsTest.java +++ /dev/null @@ -1,43 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.dataprepper.plugins.sink.http.configuration; - -import org.junit.jupiter.api.Test; -import static org.junit.jupiter.api.Assertions.assertNull; - -public class CustomHeaderOptionsTest { - - @Test - public void get_custom_attributes_test() { - assertNull(new CustomHeaderOptions().getCustomAttributes()); - } - - @Test - public void get_target_model_test() { - assertNull(new CustomHeaderOptions().getTargetModel()); - } - - @Test - public void get_target_variant_test() { - assertNull(new CustomHeaderOptions().getTargetVariant()); - } - - @Test - public void get_target_container_hostname_test() { - assertNull(new CustomHeaderOptions().getTargetContainerHostname()); - } - - @Test - public void get_inference_id_test() { - assertNull(new CustomHeaderOptions().getInferenceId()); - } - - @Test - public void get_enable_explanations_test() { - assertNull(new CustomHeaderOptions().getEnableExplanations()); - } - - -} diff --git a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/configuration/HttpSinkConfigurationTest.java b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/configuration/HttpSinkConfigurationTest.java index c487d7dd55..625a3b3e04 100644 --- a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/configuration/HttpSinkConfigurationTest.java +++ b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/configuration/HttpSinkConfigurationTest.java @@ -9,11 +9,15 @@ import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; import com.fasterxml.jackson.dataformat.yaml.YAMLGenerator; import org.hamcrest.CoreMatchers; +import org.hamcrest.MatcherAssert; import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.model.configuration.PluginModel; import org.opensearch.dataprepper.model.types.ByteCount; import org.opensearch.dataprepper.plugins.accumulator.BufferTypeOptions; import software.amazon.awssdk.regions.Region; +import java.util.HashMap; + import java.util.List; import java.util.Map; @@ -21,6 +25,8 @@ import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.instanceOf; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.opensearch.dataprepper.plugins.sink.http.configuration.AuthTypeOptions.HTTP_BASIC; import static org.opensearch.dataprepper.plugins.sink.http.configuration.AuthTypeOptions.UNAUTHENTICATED; @@ -38,10 +44,17 @@ public class HttpSinkConfigurationTest { " username: \"username\"\n" + " password: \"vip\"\n" + " bearer_token:\n" + - " token: \"\"\n" + - " ssl: false\n" + + " client_id: 0oaafr4j79segrYGC5d7\n" + + " client_secret: fFel-3FutCXAOndezEsOVlght6D6DR4OIt7G5D1_oJ6w0wNoaYtgU17JdyXmGf0M\n" + + " token_url: https://localhost/oauth2/default/v1/token\n" + + " grant_type: client_credentials\n" + + " scope: httpSink\n"+ + " insecure_skip_verify: true\n" + " dlq_file: \"/your/local/dlq-file\"\n" + " dlq:\n" + + " s3:\n" + + " bucket: dlq.test\n" + + " key_path_prefix: \\dlq\"\n" + " ssl_certificate_file: \"/full/path/to/certfile.crt\"\n" + " ssl_key_file: \"/full/path/to/keyfile.key\"\n" + " buffer_type: \"in_memory\"\n" + @@ -55,6 +68,7 @@ public class HttpSinkConfigurationTest { " maximum_size: 2mb\n" + " max_retries: 5\n" + " aws_sigv4: true\n" + + " webhook_url: \"http://localhost:8080/webhook\"\n" + " custom_header:\n" + " X-Amzn-SageMaker-Custom-Attributes: [\"test-attribute\"]\n" + " X-Amzn-SageMaker-Target-Model: [\"test-target-model\"]\n" + @@ -68,7 +82,7 @@ public class HttpSinkConfigurationTest { @Test void default_worker_test() { - assertThat(new HttpSinkConfiguration().getWorkers(), CoreMatchers.equalTo(1)); + MatcherAssert.assertThat(new HttpSinkConfiguration().getWorkers(), CoreMatchers.equalTo(1)); } @Test @@ -103,7 +117,7 @@ void get_authentication_test() { @Test void default_ssl_test() { - assertThat(new HttpSinkConfiguration().isSsl(), equalTo(false)); + assertThat(new HttpSinkConfiguration().isInsecureSkipVerify(), equalTo(false)); } @Test @@ -146,6 +160,19 @@ void get_custom_header_options_test() { assertNull(new HttpSinkConfiguration().getCustomHeaderOptions()); } + @Test + void get_http_retry_interval_test() { + assertThat(new HttpSinkConfiguration().getHttpRetryInterval(),equalTo(HttpSinkConfiguration.DEFAULT_HTTP_RETRY_INTERVAL)); + } + @Test + void get_acm_private_key_password_test() {assertNull(new HttpSinkConfiguration().getAcmPrivateKeyPassword());} + + @Test + void get_is_ssl_cert_and_key_file_in_s3_test() {assertThat(new HttpSinkConfiguration().isSslCertAndKeyFileInS3(), equalTo(false));} + + @Test + void get_acm_cert_issue_time_out_millis_test() {assertThat(new HttpSinkConfiguration().getAcmCertIssueTimeOutMillis(), equalTo(new Long(HttpSinkConfiguration.DEFAULT_ACM_CERT_ISSUE_TIME_OUT_MILLIS)));} + @Test void http_sink_pipeline_test_with_provided_config_options() throws JsonProcessingException { final HttpSinkConfiguration httpSinkConfiguration = objectMapper.readValue(SINK_YAML, HttpSinkConfiguration.class); @@ -160,6 +187,7 @@ void http_sink_pipeline_test_with_provided_config_options() throws JsonProcessin assertThat(httpSinkConfiguration.getSslKeyFile(),equalTo("/full/path/to/keyfile.key")); assertThat(httpSinkConfiguration.getWorkers(),equalTo(1)); assertThat(httpSinkConfiguration.getDlqFile(),equalTo("/your/local/dlq-file")); + assertThat(httpSinkConfiguration.getWebhookURL(),equalTo("http://localhost:8080/webhook")); final Map> customHeaderOptions = httpSinkConfiguration.getCustomHeaderOptions(); assertThat(customHeaderOptions.get("X-Amzn-SageMaker-Custom-Attributes"),equalTo(List.of("test-attribute"))); @@ -180,5 +208,55 @@ void http_sink_pipeline_test_with_provided_config_options() throws JsonProcessin final ThresholdOptions thresholdOptions = httpSinkConfiguration.getThresholdOptions(); assertThat(thresholdOptions.getEventCount(),equalTo(2000)); assertThat(thresholdOptions.getMaximumSize(),instanceOf(ByteCount.class)); + + Map pluginSettings = new HashMap<>(); + pluginSettings.put("bucket", "dlq.test"); + pluginSettings.put("key_path_prefix", "dlq"); + final PluginModel pluginModel = new PluginModel("s3", pluginSettings); + assertThat(httpSinkConfiguration.getDlq(), instanceOf(PluginModel.class)); + } + + @Test + public void validate_and_initialize_cert_and_key_file_in_s3_test() throws JsonProcessingException { + final String SINK_YAML = + " url: \"https://httpbin.org/post\"\n" + + " http_method: \"POST\"\n" + + " auth_type: \"http_basic\"\n" + + " authentication:\n" + + " http_basic:\n" + + " username: \"username\"\n" + + " password: \"vip\"\n" + + " insecure_skip_verify: true\n" + + " use_acm_cert_for_ssl: false\n"+ + " acm_certificate_arn: acm_cert\n" + + " ssl_certificate_file: \"/full/path/to/certfile.crt\"\n" + + " ssl_key_file: \"/full/path/to/keyfile.key\"\n" + + " buffer_type: \"in_memory\"\n" + + " threshold:\n" + + " event_count: 2000\n" + + " maximum_size: 2mb\n" + + " max_retries: 5\n"; + final HttpSinkConfiguration httpSinkConfiguration = objectMapper.readValue(SINK_YAML, HttpSinkConfiguration.class); + httpSinkConfiguration.validateAndInitializeCertAndKeyFileInS3(); + } + + @Test + public void is_valid_aws_url_positive_test() throws JsonProcessingException { + + final String SINK_YAML = + " url: \"https://eihycslfo6g2hwrrytyckjkkok.lambda-url.us-east-2.on.aws/\"\n"; + final HttpSinkConfiguration httpSinkConfiguration = objectMapper.readValue(SINK_YAML, HttpSinkConfiguration.class); + + assertTrue(httpSinkConfiguration.isValidAWSUrl()); + } + + @Test + public void is_valid_aws_url_negative_test() throws JsonProcessingException { + + final String SINK_YAML = + " url: \"http://localhost:8080/post\"\n"; + final HttpSinkConfiguration httpSinkConfiguration = objectMapper.readValue(SINK_YAML, HttpSinkConfiguration.class); + + assertFalse(httpSinkConfiguration.isValidAWSUrl()); } } diff --git a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/configuration/ThresholdOptionsTest.java b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/configuration/ThresholdOptionsTest.java index 6bcfdd5a14..a741fa73f9 100644 --- a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/configuration/ThresholdOptionsTest.java +++ b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/configuration/ThresholdOptionsTest.java @@ -4,6 +4,7 @@ */ package org.opensearch.dataprepper.plugins.sink.http.configuration; +import org.hamcrest.MatcherAssert; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.model.types.ByteCount; @@ -16,7 +17,7 @@ class ThresholdOptionsTest { @Test void default_byte_capacity_test() { - assertThat(new ThresholdOptions().getMaximumSize().getBytes(), + MatcherAssert.assertThat(new ThresholdOptions().getMaximumSize().getBytes(), equalTo(ByteCount.parse(DEFAULT_BYTE_CAPACITY).getBytes())); } diff --git a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/configuration/UrlConfigurationOptionTest.java b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/configuration/UrlConfigurationOptionTest.java deleted file mode 100644 index 249216412d..0000000000 --- a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/configuration/UrlConfigurationOptionTest.java +++ /dev/null @@ -1,44 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.dataprepper.plugins.sink.http.configuration; - -import org.junit.jupiter.api.Test; - -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.junit.jupiter.api.Assertions.assertNull; - -public class UrlConfigurationOptionTest { - - @Test - void default_worker_test() { - assertThat(new UrlConfigurationOption().getWorkers(),equalTo(1)); - } - - @Test - void default_codec_test() { - assertNull(new UrlConfigurationOption().getCodec()); - } - - @Test - void default_proxy_test() { - assertNull(new UrlConfigurationOption().getProxy()); - } - - @Test - void default_http_method_test() { - assertNull(new UrlConfigurationOption().getAuthType()); - } - - @Test - void default_auth_type_test() { - assertNull(new UrlConfigurationOption().getAuthType()); - } - - @Test - void default_url_test() { - assertThat(new UrlConfigurationOption().getUrl(), equalTo(null)); - } -} diff --git a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/dlq/DlqPushHandlerTest.java b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/dlq/DlqPushHandlerTest.java new file mode 100644 index 0000000000..d289b78ac4 --- /dev/null +++ b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/dlq/DlqPushHandlerTest.java @@ -0,0 +1,96 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.http.dlq; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; +import org.opensearch.dataprepper.model.configuration.PluginModel; +import org.opensearch.dataprepper.model.configuration.PluginSetting; +import org.opensearch.dataprepper.model.plugin.PluginFactory; +import org.opensearch.dataprepper.plugins.dlq.DlqProvider; +import org.opensearch.dataprepper.plugins.dlq.DlqWriter; +import org.opensearch.dataprepper.plugins.sink.http.configuration.AwsAuthenticationOptions; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.mockito.Mockito.any; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.anyList; +import static org.mockito.Mockito.anyString; + +public class DlqPushHandlerTest { + + private static final String BUCKET = "bucket"; + private static final String BUCKET_VALUE = "test"; + private static final String ROLE = "arn:aws:iam::524239988944:role/app-test"; + + private static final String REGION = "ap-south-1"; + private static final String S3_PLUGIN_NAME = "s3"; + private static final String KEY_PATH_PREFIX = "key_path_prefix"; + + private static final String KEY_PATH_PREFIX_VALUE = "dlq/"; + + private static final String PIPELINE_NAME = "log-pipeline"; + + private static final String DLQ_FILE = "local_dlq_file"; + + private PluginModel pluginModel; + + private DlqPushHandler dlqPushHandler; + private PluginFactory pluginFactory; + + private AwsAuthenticationOptions awsAuthenticationOptions; + + private DlqProvider dlqProvider; + + private DlqWriter dlqWriter; + + @BeforeEach + public void setUp() throws Exception{ + this.pluginFactory = mock(PluginFactory.class); + this.pluginModel = mock(PluginModel.class); + this.awsAuthenticationOptions = mock(AwsAuthenticationOptions.class); + this.dlqProvider = mock(DlqProvider.class); + this.dlqWriter = mock(DlqWriter.class); + } + + @Test + public void perform_for_dlq_s3_success() throws IOException { + Map props = new HashMap<>(); + props.put(BUCKET,BUCKET_VALUE); + props.put(KEY_PATH_PREFIX,KEY_PATH_PREFIX_VALUE); + + when(pluginFactory.loadPlugin(any(Class.class), any(PluginSetting.class))).thenReturn(dlqProvider); + + when(dlqProvider.getDlqWriter(Mockito.anyString())).thenReturn(Optional.of(dlqWriter)); + doNothing().when(dlqWriter).write(anyList(), anyString(), anyString()); + FailedDlqData failedDlqData = FailedDlqData.builder().build(); + dlqPushHandler = new DlqPushHandler(null,pluginFactory, BUCKET_VALUE, ROLE, REGION,KEY_PATH_PREFIX_VALUE); + + PluginSetting pluginSetting = new PluginSetting(S3_PLUGIN_NAME, props); + pluginSetting.setPipelineName(PIPELINE_NAME); + dlqPushHandler.perform(pluginSetting, failedDlqData); + verify(dlqWriter).write(anyList(), anyString(), anyString()); + } + + + @Test + public void perform_for_dlq_local_file_success(){ + + FailedDlqData failedDlqData = FailedDlqData.builder().build(); + dlqPushHandler = new DlqPushHandler(DLQ_FILE,pluginFactory,null, ROLE, REGION,null); + + PluginSetting pluginSetting = new PluginSetting(S3_PLUGIN_NAME, null); + pluginSetting.setPipelineName(PIPELINE_NAME); + dlqPushHandler.perform(pluginSetting, failedDlqData); + } +} diff --git a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/handler/BasicAuthHttpSinkHandlerTest.java b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/handler/BasicAuthHttpSinkHandlerTest.java new file mode 100644 index 0000000000..cfe6eb06a0 --- /dev/null +++ b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/handler/BasicAuthHttpSinkHandlerTest.java @@ -0,0 +1,52 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.http.handler; + +import org.apache.hc.client5.http.impl.classic.HttpClients; +import org.apache.hc.client5.http.impl.io.PoolingHttpClientConnectionManager; +import org.apache.hc.client5.http.impl.io.PoolingHttpClientConnectionManagerBuilder; +import org.apache.hc.core5.http.HttpHost; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; +import org.opensearch.dataprepper.plugins.sink.http.util.HttpSinkUtil; + +import java.net.URL; + +import static org.mockito.ArgumentMatchers.any; + +public class BasicAuthHttpSinkHandlerTest { + + private MockedStatic httpSinkUtilStatic; + + private String urlString = "http://localhost:8080"; + @BeforeEach + public void setUp() throws Exception{ + URL url = new URL(urlString); + httpSinkUtilStatic = Mockito.mockStatic(HttpSinkUtil.class); + httpSinkUtilStatic.when(() -> HttpSinkUtil.getURLByUrlString(any())) + .thenReturn(url); + HttpHost targetHost = new HttpHost(url.toURI().getScheme(), url.getHost(), url.getPort()); + httpSinkUtilStatic.when(() -> HttpSinkUtil.getHttpHostByURL(any(URL.class))) + .thenReturn(targetHost); + } + + @AfterEach + public void tearDown() { + httpSinkUtilStatic.close(); + } + + @Test + public void authenticateTest() { + HttpAuthOptions.Builder httpAuthOptionsBuilder = new HttpAuthOptions.Builder(); + httpAuthOptionsBuilder.setUrl(urlString); + httpAuthOptionsBuilder.setHttpClientBuilder(HttpClients.custom()); + httpAuthOptionsBuilder.setHttpClientConnectionManager(PoolingHttpClientConnectionManagerBuilder.create().build()); + Assertions.assertEquals(urlString, new BasicAuthHttpSinkHandler("test", "test", new PoolingHttpClientConnectionManager()).authenticate(httpAuthOptionsBuilder).getUrl()); + } +} diff --git a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/handler/BearerTokenAuthHttpSinkHandlerTest.java b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/handler/BearerTokenAuthHttpSinkHandlerTest.java new file mode 100644 index 0000000000..bf4406ea41 --- /dev/null +++ b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/handler/BearerTokenAuthHttpSinkHandlerTest.java @@ -0,0 +1,44 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.http.handler; + +import org.apache.hc.client5.http.impl.classic.HttpClients; +import org.apache.hc.client5.http.impl.io.PoolingHttpClientConnectionManager; +import org.apache.hc.client5.http.impl.io.PoolingHttpClientConnectionManagerBuilder; +import org.apache.hc.core5.http.io.support.ClassicRequestBuilder; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.plugins.sink.http.OAuthAccessTokenManager; +import org.opensearch.dataprepper.plugins.sink.http.configuration.BearerTokenOptions; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class BearerTokenAuthHttpSinkHandlerTest { + + private String urlString = "http://localhost:8080"; + + private OAuthAccessTokenManager oAuthAccessTokenManager; + + private BearerTokenOptions bearerTokenOptions; + @BeforeEach + public void setUp() throws Exception{ + bearerTokenOptions = new BearerTokenOptions(); + oAuthAccessTokenManager = mock(OAuthAccessTokenManager.class); + when(oAuthAccessTokenManager.getAccessToken(bearerTokenOptions)).thenReturn("access_token_test"); + } + + @Test + public void authenticateTest() { + + HttpAuthOptions.Builder httpAuthOptionsBuilder = new HttpAuthOptions.Builder(); + httpAuthOptionsBuilder.setUrl(urlString); + httpAuthOptionsBuilder.setHttpClientBuilder(HttpClients.custom()); + httpAuthOptionsBuilder.setHttpClientConnectionManager(PoolingHttpClientConnectionManagerBuilder.create().build()); + httpAuthOptionsBuilder.setClassicHttpRequestBuilder(ClassicRequestBuilder.post()); + Assertions.assertEquals(urlString, new BearerTokenAuthHttpSinkHandler(bearerTokenOptions, new PoolingHttpClientConnectionManager(),oAuthAccessTokenManager).authenticate(httpAuthOptionsBuilder).getUrl()); + } +} diff --git a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkAwsServiceTest.java b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkAwsServiceTest.java new file mode 100644 index 0000000000..4d41c463dc --- /dev/null +++ b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkAwsServiceTest.java @@ -0,0 +1,56 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.http.service; + +import org.apache.hc.client5.http.impl.classic.HttpClientBuilder; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; +import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; +import org.opensearch.dataprepper.plugins.sink.http.configuration.AwsAuthenticationOptions; +import org.opensearch.dataprepper.plugins.sink.http.configuration.HttpSinkConfiguration; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.regions.Region; + +import java.io.IOException; +import java.util.HashMap; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class HttpSinkAwsServiceTest { + + private HttpSinkConfiguration httpSinkConfiguration; + + private HttpClientBuilder httpClientBuilder; + + private AwsCredentialsSupplier awsCredentialsSupplier; + + private AwsAuthenticationOptions awsAuthenticationOptions; + + private AwsCredentialsProvider awsCredentialsProvider; + + @BeforeEach + public void setup() throws IOException { + httpSinkConfiguration = mock(HttpSinkConfiguration.class); + httpClientBuilder = mock(HttpClientBuilder.class); + awsCredentialsSupplier = mock(AwsCredentialsSupplier.class); + awsAuthenticationOptions = mock(AwsAuthenticationOptions.class); + awsCredentialsProvider = mock(AwsCredentialsProvider.class); + when(awsAuthenticationOptions.getAwsStsRoleArn()).thenReturn("arn:aws:iam::1234567890:role/app-test"); + when(awsAuthenticationOptions.getAwsStsExternalId()).thenReturn("test"); + when(awsAuthenticationOptions.getAwsRegion()).thenReturn(Region.of("ap-south-1")); + when(awsAuthenticationOptions.getServiceName()).thenReturn("lambda"); + when(awsAuthenticationOptions.getAwsStsHeaderOverrides()).thenReturn(new HashMap<>()); + when(httpSinkConfiguration.getAwsAuthenticationOptions()).thenReturn(awsAuthenticationOptions); + when(awsCredentialsSupplier.getProvider(Mockito.any())).thenReturn(awsCredentialsProvider); + + } + + @Test + public void attachSigV4Test() { + HttpSinkAwsService.attachSigV4(httpSinkConfiguration,httpClientBuilder,awsCredentialsSupplier); + } +} diff --git a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkServiceTest.java b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkServiceTest.java index eae25204ee..1a37476974 100644 --- a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkServiceTest.java +++ b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkServiceTest.java @@ -17,10 +17,10 @@ import org.apache.hc.core5.http.ClassicHttpRequest; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.Mockito; import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.codec.OutputCodec; -import org.opensearch.dataprepper.model.configuration.PluginModel; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.event.EventHandle; @@ -29,6 +29,8 @@ import org.opensearch.dataprepper.plugins.accumulator.BufferFactory; import org.opensearch.dataprepper.plugins.accumulator.InMemoryBufferFactory; import org.opensearch.dataprepper.plugins.sink.http.FailedHttpResponseInterceptor; +import org.opensearch.dataprepper.plugins.sink.http.configuration.AuthenticationOptions; +import org.opensearch.dataprepper.plugins.sink.http.configuration.AuthTypeOptions; import org.opensearch.dataprepper.plugins.sink.http.configuration.HttpSinkConfiguration; import org.opensearch.dataprepper.plugins.sink.http.configuration.ThresholdOptions; import org.opensearch.dataprepper.plugins.sink.http.dlq.DlqPushHandler; @@ -50,6 +52,7 @@ import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; + public class HttpSinkServiceTest { private ObjectMapper objectMapper = new ObjectMapper(new YAMLFactory().enable(YAMLGenerator.Feature.USE_PLATFORM_LINE_BREAKS)); @@ -66,8 +69,12 @@ public class HttpSinkServiceTest { " username: \"username\"\n" + " password: \"vip\"\n" + " bearer_token:\n" + - " token: \"test\"\n" + - " ssl: false\n" + + " client_id: 0oaafr4j79segrYGC5d7\n" + + " client_secret: fFel-3FutCXAOndezEsOVlght6D6DR4OIt7G5D1_oJ6w0wNoaYtgU17JdyXmGf0M\n" + + " token_url: https://localhost/oauth2/default/v1/token\n" + + " grant_type: client_credentials\n" + + " scope: httpSink\n"+ + " insecure_skip_verify: true\n" + " dlq_file: \"/your/local/dlq-file\"\n" + " dlq:\n" + " ssl_certificate_file: \"/full/path/to/certfile.crt\"\n" + @@ -110,10 +117,8 @@ public class HttpSinkServiceTest { private CloseableHttpResponse closeableHttpResponse; - private String tagsTargetKey; - @BeforeEach - void setup() throws IOException { + void setup() throws Exception { this.codec = mock(OutputCodec.class); this.pluginMetrics = mock(PluginMetrics.class); this.httpSinkConfiguration = objectMapper.readValue(SINK_YAML,HttpSinkConfiguration.class); @@ -128,13 +133,12 @@ void setup() throws IOException { this.closeableHttpResponse = mock(CloseableHttpResponse.class); this.bufferFactory = new InMemoryBufferFactory(); - lenient().when(httpClientBuilder.setConnectionManager(null)).thenReturn(httpClientBuilder); + lenient().when(httpClientBuilder.setConnectionManager(Mockito.any())).thenReturn(httpClientBuilder); lenient().when(httpClientBuilder.addResponseInterceptorLast(any(FailedHttpResponseInterceptor.class))).thenReturn(httpClientBuilder); lenient().when(httpClientBuilder.build()).thenReturn(closeableHttpClient); lenient().when(closeableHttpClient.execute(any(ClassicHttpRequest.class),any(HttpClientContext.class))).thenReturn(closeableHttpResponse); when(pluginMetrics.counter(HttpSinkService.HTTP_SINK_RECORDS_SUCCESS_COUNTER)).thenReturn(httpSinkRecordsSuccessCounter); when(pluginMetrics.counter(HttpSinkService.HTTP_SINK_RECORDS_FAILED_COUNTER)).thenReturn(httpSinkRecordsFailedCounter); - } HttpSinkService createObjectUnderTest(final int eventCount,final HttpSinkConfiguration httpSinkConfig) throws NoSuchFieldException, IllegalAccessException { @@ -148,7 +152,9 @@ HttpSinkService createObjectUnderTest(final int eventCount,final HttpSinkConfigu webhookService, httpClientBuilder, pluginMetrics, - pluginSetting); + pluginSetting, + codec, + null); } @Test @@ -181,10 +187,16 @@ void http_sink_service_test_with_internal_server_error() throws NoSuchFieldExcep } @Test - void http_sink_service_test_with_single_record_with_basic_authentication() throws NoSuchFieldException, IllegalAccessException { + void http_sink_service_test_with_single_record_with_basic_authentication() throws NoSuchFieldException, IllegalAccessException, JsonProcessingException { + + final String basicAuthYaml = " http_basic:\n" + + " username: \"username\"\n" + + " password: \"vip\"\n" ; + ReflectivelySetField.setField(HttpSinkConfiguration.class,httpSinkConfiguration,"authentication", objectMapper.readValue(basicAuthYaml, AuthenticationOptions.class)); + ReflectivelySetField.setField(HttpSinkConfiguration.class,httpSinkConfiguration,"authType", AuthTypeOptions.HTTP_BASIC); + final Record eventRecord = new Record<>(JacksonEvent.fromMessage("{\"message\":\"c3f847eb-333a-49c3-a4cd-54715ad1b58a\"}")); lenient().when(httpClientBuilder.setDefaultCredentialsProvider(any(BasicCredentialsProvider.class))).thenReturn(httpClientBuilder); final HttpSinkService objectUnderTest = createObjectUnderTest(1,httpSinkConfiguration); - final Record eventRecord = new Record<>(JacksonEvent.fromMessage("{\"message\":\"c3f847eb-333a-49c3-a4cd-54715ad1b58a\"}")); objectUnderTest.output(List.of(eventRecord)); verify(httpSinkRecordsSuccessCounter).increment(1); } @@ -193,8 +205,12 @@ void http_sink_service_test_with_single_record_with_basic_authentication() throw void http_sink_service_test_with_single_record_with_bearer_token() throws NoSuchFieldException, IllegalAccessException, JsonProcessingException { lenient().when(httpClientBuilder.setDefaultCredentialsProvider(any(BasicCredentialsProvider.class))).thenReturn(httpClientBuilder); final String authentication = " bearer_token:\n" + - " token: \"test\"" ; - ReflectivelySetField.setField(HttpSinkConfiguration.class,httpSinkConfiguration,"authentication", objectMapper.readValue(authentication, PluginModel.class)); + " client_id: 0oaafr4j79segrYGC5d7\n" + + " client_secret: fFel-3FutCXAOndezEsOVlght6D6DR4OIt7G5D1_oJ6w0wNoaYtgU17JdyXmGf0M\n" + + " token_url: https://localhost/oauth2/default/v1/token\n" + + " grant_type: client_credentials\n" + + " scope: httpSink" ; + ReflectivelySetField.setField(HttpSinkConfiguration.class,httpSinkConfiguration,"authentication", objectMapper.readValue(authentication, AuthenticationOptions.class)); final HttpSinkService objectUnderTest = createObjectUnderTest(1,httpSinkConfiguration); final Record eventRecord = new Record<>(JacksonEvent.fromMessage("{\"message\":\"c3f847eb-333a-49c3-a4cd-54715ad1b58a\"}")); objectUnderTest.output(List.of(eventRecord)); @@ -216,6 +232,7 @@ void http_sink_service_test_output_with_single_record_ack_release() throws NoSuc final Event event = mock(Event.class); given(event.toJsonString()).willReturn("{\"message\":\"c3f847eb-333a-49c3-a4cd-54715ad1b58a\"}"); given(event.getEventHandle()).willReturn(mock(EventHandle.class)); + given(event.jsonBuilder()).willReturn(mock(Event.JsonStringBuilder.class)); objectUnderTest.output(List.of(new Record<>(event))); verify(httpSinkRecordsSuccessCounter).increment(1); } diff --git a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/service/WebhookServiceTest.java b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/service/WebhookServiceTest.java index 25aa1c53c5..e42504d7fe 100644 --- a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/service/WebhookServiceTest.java +++ b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/service/WebhookServiceTest.java @@ -17,6 +17,7 @@ import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.plugins.sink.http.FailedHttpResponseInterceptor; import org.opensearch.dataprepper.plugins.sink.http.HttpEndPointResponse; + import org.opensearch.dataprepper.plugins.sink.http.configuration.HttpSinkConfiguration; import org.opensearch.dataprepper.plugins.sink.http.dlq.FailedDlqData; @@ -71,7 +72,11 @@ private WebhookService createObjectUnderTest(){ public void http_sink_webhook_service_test_with_one_webhook_success_push() throws IOException { lenient().when(closeableHttpClient.execute(any(ClassicHttpRequest.class),any(HttpClientContext.class))).thenReturn(closeableHttpResponse); HttpEndPointResponse httpEndPointResponse = new HttpEndPointResponse(TEST_URL,200); - FailedDlqData failedDlqData = FailedDlqData.builder().withBufferData("Test Data").withEndPointResponses(httpEndPointResponse).build(); + FailedDlqData failedDlqData = + FailedDlqData.builder() + .withUrl(httpEndPointResponse.getUrl()) + .withMessage("Test Data") + .withStatus(httpEndPointResponse.getStatusCode()).build(); WebhookService webhookService = createObjectUnderTest(); webhookService.pushWebhook(failedDlqData); verify(httpSinkWebhookSuccessCounter).increment(); @@ -81,7 +86,11 @@ public void http_sink_webhook_service_test_with_one_webhook_success_push() throw public void http_sink_webhook_service_test_with_one_webhook_failed_to_push() throws IOException { when(closeableHttpClient.execute(any(HttpHost.class),any(ClassicHttpRequest.class),any(HttpClientContext.class))).thenThrow(new IOException("Internal Server Error")); HttpEndPointResponse httpEndPointResponse = new HttpEndPointResponse(TEST_URL,500); - FailedDlqData failedDlqData = FailedDlqData.builder().withBufferData("Test Data").withEndPointResponses(httpEndPointResponse).build(); + FailedDlqData failedDlqData = FailedDlqData.builder() + .withMessage("Test Data") + .withStatus(httpEndPointResponse.getStatusCode()) + .withUrl(httpEndPointResponse.getUrl()) + .build(); WebhookService webhookService = createObjectUnderTest(); webhookService.pushWebhook(failedDlqData); verify(httpSinkWebhookFailedCounter).increment(); diff --git a/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/util/HttpSinkUtilTest.java b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/util/HttpSinkUtilTest.java new file mode 100644 index 0000000000..61fda4fb0a --- /dev/null +++ b/data-prepper-plugins/http-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/http/util/HttpSinkUtilTest.java @@ -0,0 +1,34 @@ +package org.opensearch.dataprepper.plugins.sink.http.util; + +import org.apache.hc.core5.http.HttpHost; +import org.junit.jupiter.api.Test; + +import java.net.MalformedURLException; +import java.net.URL; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +public class HttpSinkUtilTest { + + @Test + public void get_url_by_url_string_positive_test() throws MalformedURLException { + assertEquals(HttpSinkUtil.getURLByUrlString("http://localhost:8080"), new URL("http://localhost:8080")); + } + + @Test + public void get_http_host_by_url_positive_test() throws MalformedURLException { + assertEquals(HttpSinkUtil.getHttpHostByURL(new URL("http://localhost:8080")), new HttpHost(null, "localhost", 8080)); + } + + @Test + public void get_url_by_url_string_negative_test() { + assertThrows(RuntimeException.class, () -> HttpSinkUtil.getURLByUrlString("ht://localhost:8080")); + } + + @Test + public void get_http_host_by_url_negative_test() { + assertThrows(RuntimeException.class, () -> HttpSinkUtil.getHttpHostByURL(new URL("http://localhost:8080/h?s=^IXIC"))); + } + +} diff --git a/data-prepper-plugins/http-sink/src/test/resources/test_cert.crt b/data-prepper-plugins/http-sink/src/test/resources/test_cert.crt new file mode 100644 index 0000000000..26c78d1411 --- /dev/null +++ b/data-prepper-plugins/http-sink/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/http-sink/src/test/resources/test_decrypted_key.key b/data-prepper-plugins/http-sink/src/test/resources/test_decrypted_key.key new file mode 100644 index 0000000000..479b877131 --- /dev/null +++ b/data-prepper-plugins/http-sink/src/test/resources/test_decrypted_key.key @@ -0,0 +1,15 @@ +-----BEGIN RSA PRIVATE KEY----- +MIICXAIBAAKBgQCq292IXSm0OT7Sx2pZddIi2QvWXQ1a6hzr6lUjPDR5S/H/NoCd +5dJ/ZO4WHOXiorLrWkvXg05ifg55/MZ/6j0ikyKvVofjmTvnvWBBKSPxuBfSqRlC +4DMkYPbCStBpWXRLfaORiTSSCVJKbSVOXoi7hNCFfo8iaDb5vDhOhCT6sQIDAQAB +AoGANrrhFqpJDpr7vcb1ER0Fp/YArbT27zVo+EUC6puBb41dQlQyFOImcHpjLaAq +H1PgnjU5cBp2hGQ+vOK0rwrYc/HNl6vfh6N3NbDptMiuoBafRJA9JzYourAM09BU +zmXyr61Yn3KHzx1PRwWe37icX93oXP3P0qHb3dI1ZF4jG0ECQQDU5N/a7ogoz2zn +ZssD6FvUOUQDsdBWdXmhUvg+YdZrV44e4xk+FVzwEONoRktEYKz9MFXlsgNHr445 +KRguHWcJAkEAzXQkwOkN8WID1wrwoobUIMbZSGAZzofwkKXgTTnllnT1qOQXuRbS +aCMejFEymBBef4aXP6N4+va2FKW/MF34aQJAO2oMl1sOoOUSrZngepy0VAwPUUCk +thxe74jqQu6nGpn6zd/vQYZQw6bS8Fz90H1yic6dilcd1znFZWp0lxoZkQJBALeI +xoBycRsuFQIYasi1q3AwUtBd0Q/3zkZZeBtk2hzjFMUwJaUZpxKSNOrialD/ZnuD +jz+xWBTRKe0d98JMX+kCQCmsJEj/HYQAC1GamZ7JQWogRSRF2KTgTWRaDXDxy0d4 +yUQgwHB+HZLFcbi1JEK6eIixCsX8iifrrkteh+1npJ0= +-----END RSA PRIVATE KEY----- From 8b674cd1bd26447685e0739454ce134c5d2dae50 Mon Sep 17 00:00:00 2001 From: Taylor Gray Date: Wed, 25 Oct 2023 09:52:13 -0500 Subject: [PATCH 09/53] Add obfuscate_when parameter and tags_on_match failure to obfuscate processor (#3544) Add obfuscate_when parameter to obfuscate processor Signed-off-by: Taylor Gray --- .../obfuscation/ObfuscationProcessor.java | 24 +++++- .../ObfuscationProcessorConfig.java | 25 ++++++- .../obfuscation/ObfuscationProcessorTest.java | 75 +++++++++++++++---- 3 files changed, 107 insertions(+), 17 deletions(-) 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 f349ce0ebc..0ffc309797 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 @@ -5,6 +5,7 @@ package org.opensearch.dataprepper.plugins.processor.obfuscation; +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; @@ -33,6 +34,10 @@ public class ObfuscationProcessor extends AbstractProcessor, Recor private static final String COMMON_PATTERN_REGEX = "^%\\{([A-Z_0-9]+)}$"; private static final Logger LOG = LoggerFactory.getLogger(ObfuscationProcessor.class); + + private final ExpressionEvaluator expressionEvaluator; + private final ObfuscationProcessorConfig obfuscationProcessorConfig; + private final String source; private final String target; @@ -41,13 +46,20 @@ public class ObfuscationProcessor extends AbstractProcessor, Recor @DataPrepperPluginConstructor - public ObfuscationProcessor(final PluginMetrics pluginMetrics, final ObfuscationProcessorConfig config, final PluginFactory pluginFactory) { + public ObfuscationProcessor(final PluginMetrics pluginMetrics, + final ObfuscationProcessorConfig config, + final PluginFactory pluginFactory, + final ExpressionEvaluator expressionEvaluator) { // No special metrics generate by this processor. super(pluginMetrics); this.source = config.getSource(); this.target = config.getTarget(); this.patterns = new ArrayList<>(); + this.expressionEvaluator = expressionEvaluator; + this.obfuscationProcessorConfig = config; + + config.validateObfuscateWhen(expressionEvaluator); final PluginModel actionPlugin = config.getAction(); if (actionPlugin == null) { @@ -94,14 +106,24 @@ 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; + } + if (!recordEvent.containsKey(source)) { continue; } String rawValue = recordEvent.get(source, String.class); + // 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); diff --git a/data-prepper-plugins/obfuscate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/obfuscation/ObfuscationProcessorConfig.java b/data-prepper-plugins/obfuscate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/obfuscation/ObfuscationProcessorConfig.java index 714f41cb85..56defb6baf 100644 --- a/data-prepper-plugins/obfuscate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/obfuscation/ObfuscationProcessorConfig.java +++ b/data-prepper-plugins/obfuscate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/obfuscation/ObfuscationProcessorConfig.java @@ -8,7 +8,9 @@ import com.fasterxml.jackson.annotation.JsonProperty; import jakarta.validation.constraints.NotEmpty; import jakarta.validation.constraints.NotNull; +import org.opensearch.dataprepper.expression.ExpressionEvaluator; import org.opensearch.dataprepper.model.configuration.PluginModel; +import org.opensearch.dataprepper.model.plugin.InvalidPluginConfigurationException; import java.util.List; @@ -28,14 +30,21 @@ public class ObfuscationProcessorConfig { @JsonProperty("action") private PluginModel action; + @JsonProperty("obfuscate_when") + private String obfuscateWhen; + + @JsonProperty("tags_on_match_failure") + private List tagsOnMatchFailure; + public ObfuscationProcessorConfig() { } - public ObfuscationProcessorConfig(String source, List patterns, String target, PluginModel action) { + public ObfuscationProcessorConfig(String source, List patterns, String target, PluginModel action, List tagsOnMatchFailure) { this.source = source; this.patterns = patterns; this.target = target; this.action = action; + this.tagsOnMatchFailure = tagsOnMatchFailure; } public String getSource() { @@ -53,4 +62,18 @@ public String getTarget() { public PluginModel getAction() { return action; } + + public String getObfuscateWhen() { + return obfuscateWhen; + } + + public List getTagsOnMatchFailure() { + return tagsOnMatchFailure; + } + + void validateObfuscateWhen(final ExpressionEvaluator expressionEvaluator) { + if (obfuscateWhen != null && !expressionEvaluator.isValidExpressionStatement(obfuscateWhen)) { + throw new InvalidPluginConfigurationException(String.format("obfuscate_when value %s is not a valid Data Prepper expression statement", obfuscateWhen)); + } + } } diff --git a/data-prepper-plugins/obfuscate-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/obfuscation/ObfuscationProcessorTest.java b/data-prepper-plugins/obfuscate-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/obfuscation/ObfuscationProcessorTest.java index 50e232d75d..b29ad3b0f4 100644 --- a/data-prepper-plugins/obfuscate-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/obfuscation/ObfuscationProcessorTest.java +++ b/data-prepper-plugins/obfuscate-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/obfuscation/ObfuscationProcessorTest.java @@ -13,6 +13,7 @@ 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.configuration.PluginModel; import org.opensearch.dataprepper.model.configuration.PluginSetting; @@ -27,9 +28,11 @@ import java.util.HashMap; 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.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; @@ -53,6 +56,9 @@ class ObfuscationProcessorTest { @Mock private ObfuscationProcessorConfig mockConfig; + @Mock + private ExpressionEvaluator expressionEvaluator; + private ObfuscationProcessor obfuscationProcessor; static Record buildRecordWithEvent(final Map data) { @@ -70,12 +76,51 @@ private Record createRecord(String message) { @BeforeEach void setup() { - final ObfuscationProcessorConfig defaultConfig = new ObfuscationProcessorConfig("message", null, null, null); + final ObfuscationProcessorConfig defaultConfig = new ObfuscationProcessorConfig("message", null, null, null, null); lenient().when(mockConfig.getSource()).thenReturn(defaultConfig.getSource()); lenient().when(mockConfig.getAction()).thenReturn(defaultConfig.getAction()); lenient().when(mockConfig.getPatterns()).thenReturn(defaultConfig.getPatterns()); lenient().when(mockConfig.getTarget()).thenReturn(defaultConfig.getTarget()); - obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory); + lenient().when(mockConfig.getObfuscateWhen()).thenReturn(null); + lenient().when(mockConfig.getTagsOnMatchFailure()).thenReturn(List.of(UUID.randomUUID().toString())); + obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory, expressionEvaluator); + } + + @Test + void obfuscate_when_evaluates_to_false_does_not_modify_event() { + final String expression = "/test == success"; + final Record record = createRecord(UUID.randomUUID().toString()); + when(mockConfig.getObfuscateWhen()).thenReturn(expression); + when(expressionEvaluator.evaluateConditional(expression, record.getData())).thenReturn(false); + + final ObfuscationProcessor objectUnderTest = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory, expressionEvaluator); + + final Map expectedEventMap = record.getData().toMap(); + final List> editedRecords = (List>) objectUnderTest.doExecute(Collections.singletonList(record)); + + assertThat(editedRecords.size(), equalTo(1)); + assertThat(editedRecords.get(0).getData().toMap(), equalTo(expectedEventMap)); + } + + @Test + void event_is_tagged_with_match_failure_tags_when_it_does_not_match_any_patterns_and_when_condition_is_true() { + final Record record = createRecord(UUID.randomUUID().toString()); + + final String expression = UUID.randomUUID().toString(); + when(mockConfig.getObfuscateWhen()).thenReturn(expression); + when(expressionEvaluator.evaluateConditional(expression, record.getData())).thenReturn(true); + when(mockConfig.getPatterns()).thenReturn(List.of(UUID.randomUUID().toString())); + + final ObfuscationProcessor objectUnderTest = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory, expressionEvaluator); + + final Map expectedEventMap = record.getData().toMap(); + final List> editedRecords = (List>) objectUnderTest.doExecute(Collections.singletonList(record)); + + assertThat(editedRecords.size(), equalTo(1)); + assertThat(editedRecords.get(0).getData().toMap(), equalTo(expectedEventMap)); + assertThat(editedRecords.get(0).getData().getMetadata().getTags(), notNullValue()); + assertThat(editedRecords.get(0).getData().getMetadata().getTags().size(), equalTo(1)); + assertThat(editedRecords.get(0).getData().getMetadata().getTags().contains(mockConfig.getTagsOnMatchFailure().get(0)), equalTo(true)); } @@ -102,7 +147,7 @@ void testProcessorWithDifferentAction() { when(mockFactory.loadPlugin(eq(ObfuscationAction.class), any(PluginSetting.class))) .thenReturn(mockAction); - obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory); + obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory, expressionEvaluator); final Record record = createRecord("Hello"); final List> editedRecords = (List>) obfuscationProcessor.doExecute(Collections.singletonList(record)); @@ -117,7 +162,7 @@ void testProcessorWithDifferentAction() { @ValueSource(strings = {"hello", "hello, world", "This is a message", "123", "你好"}) void testProcessorWithTarget(String message) { when(mockConfig.getTarget()).thenReturn("new_message"); - obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory); + obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory, expressionEvaluator); final Record record = createRecord(message); final List> editedRecords = (List>) obfuscationProcessor.doExecute(Collections.singletonList(record)); @@ -135,7 +180,7 @@ void testProcessorWithTarget(String message) { @Test void testProcessorWithUnknownSource() { when(mockConfig.getSource()).thenReturn("email"); - obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory); + obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory, expressionEvaluator); final Record record = createRecord("Hello"); final List> editedRecords = (List>) obfuscationProcessor.doExecute(Collections.singletonList(record)); @@ -158,7 +203,7 @@ void testProcessorWithUnknownSource() { }) void testProcessorWithPattern(String message, String pattern, String expected) { when(mockConfig.getPatterns()).thenReturn(List.of(pattern)); - obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory); + obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory, expressionEvaluator); final Record record = createRecord(message); final List> editedRecords = (List>) obfuscationProcessor.doExecute(Collections.singletonList(record)); @@ -171,13 +216,13 @@ void testProcessorWithPattern(String message, String pattern, String expected) { @Test void testProcessorWithUnknownPattern() { when(mockConfig.getPatterns()).thenReturn(List.of("%{UNKNOWN}")); - assertThrows(InvalidPluginConfigurationException.class, () -> new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory)); + assertThrows(InvalidPluginConfigurationException.class, () -> new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory, expressionEvaluator)); } @Test void testProcessorInvalidPattern() { when(mockConfig.getPatterns()).thenReturn(List.of("[")); - assertThrows(InvalidPluginConfigurationException.class, () -> new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory)); + assertThrows(InvalidPluginConfigurationException.class, () -> new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory, expressionEvaluator)); } @ParameterizedTest @@ -194,7 +239,7 @@ void testProcessorInvalidPattern() { }) void testProcessorWithEmailAddressPattern(String message, String expected) { when(mockConfig.getPatterns()).thenReturn(List.of("%{EMAIL_ADDRESS}")); - obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory); + obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory, expressionEvaluator); final Record record = createRecord(message); @@ -221,7 +266,7 @@ void testProcessorWithEmailAddressPattern(String message, String expected) { }) void testProcessorWithUSPhoneNumberPattern(String message, String expected) { when(mockConfig.getPatterns()).thenReturn(List.of("%{US_PHONE_NUMBER}")); - obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory); + obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory, expressionEvaluator); final Record record = createRecord(message); @@ -247,7 +292,7 @@ void testProcessorWithUSPhoneNumberPattern(String message, String expected) { }) void testProcessorWithCreditNumberPattern(String message, String expected) { when(mockConfig.getPatterns()).thenReturn(List.of("%{CREDIT_CARD_NUMBER}")); - obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory); + obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory, expressionEvaluator); final Record record = createRecord(message); @@ -271,7 +316,7 @@ void testProcessorWithCreditNumberPattern(String message, String expected) { }) void testProcessorWithIPAddressV4Pattern(String message, String expected) { when(mockConfig.getPatterns()).thenReturn(List.of("%{IP_ADDRESS_V4}")); - obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory); + obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory, expressionEvaluator); final Record record = createRecord(message); final List> editedRecords = (List>) obfuscationProcessor.doExecute(Collections.singletonList(record)); @@ -291,7 +336,7 @@ void testProcessorWithIPAddressV4Pattern(String message, String expected) { }) void testProcessorWithUSSSNPattern(String message, String expected) { when(mockConfig.getPatterns()).thenReturn(List.of("%{US_SSN_NUMBER}")); - obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory); + obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory, expressionEvaluator); final Record record = createRecord(message); final List> editedRecords = (List>) obfuscationProcessor.doExecute(Collections.singletonList(record)); @@ -314,7 +359,7 @@ void testProcessorWithUSSSNPattern(String message, String expected) { }) void testProcessorWithBaseNumberPattern(String message, String expected) { when(mockConfig.getPatterns()).thenReturn(List.of("%{BASE_NUMBER}")); - obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory); + obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory, expressionEvaluator); final Record record = createRecord(message); final List> editedRecords = (List>) obfuscationProcessor.doExecute(Collections.singletonList(record)); @@ -333,7 +378,7 @@ void testProcessorWithBaseNumberPattern(String message, String expected) { }) void testProcessorWithMultiplePatterns(String message, String expected) { when(mockConfig.getPatterns()).thenReturn(List.of("%{EMAIL_ADDRESS}", "%{IP_ADDRESS_V4}")); - obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory); + obfuscationProcessor = new ObfuscationProcessor(pluginMetrics, mockConfig, mockFactory, expressionEvaluator); final Record record = createRecord(message); final List> editedRecords = (List>) obfuscationProcessor.doExecute(Collections.singletonList(record)); From 4a112d40338b230781e8d7b04d8af4b8d80ce4e8 Mon Sep 17 00:00:00 2001 From: kkondaka <41027584+kkondaka@users.noreply.github.com> Date: Wed, 25 Oct 2023 14:34:18 -0700 Subject: [PATCH 10/53] Add tests to InMemoryBufferTest and LocalFileBufferTest (#3550) Signed-off-by: Krishna Kondaka Co-authored-by: Krishna Kondaka --- .../accumulator/InMemoryBufferTest.java | 21 ++++++++++++++++++- .../accumulator/LocalFileBufferTest.java | 19 ++++++++++++++++- 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBufferTest.java b/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBufferTest.java index 88d5205dca..2222919ed0 100644 --- a/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBufferTest.java +++ b/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/InMemoryBufferTest.java @@ -11,6 +11,11 @@ import org.mockito.junit.jupiter.MockitoExtension; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.hamcrest.Matchers.greaterThanOrEqualTo; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +import java.io.IOException; @ExtendWith(MockitoExtension.class) class InMemoryBufferTest { @@ -19,6 +24,20 @@ class InMemoryBufferTest { private InMemoryBuffer inMemoryBuffer; + @Test + void test_with_write_event_into_buffer() throws IOException { + inMemoryBuffer = new InMemoryBuffer(); + + int i = 0; + while (inMemoryBuffer.getEventCount() < MAX_EVENTS) { + inMemoryBuffer.getOutputStream().write(generateByteArray()); + inMemoryBuffer.setEventCount(++i); + } + assertThat(inMemoryBuffer.getSize(), greaterThanOrEqualTo(54110L)); + assertThat(inMemoryBuffer.getEventCount(), equalTo(MAX_EVENTS)); + assertThat(inMemoryBuffer.getDuration(), greaterThanOrEqualTo(0L)); + + } @Test void test_getSinkData_success() { @@ -36,4 +55,4 @@ private byte[] generateByteArray() { } return bytes; } -} \ No newline at end of file +} diff --git a/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBufferTest.java b/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBufferTest.java index 6090e9b57c..e6d65a23aa 100644 --- a/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBufferTest.java +++ b/data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/accumulator/LocalFileBufferTest.java @@ -16,6 +16,7 @@ import java.util.UUID; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.greaterThan; import static org.hamcrest.Matchers.equalTo; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; @@ -23,6 +24,7 @@ @ExtendWith(MockitoExtension.class) class LocalFileBufferTest { + public static final int MAX_EVENTS = 55; public static final String KEY = UUID.randomUUID().toString() + ".log"; public static final String PREFIX = "local"; @@ -37,6 +39,21 @@ void setUp() throws IOException { localFileBuffer = new LocalFileBuffer(tempFile); } + @Test + void test_with_write_events_into_buffer() throws IOException { + int i = 0; + while (localFileBuffer.getEventCount() < MAX_EVENTS) { + localFileBuffer.getOutputStream().write(generateByteArray()); + localFileBuffer.setEventCount(++i); + } + assertThat(localFileBuffer.getSize(), greaterThan(1l)); + assertThat(localFileBuffer.getEventCount(), equalTo(55)); + assertThat(localFileBuffer.getDuration(), equalTo(0L)); + localFileBuffer.flushAndCloseStream(); + localFileBuffer.removeTemporaryFile(); + assertFalse(tempFile.exists(), "The temp file has not been deleted."); + } + @Test void test_without_write_events_into_buffer() { assertThat(localFileBuffer.getSize(), equalTo(0L)); @@ -67,4 +84,4 @@ private byte[] generateByteArray() { } return bytes; } -} \ No newline at end of file +} From ad45c2fb728db9e4cab81e14ceb1169bf181e921 Mon Sep 17 00:00:00 2001 From: David Venable Date: Thu, 26 Oct 2023 09:50:44 -0700 Subject: [PATCH 11/53] Implement writeBytes and isByteBuffer in the CircuitBreakingBuffer. Also update the Buffer to throw UnsupportedOperationException when writeBytes is called. (#3553) Signed-off-by: David Venable --- .../dataprepper/model/buffer/Buffer.java | 2 +- .../dataprepper/model/buffer/BufferTest.java | 2 +- .../parser/CircuitBreakingBuffer.java | 12 ++++++ .../parser/CircuitBreakingBufferTest.java | 43 +++++++++++++++++++ 4 files changed, 57 insertions(+), 2 deletions(-) diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/Buffer.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/Buffer.java index ff0d712889..f8e3137285 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/Buffer.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/Buffer.java @@ -49,7 +49,7 @@ public interface Buffer> { * @throws RuntimeException Other exceptions */ default void writeBytes(final byte[] bytes, final String key, int timeoutInMillis) throws Exception { - throw new RuntimeException("Not supported"); + throw new UnsupportedOperationException("This buffer type does not support bytes."); } /** diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/BufferTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/BufferTest.java index bc9df29ba3..b7fc8505d4 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/BufferTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/BufferTest.java @@ -41,7 +41,7 @@ public void testWriteBytes() { final Buffer> buffer = spy(Buffer.class); byte[] bytes = new byte[2]; - Assert.assertThrows(RuntimeException.class, () -> buffer.writeBytes(bytes, "", 10)); + Assert.assertThrows(UnsupportedOperationException.class, () -> buffer.writeBytes(bytes, "", 10)); } diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/CircuitBreakingBuffer.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/CircuitBreakingBuffer.java index 0e63c682c4..2242f2db65 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/CircuitBreakingBuffer.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/CircuitBreakingBuffer.java @@ -53,6 +53,13 @@ public void writeAll(final Collection records, final int timeoutInMillis) thr buffer.writeAll(records, timeoutInMillis); } + @Override + public void writeBytes(byte[] bytes, String key, int timeoutInMillis) throws Exception { + checkBreaker(); + + buffer.writeBytes(bytes, key, timeoutInMillis); + } + private void checkBreaker() throws TimeoutException { if(circuitBreaker.isOpen()) throw new TimeoutException("Circuit breaker is open. Unable to write to buffer."); @@ -73,6 +80,11 @@ public boolean isEmpty() { return buffer.isEmpty(); } + @Override + public boolean isByteBuffer() { + return buffer.isByteBuffer(); + } + @Override public Duration getDrainTimeout() { return buffer.getDrainTimeout(); diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/CircuitBreakingBufferTest.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/CircuitBreakingBufferTest.java index 33e8832662..346d4fb8bf 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/CircuitBreakingBufferTest.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/CircuitBreakingBufferTest.java @@ -25,6 +25,7 @@ import java.util.List; import java.util.Map; import java.util.Random; +import java.util.UUID; import java.util.concurrent.TimeoutException; import static org.hamcrest.MatcherAssert.assertThat; @@ -81,6 +82,14 @@ void shutdown_calls_buffer_shutdown() { verify(buffer).shutdown(); } + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void isByteBuffer_returns_value_of_inner_buffer(boolean innerIsByteBuffer) { + when(buffer.isByteBuffer()).thenReturn(innerIsByteBuffer); + + assertThat(createObjectUnderTest().isByteBuffer(), equalTo(innerIsByteBuffer)); + } + @Nested class NoCircuitBreakerChecks { @AfterEach @@ -159,4 +168,38 @@ void writeAll_should_check_CircuitBreaker_and_throw_if_open() { verify(circuitBreaker).isOpen(); } } + + @Nested + class WithBytes { + private byte[] bytes; + + private String key; + + @BeforeEach + void setUp() { + bytes = UUID.randomUUID().toString().getBytes(); + key = UUID.randomUUID().toString(); + } + + @Test + void writeBytes_should_check_CircuitBreaker_and_call_inner_write_if_not_open() throws Exception { + when(circuitBreaker.isOpen()).thenReturn(false); + + createObjectUnderTest().writeBytes(bytes, key, timeoutMillis); + + verify(buffer).writeBytes(bytes, key, timeoutMillis); + verify(circuitBreaker).isOpen(); + } + + @Test + void writeBytes_should_check_CircuitBreaker_and_throw_if_open() { + when(circuitBreaker.isOpen()).thenReturn(true); + + CircuitBreakingBuffer> objectUnderTest = createObjectUnderTest(); + assertThrows(TimeoutException.class, () -> objectUnderTest.writeBytes(bytes, key, timeoutMillis)); + + verifyNoInteractions(buffer); + verify(circuitBreaker).isOpen(); + } + } } \ No newline at end of file From 267d3bcd6c0e3fedd968a659f57c10af9ff3cd51 Mon Sep 17 00:00:00 2001 From: wanghd89 Date: Fri, 27 Oct 2023 02:12:39 +0800 Subject: [PATCH 12/53] Fix CVE error for Jetty version 11.0.12 (#3548) * Fix CVE error for Jetty version 11.0.12 Signed-off-by: Haidong --------- Signed-off-by: Haidong Co-authored-by: Haidong --- data-prepper-plugins/s3-source/build.gradle | 1 + 1 file changed, 1 insertion(+) diff --git a/data-prepper-plugins/s3-source/build.gradle b/data-prepper-plugins/s3-source/build.gradle index 20f4d8ef0c..fe66c4a0ca 100644 --- a/data-prepper-plugins/s3-source/build.gradle +++ b/data-prepper-plugins/s3-source/build.gradle @@ -36,6 +36,7 @@ dependencies { implementation 'org.apache.httpcomponents:httpcore:4.4.15' testImplementation libs.commons.lang3 testImplementation 'com.github.tomakehurst:wiremock:3.0.0-beta-8' + testImplementation 'org.eclipse.jetty:jetty-bom:11.0.17' testImplementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml' testImplementation testLibs.junit.vintage testImplementation project(':data-prepper-test-common') From ca3d6acb077d1305d0214a07b9843814ca5190b2 Mon Sep 17 00:00:00 2001 From: Taylor Gray Date: Fri, 27 Oct 2023 10:37:39 -0500 Subject: [PATCH 13/53] Catch exceptions and backoff and retry ddb source threads instead of shutting down on exception (#3554) Signed-off-by: Taylor Gray --- .../dynamodb/converter/RecordConverter.java | 2 +- .../dynamodb/export/DataFileScheduler.java | 39 +++++++++------ .../dynamodb/export/ExportScheduler.java | 49 +++++++++++-------- .../dynamodb/stream/StreamScheduler.java | 37 +++++++++----- .../converter/ExportRecordConverterTest.java | 2 +- .../converter/StreamRecordConverterTest.java | 2 +- .../export/DataFileSchedulerTest.java | 15 ++++++ .../dynamodb/export/ExportSchedulerTest.java | 47 ++++++++++++------ .../dynamodb/stream/StreamSchedulerTest.java | 17 ++++++- 9 files changed, 144 insertions(+), 66 deletions(-) diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java index d5c5a059a8..4d5693db64 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java @@ -86,7 +86,7 @@ public void addToBuffer(Map data, Map keys, long String sortKey = getAttributeValue(keys, tableInfo.getMetadata().getSortKeyAttributeName()); if (sortKey != null) { eventMetadata.setAttribute(SORT_KEY_METADATA_ATTRIBUTE, sortKey); - eventMetadata.setAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE, partitionKey + "_" + sortKey); + eventMetadata.setAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE, partitionKey + "|" + sortKey); } else { eventMetadata.setAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE, partitionKey); } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java index 8c1979e320..91432f3dad 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java @@ -88,23 +88,32 @@ public void run() { LOG.debug("Start running Data File Scheduler"); while (!Thread.currentThread().isInterrupted()) { - if (numOfWorkers.get() < MAX_JOB_COUNT) { - final Optional sourcePartition = coordinator.acquireAvailablePartition(DataFilePartition.PARTITION_TYPE); - - if (sourcePartition.isPresent()) { - activeExportS3ObjectConsumersGauge.incrementAndGet(); - DataFilePartition dataFilePartition = (DataFilePartition) sourcePartition.get(); - processDataFilePartition(dataFilePartition); - activeExportS3ObjectConsumersGauge.decrementAndGet(); - } - } try { - Thread.sleep(DEFAULT_LEASE_INTERVAL_MILLIS); - } catch (final InterruptedException e) { - LOG.info("InterruptedException occurred"); - break; + if (numOfWorkers.get() < MAX_JOB_COUNT) { + final Optional sourcePartition = coordinator.acquireAvailablePartition(DataFilePartition.PARTITION_TYPE); + + if (sourcePartition.isPresent()) { + activeExportS3ObjectConsumersGauge.incrementAndGet(); + DataFilePartition dataFilePartition = (DataFilePartition) sourcePartition.get(); + processDataFilePartition(dataFilePartition); + activeExportS3ObjectConsumersGauge.decrementAndGet(); + } + } + try { + Thread.sleep(DEFAULT_LEASE_INTERVAL_MILLIS); + } catch (final InterruptedException e) { + LOG.info("The DataFileScheduler was interrupted while waiting to retry, stopping processing"); + break; + } + } catch (final Exception e) { + LOG.error("Received an exception while processing an S3 data file, backing off and retrying", e); + try { + Thread.sleep(DEFAULT_LEASE_INTERVAL_MILLIS); + } catch (final InterruptedException ex) { + LOG.info("The DataFileScheduler was interrupted while waiting to retry, stopping processing"); + break; + } } - } LOG.warn("Data file scheduler is interrupted, Stop all data file loaders..."); // Cannot call executor.shutdownNow() here diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java index 93f9e5b51a..9ea0d0dde9 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java @@ -94,32 +94,41 @@ public ExportScheduler(EnhancedSourceCoordinator enhancedSourceCoordinator, Dyna public void run() { LOG.debug("Start running Export Scheduler"); while (!Thread.currentThread().isInterrupted()) { - // Does not have limit on max leases - // As most of the time it's just to wait - final Optional sourcePartition = enhancedSourceCoordinator.acquireAvailablePartition(ExportPartition.PARTITION_TYPE); + try { + // Does not have limit on max leases + // As most of the time it's just to wait + final Optional sourcePartition = enhancedSourceCoordinator.acquireAvailablePartition(ExportPartition.PARTITION_TYPE); - if (sourcePartition.isPresent()) { + if (sourcePartition.isPresent()) { - ExportPartition exportPartition = (ExportPartition) sourcePartition.get(); - LOG.debug("Acquired an export partition: " + exportPartition.getPartitionKey()); + ExportPartition exportPartition = (ExportPartition) sourcePartition.get(); + LOG.debug("Acquired an export partition: " + exportPartition.getPartitionKey()); - String exportArn = getOrCreateExportArn(exportPartition); + String exportArn = getOrCreateExportArn(exportPartition); - if (exportArn == null) { - closeExportPartitionWithError(exportPartition); - } else { - CompletableFuture checkStatus = CompletableFuture.supplyAsync(() -> checkExportStatus(exportPartition), executor); - checkStatus.whenComplete(completeExport(exportPartition)); - } + if (exportArn == null) { + closeExportPartitionWithError(exportPartition); + } else { + CompletableFuture checkStatus = CompletableFuture.supplyAsync(() -> checkExportStatus(exportPartition), executor); + checkStatus.whenComplete(completeExport(exportPartition)); + } + } + try { + Thread.sleep(DEFAULT_TAKE_LEASE_INTERVAL_MILLIS); + } catch (final InterruptedException e) { + LOG.info("The ExportScheduler was interrupted while waiting to retry, stopping processing"); + break; + } + } catch (final Exception e) { + LOG.error("Received an exception during export from DynamoDB to S3, backing off and retrying", e); + try { + Thread.sleep(DEFAULT_TAKE_LEASE_INTERVAL_MILLIS); + } catch (final InterruptedException ex) { + LOG.info("The ExportScheduler was interrupted while waiting to retry, stopping processing"); + break; + } } - try { - Thread.sleep(DEFAULT_TAKE_LEASE_INTERVAL_MILLIS); - } catch (final InterruptedException e) { - LOG.info("InterruptedException occurred"); - break; - } - } LOG.warn("Export scheduler interrupted, looks like shutdown has triggered"); executor.shutdownNow(); diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamScheduler.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamScheduler.java index 19e4dd28e0..cd4d43baaf 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamScheduler.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamScheduler.java @@ -83,21 +83,32 @@ private void processStreamPartition(StreamPartition streamPartition) { public void run() { LOG.debug("Stream Scheduler start to run..."); while (!Thread.currentThread().isInterrupted()) { - if (numOfWorkers.get() < MAX_JOB_COUNT) { - final Optional sourcePartition = coordinator.acquireAvailablePartition(StreamPartition.PARTITION_TYPE); - if (sourcePartition.isPresent()) { - activeChangeEventConsumers.incrementAndGet(); - StreamPartition streamPartition = (StreamPartition) sourcePartition.get(); - processStreamPartition(streamPartition); - activeChangeEventConsumers.decrementAndGet(); + try { + + if (numOfWorkers.get() < MAX_JOB_COUNT) { + final Optional sourcePartition = coordinator.acquireAvailablePartition(StreamPartition.PARTITION_TYPE); + if (sourcePartition.isPresent()) { + activeChangeEventConsumers.incrementAndGet(); + StreamPartition streamPartition = (StreamPartition) sourcePartition.get(); + processStreamPartition(streamPartition); + activeChangeEventConsumers.decrementAndGet(); + } } - } - try { - Thread.sleep(DEFAULT_LEASE_INTERVAL_MILLIS); - } catch (final InterruptedException e) { - LOG.info("InterruptedException occurred"); - break; + try { + Thread.sleep(DEFAULT_LEASE_INTERVAL_MILLIS); + } catch (final InterruptedException e) { + LOG.info("InterruptedException occurred"); + break; + } + } catch (final Exception e) { + LOG.error("Received an exception while trying to get an S3 data file to process, backing off and retrying", e); + try { + Thread.sleep(DEFAULT_LEASE_INTERVAL_MILLIS); + } catch (final InterruptedException ex) { + LOG.info("The StreamScheduler was interrupted while waiting to retry, stopping processing"); + break; + } } } // Should Stop diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java index e79d231cb3..a31a642b8e 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java @@ -137,7 +137,7 @@ void test_writeSingleRecordToBuffer() throws Exception { assertThat(event.getMetadata().getAttribute(PARTITION_KEY_METADATA_ATTRIBUTE), equalTo(pk)); assertThat(event.getMetadata().getAttribute(SORT_KEY_METADATA_ATTRIBUTE), equalTo(sk)); - assertThat(event.getMetadata().getAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE), equalTo(pk + "_" + sk)); + assertThat(event.getMetadata().getAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE), equalTo(pk + "|" + sk)); assertThat(event.getMetadata().getAttribute(EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE), equalTo(OpenSearchBulkActions.INDEX.toString())); assertThat(event.getMetadata().getAttribute(EVENT_TIMESTAMP_METADATA_ATTRIBUTE), notNullValue()); } diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java index 23ac32f6df..cd02172da2 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java @@ -133,7 +133,7 @@ void test_writeSingleRecordToBuffer() throws Exception { String sortKey = record.dynamodb().keys().get(sortKeyAttrName).s(); assertThat(event.getMetadata().getAttribute(PARTITION_KEY_METADATA_ATTRIBUTE), equalTo(partitionKey)); assertThat(event.getMetadata().getAttribute(SORT_KEY_METADATA_ATTRIBUTE), equalTo(sortKey)); - assertThat(event.getMetadata().getAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE), equalTo(partitionKey + "_" + sortKey)); + assertThat(event.getMetadata().getAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE), equalTo(partitionKey + "|" + sortKey)); assertThat(event.getMetadata().getAttribute(EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE), equalTo(OpenSearchBulkActions.CREATE.toString())); assertThat(event.getMetadata().getAttribute(EVENT_TIMESTAMP_METADATA_ATTRIBUTE), equalTo(record.dynamodb().approximateCreationDateTime().toEpochMilli())); diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileSchedulerTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileSchedulerTest.java index dd7562341b..66378ee406 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileSchedulerTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileSchedulerTest.java @@ -151,5 +151,20 @@ public void test_run_DataFileLoader_correctly() throws InterruptedException { } + @Test + void run_catches_exception_and_retries_when_exception_is_thrown_during_processing() throws InterruptedException { + given(coordinator.acquireAvailablePartition(DataFilePartition.PARTITION_TYPE)).willThrow(RuntimeException.class); + + scheduler = new DataFileScheduler(coordinator, loaderFactory, pluginMetrics); + + ExecutorService executorService = Executors.newSingleThreadExecutor(); + final Future future = executorService.submit(() -> scheduler.run()); + Thread.sleep(100); + assertThat(future.isDone(), equalTo(false)); + executorService.shutdown(); + future.cancel(true); + assertThat(executorService.awaitTermination(1000, TimeUnit.MILLISECONDS), equalTo(true)); + } + } \ No newline at end of file diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportSchedulerTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportSchedulerTest.java index 2a1506643f..0a74f2c94f 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportSchedulerTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportSchedulerTest.java @@ -33,7 +33,11 @@ import java.util.UUID; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.BDDMockito.given; @@ -43,10 +47,10 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; -import static org.opensearch.dataprepper.plugins.source.dynamodb.export.ExportScheduler.EXPORT_S3_OBJECTS_TOTAL_COUNT; import static org.opensearch.dataprepper.plugins.source.dynamodb.export.ExportScheduler.EXPORT_JOB_FAILURE_COUNT; import static org.opensearch.dataprepper.plugins.source.dynamodb.export.ExportScheduler.EXPORT_JOB_SUCCESS_COUNT; import static org.opensearch.dataprepper.plugins.source.dynamodb.export.ExportScheduler.EXPORT_RECORDS_TOTAL_COUNT; +import static org.opensearch.dataprepper.plugins.source.dynamodb.export.ExportScheduler.EXPORT_S3_OBJECTS_TOTAL_COUNT; @ExtendWith(MockitoExtension.class) @@ -99,19 +103,6 @@ class ExportSchedulerTest { @BeforeEach void setup() { - when(exportPartition.getTableArn()).thenReturn(tableArn); - when(exportPartition.getExportTime()).thenReturn(exportTime); - - ExportProgressState state = new ExportProgressState(); - state.setBucket(bucketName); - state.setPrefix(prefix); - when(exportPartition.getProgressState()).thenReturn(Optional.of(state)); - - given(pluginMetrics.counter(EXPORT_JOB_SUCCESS_COUNT)).willReturn(exportJobSuccess); - given(pluginMetrics.counter(EXPORT_JOB_FAILURE_COUNT)).willReturn(exportJobErrors); - given(pluginMetrics.counter(EXPORT_S3_OBJECTS_TOTAL_COUNT)).willReturn(exportFilesTotal); - given(pluginMetrics.counter(EXPORT_RECORDS_TOTAL_COUNT)).willReturn(exportRecordsTotal); - ExportSummary summary = mock(ExportSummary.class); lenient().when(manifestFileReader.parseSummaryFile(anyString(), anyString())).thenReturn(summary); lenient().when(summary.getS3Bucket()).thenReturn(bucketName); @@ -127,6 +118,19 @@ void setup() { @Test public void test_run_exportJob_correctly() throws InterruptedException { + when(exportPartition.getTableArn()).thenReturn(tableArn); + when(exportPartition.getExportTime()).thenReturn(exportTime); + + ExportProgressState state = new ExportProgressState(); + state.setBucket(bucketName); + state.setPrefix(prefix); + when(exportPartition.getProgressState()).thenReturn(Optional.of(state)); + + given(pluginMetrics.counter(EXPORT_JOB_SUCCESS_COUNT)).willReturn(exportJobSuccess); + given(pluginMetrics.counter(EXPORT_JOB_FAILURE_COUNT)).willReturn(exportJobErrors); + given(pluginMetrics.counter(EXPORT_S3_OBJECTS_TOTAL_COUNT)).willReturn(exportFilesTotal); + given(pluginMetrics.counter(EXPORT_RECORDS_TOTAL_COUNT)).willReturn(exportRecordsTotal); + given(coordinator.acquireAvailablePartition(ExportPartition.PARTITION_TYPE)).willReturn(Optional.of(exportPartition)).willReturn(Optional.empty()); // Set up mock behavior @@ -169,4 +173,19 @@ public void test_run_exportJob_correctly() throws InterruptedException { } + @Test + void run_catches_exception_and_retries_when_exception_is_thrown_during_processing() throws InterruptedException { + given(coordinator.acquireAvailablePartition(ExportPartition.PARTITION_TYPE)).willThrow(RuntimeException.class); + + scheduler = new ExportScheduler(coordinator, dynamoDBClient, manifestFileReader, pluginMetrics); + + ExecutorService executorService = Executors.newSingleThreadExecutor(); + final Future future = executorService.submit(() -> scheduler.run()); + Thread.sleep(100); + assertThat(future.isDone(), equalTo(false)); + executorService.shutdown(); + future.cancel(true); + assertThat(executorService.awaitTermination(1000, TimeUnit.MILLISECONDS), equalTo(true)); + } + } \ No newline at end of file diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamSchedulerTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamSchedulerTest.java index c041323a0d..586a3c188a 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamSchedulerTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamSchedulerTest.java @@ -40,7 +40,6 @@ import static org.opensearch.dataprepper.plugins.source.dynamodb.stream.StreamScheduler.ACTIVE_CHANGE_EVENT_CONSUMERS; @ExtendWith(MockitoExtension.class) -@Disabled class StreamSchedulerTest { @Mock @@ -104,6 +103,7 @@ void setup() { @Test + @Disabled public void test_normal_run() throws InterruptedException { given(coordinator.acquireAvailablePartition(StreamPartition.PARTITION_TYPE)).willReturn(Optional.of(streamPartition)).willReturn(Optional.empty()); @@ -128,4 +128,19 @@ public void test_normal_run() throws InterruptedException { executorService.shutdownNow(); } + + @Test + void run_catches_exception_and_retries_when_exception_is_thrown_during_processing() throws InterruptedException { + given(coordinator.acquireAvailablePartition(StreamPartition.PARTITION_TYPE)).willThrow(RuntimeException.class); + + scheduler = new StreamScheduler(coordinator, consumerFactory, shardManager, pluginMetrics); + + ExecutorService executorService = Executors.newSingleThreadExecutor(); + final Future future = executorService.submit(() -> scheduler.run()); + Thread.sleep(100); + assertThat(future.isDone(), equalTo(false)); + executorService.shutdown(); + future.cancel(true); + assertThat(executorService.awaitTermination(1000, TimeUnit.MILLISECONDS), equalTo(true)); + } } \ No newline at end of file From cb765d8478448482dffcfe16df30bdc970025ed3 Mon Sep 17 00:00:00 2001 From: David Venable Date: Fri, 27 Oct 2023 14:39:02 -0700 Subject: [PATCH 14/53] Fix for Kafka buffer encryption with bytes serde_format by returning null for null input in EncryptionSerializer. (#3556) Signed-off-by: David Venable --- .../common/serialization/EncryptionSerializer.java | 3 +++ .../common/serialization/EncryptionSerializerTest.java | 10 ++++++++++ 2 files changed, 13 insertions(+) diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/common/serialization/EncryptionSerializer.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/common/serialization/EncryptionSerializer.java index 7710b78c8a..d9cce50677 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/common/serialization/EncryptionSerializer.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/common/serialization/EncryptionSerializer.java @@ -28,6 +28,9 @@ class EncryptionSerializer implements Serializer { @Override public byte[] serialize(String topic, T data) { + if(data == null) + return null; + byte[] unencryptedBytes = innerSerializer.serialize(topic, data); try { return cipher.doFinal(unencryptedBytes); diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/common/serialization/EncryptionSerializerTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/common/serialization/EncryptionSerializerTest.java index 941b27e1ea..342db03d9f 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/common/serialization/EncryptionSerializerTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/common/serialization/EncryptionSerializerTest.java @@ -16,7 +16,9 @@ import java.util.UUID; import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.nullValue; import static org.hamcrest.MatcherAssert.assertThat; +import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; @ExtendWith(MockitoExtension.class) @@ -54,4 +56,12 @@ void serialize_performs_cipher_encryption_on_serialized_data() throws IllegalBlo assertThat(createObjectUnderTest().serialize(topicName, input), equalTo(encryptedData)); } + + @Test + void serialize_returns_null_and_does_not_call_cipher_if_input_is_null() throws NoSuchPaddingException, NoSuchAlgorithmException, InvalidKeyException { + assertThat(createObjectUnderTest().serialize(topicName, null), + nullValue()); + + verifyNoInteractions(cipher); + } } \ No newline at end of file From 7323c9e29a0ea29167bcb0b706e4b9ea80c9f587 Mon Sep 17 00:00:00 2001 From: David Venable Date: Tue, 31 Oct 2023 10:20:09 -0700 Subject: [PATCH 15/53] Updates Kafka configurations such that plugin has its own topic configurations (#3551) Updates Kafka buffer configurations to disallow configurations which are not needed - e.g. schema and the serde_format for the topic configuration. As part of this change, I also split the TopicConfig into three distinct interfaces and classes. This allows each plugin to either accept a configuration or provide a value of the plugin's own choosing. Also adds copyright headers to all files modified as part of this commit. Renamed is_topic_create to create_topic. Also made this a boolean internally instead of Boolean since it will have a value. Adds a zeroBytes() static method to ByteCount as a convenience. Signed-off-by: David Venable --- .../dataprepper/model/types/ByteCount.java | 5 + .../model/types/ByteCountTest.java | 14 + .../plugins/kafka/buffer/KafkaBufferIT.java | 11 +- .../kafka/sink/KafkaSinkAvroTypeIT.java | 31 +- .../kafka/sink/KafkaSinkJsonTypeIT.java | 31 +- .../kafka/sink/KafkaSinkPlainTextTypeIT.java | 30 +- .../ConfluentKafkaProducerConsumerIT.java | 50 ++-- ...aProducerConsumerWithSchemaRegistryIT.java | 62 ++-- .../plugins/kafka/source/JSONConsumerIT.java | 9 +- .../kafka/source/KafkaSourceJsonTypeIT.java | 16 +- .../source/KafkaSourceMultipleAuthTypeIT.java | 14 +- .../source/MskGlueRegistryMultiTypeIT.java | 13 +- .../kafka/source/PlainTextConsumerIT.java | 9 +- .../kafka/buffer/BufferTopicConfig.java | 229 +++++++++++++++ .../plugins/kafka/buffer/KafkaBuffer.java | 6 +- .../KafkaBufferConfig.java | 29 +- .../kafka/common/KafkaDataConfigAdapter.java | 8 +- .../configuration/CommonTopicConfig.java | 49 ++++ .../configuration/KafkaConsumerConfig.java | 8 +- .../configuration/KafkaProducerConfig.java | 7 +- .../kafka/configuration/TopicConfig.java | 276 +----------------- .../configuration/TopicConsumerConfig.java | 44 +++ .../configuration/TopicProducerConfig.java | 17 ++ .../kafka/consumer/KafkaCustomConsumer.java | 8 +- .../consumer/KafkaCustomConsumerFactory.java | 12 +- .../producer/KafkaCustomProducerFactory.java | 12 +- .../plugins/kafka/sink/KafkaSink.java | 9 +- .../KafkaSinkConfig.java | 18 +- .../plugins/kafka/sink/SinkTopicConfig.java | 68 +++++ .../plugins/kafka/source/KafkaSource.java | 8 +- .../KafkaSourceConfig.java | 14 +- .../kafka/source/SourceTopicConfig.java | 200 +++++++++++++ .../kafka/util/KafkaSecurityConfigurer.java | 2 +- .../kafka/buffer/BufferTopicConfigTest.java | 56 ++++ .../plugins/kafka/buffer/KafkaBufferTest.java | 11 +- .../kafka/configuration/AuthConfigTest.java | 6 + .../configuration/CommonTopicConfigTest.java | 84 ++++++ .../kafka/configuration/OAuthConfigTest.java | 6 + .../PlainTextAuthConfigTest.java | 6 + .../kafka/configuration/SchemaConfigTest.java | 6 + .../kafka/configuration/TopicConfigTest.java | 142 --------- .../consumer/KafkaCustomConsumerTest.java | 13 +- .../producer/KafkaCustomProducerTest.java | 11 +- .../plugins/kafka/sink/DLQSinkTest.java | 1 - .../KafkaSinkConfigTest.java | 6 +- .../plugins/kafka/sink/KafkaSinkTest.java | 6 +- .../KafkaSourceConfigTest.java | 16 +- .../plugins/kafka/source/KafkaSourceTest.java | 48 ++- .../kafka/source/SourceTopicConfigTest.java | 56 ++++ .../AuthenticationPropertyConfigurerTest.java | 2 +- .../util/SinkPropertyConfigurerTest.java | 9 +- .../test/resources/sample-pipelines-sink.yaml | 2 +- 52 files changed, 1118 insertions(+), 688 deletions(-) create mode 100644 data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/BufferTopicConfig.java rename data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/{configuration => buffer}/KafkaBufferConfig.java (72%) create mode 100644 data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/CommonTopicConfig.java create mode 100644 data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/TopicConsumerConfig.java create mode 100644 data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/TopicProducerConfig.java rename data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/{configuration => sink}/KafkaSinkConfig.java (82%) create mode 100644 data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/sink/SinkTopicConfig.java rename data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/{configuration => source}/KafkaSourceConfig.java (79%) create mode 100644 data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/source/SourceTopicConfig.java create mode 100644 data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/buffer/BufferTopicConfigTest.java create mode 100644 data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/CommonTopicConfigTest.java delete mode 100644 data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/TopicConfigTest.java rename data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/{configuration => sink}/KafkaSinkConfigTest.java (92%) rename data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/{configuration => source}/KafkaSourceConfigTest.java (86%) create mode 100644 data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/source/SourceTopicConfigTest.java diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/types/ByteCount.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/types/ByteCount.java index ae09e953ff..1681563675 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/types/ByteCount.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/types/ByteCount.java @@ -21,6 +21,7 @@ */ public class ByteCount { private static final Pattern BYTE_PATTERN = Pattern.compile("^(?\\d+\\.?\\d*)(?[a-z]+)?\\z"); + private static final ByteCount ZERO_BYTES = new ByteCount(0); private final long bytes; private ByteCount(final long bytes) { @@ -94,6 +95,10 @@ public static ByteCount parse(final String string) { return new ByteCount(byteCount.longValue()); } + public static ByteCount zeroBytes() { + return ZERO_BYTES; + } + private static BigDecimal scaleToBytes(final BigDecimal value, final Unit unit) { return value.multiply(BigDecimal.valueOf(unit.multiplier)); } diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/types/ByteCountTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/types/ByteCountTest.java index eda34eae69..b717289a7f 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/types/ByteCountTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/types/ByteCountTest.java @@ -5,6 +5,7 @@ package org.opensearch.dataprepper.model.types; +import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; import org.junit.jupiter.params.provider.ValueSource; @@ -13,6 +14,7 @@ import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.notNullValue; +import static org.hamcrest.Matchers.sameInstance; import static org.junit.jupiter.api.Assertions.assertThrows; class ByteCountTest { @@ -145,4 +147,16 @@ void parse_returns_rounded_bytes_for_implicit_fractional_bytes(final String byte assertThat(byteCount, notNullValue()); assertThat(byteCount.getBytes(), equalTo(expectedBytes)); } + + @Test + void zeroBytes_returns_bytes_with_getBytes_equal_to_0() { + assertThat(ByteCount.zeroBytes(), notNullValue()); + assertThat(ByteCount.zeroBytes().getBytes(), equalTo(0L)); + } + + @Test + void zeroBytes_returns_same_instance() { + assertThat(ByteCount.zeroBytes(), notNullValue()); + assertThat(ByteCount.zeroBytes(), sameInstance(ByteCount.zeroBytes())); + } } \ No newline at end of file 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 4de082d415..9f2a782ddf 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 @@ -1,3 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + package org.opensearch.dataprepper.plugins.kafka.buffer; import org.apache.commons.lang3.RandomStringUtils; @@ -15,8 +20,6 @@ import org.opensearch.dataprepper.model.plugin.PluginFactory; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaBufferConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -50,7 +53,7 @@ public class KafkaBufferIT { @Mock private AcknowledgementSetManager acknowledgementSetManager; @Mock - private TopicConfig topicConfig; + private BufferTopicConfig topicConfig; private PluginMetrics pluginMetrics; private String bootstrapServersCommaDelimited; @@ -66,7 +69,7 @@ void setUp() { String topicName = "buffer-" + RandomStringUtils.randomAlphabetic(5); when(topicConfig.getName()).thenReturn(topicName); when(topicConfig.getGroupId()).thenReturn("buffergroup-" + RandomStringUtils.randomAlphabetic(6)); - when(topicConfig.isCreate()).thenReturn(true); + when(topicConfig.isCreateTopic()).thenReturn(true); when(topicConfig.getSerdeFormat()).thenReturn(messageFormat); when(topicConfig.getWorkers()).thenReturn(1); when(topicConfig.getMaxPollInterval()).thenReturn(Duration.ofSeconds(5)); diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkAvroTypeIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkAvroTypeIT.java index 3b75384581..da901ed241 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkAvroTypeIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkAvroTypeIT.java @@ -36,14 +36,12 @@ import org.opensearch.dataprepper.plugins.dlq.DlqProvider; import org.opensearch.dataprepper.plugins.dlq.DlqWriter; import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSinkConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.PlainTextAuthConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicProducerConfig; import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; import java.io.IOException; -import java.time.Duration; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -71,7 +69,7 @@ public class KafkaSinkAvroTypeIT { private KafkaSinkConfig kafkaSinkConfig; @Mock - private TopicConfig topicConfig; + private TopicProducerConfig topicConfig; private KafkaSink kafkaSink; @@ -145,18 +143,10 @@ public void setup() throws RestClientException, IOException { when(kafkaSinkConfig.getSerdeFormat()).thenReturn(MessageFormat.AVRO.toString()); when(kafkaSinkConfig.getPartitionKey()).thenReturn("test-${name}"); - final String testGroup = "TestGroup_" + RandomStringUtils.randomAlphabetic(5); testTopic = "TestTopic_" + RandomStringUtils.randomAlphabetic(5); - topicConfig = mock(TopicConfig.class); + topicConfig = mock(TopicProducerConfig.class); when(topicConfig.getName()).thenReturn(testTopic); - when(topicConfig.getGroupId()).thenReturn(testGroup); - when(topicConfig.getWorkers()).thenReturn(1); - when(topicConfig.getSessionTimeOut()).thenReturn(Duration.ofSeconds(45)); - when(topicConfig.getHeartBeatInterval()).thenReturn(Duration.ofSeconds(3)); - when(topicConfig.getAutoCommit()).thenReturn(false); - when(topicConfig.getAutoOffsetReset()).thenReturn("earliest"); - when(topicConfig.getThreadWaitingTime()).thenReturn(Duration.ofSeconds(1)); bootstrapServers = System.getProperty("tests.kafka.bootstrap_servers"); when(kafkaSinkConfig.getBootstrapServers()).thenReturn(Collections.singletonList(bootstrapServers)); @@ -171,8 +161,7 @@ public void TestPollRecordsAvroSASLPlainText() throws Exception { configureJasConfForSASLPlainText(); final int numRecords = 1; - when(topicConfig.getConsumerMaxPollRecords()).thenReturn(numRecords); - when(topicConfig.isCreate()).thenReturn(false); + when(topicConfig.isCreateTopic()).thenReturn(false); when(kafkaSinkConfig.getTopic()).thenReturn(topicConfig); when(kafkaSinkConfig.getAuthConfig()).thenReturn(authConfig); @@ -250,16 +239,8 @@ private void deleteTopic(AtomicBoolean created, String topicName) throws Interru } private void consumeTestMessages(List> recList) { - - props.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, - topicConfig.getCommitInterval().toSecondsPart()); - props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, - topicConfig.getAutoOffsetReset()); - props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, - topicConfig.getAutoCommit()); - props.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, - topicConfig.getConsumerMaxPollRecords()); - props.put(ConsumerConfig.GROUP_ID_CONFIG, topicConfig.getGroupId()); + final String testGroup = "TestGroup_" + RandomStringUtils.randomAlphabetic(5); + props.put(ConsumerConfig.GROUP_ID_CONFIG, testGroup); props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class); props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkJsonTypeIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkJsonTypeIT.java index 3d825925ad..bf2d1cb3df 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkJsonTypeIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkJsonTypeIT.java @@ -34,12 +34,10 @@ import org.opensearch.dataprepper.plugins.dlq.DlqProvider; import org.opensearch.dataprepper.plugins.dlq.DlqWriter; import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSinkConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.PlainTextAuthConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicProducerConfig; import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; -import java.time.Duration; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -63,7 +61,7 @@ public class KafkaSinkJsonTypeIT { private KafkaSinkConfig kafkaSinkConfig; @Mock - private TopicConfig topicConfig; + private TopicProducerConfig topicConfig; private KafkaSink kafkaSink; @@ -123,18 +121,10 @@ public void setup() { when(kafkaSinkConfig.getSerdeFormat()).thenReturn(MessageFormat.JSON.toString()); when(kafkaSinkConfig.getPartitionKey()).thenReturn("test-${name}"); - final String testGroup = "TestGroup_" + RandomStringUtils.randomAlphabetic(5); testTopic = "TestTopic_" + RandomStringUtils.randomAlphabetic(5); - topicConfig = mock(TopicConfig.class); + topicConfig = mock(TopicProducerConfig.class); when(topicConfig.getName()).thenReturn(testTopic); - when(topicConfig.getGroupId()).thenReturn(testGroup); - when(topicConfig.getWorkers()).thenReturn(1); - when(topicConfig.getSessionTimeOut()).thenReturn(Duration.ofSeconds(45)); - when(topicConfig.getHeartBeatInterval()).thenReturn(Duration.ofSeconds(3)); - when(topicConfig.getAutoCommit()).thenReturn(false); - when(topicConfig.getAutoOffsetReset()).thenReturn("earliest"); - when(topicConfig.getThreadWaitingTime()).thenReturn(Duration.ofSeconds(1)); bootstrapServers = System.getProperty("tests.kafka.bootstrap_servers"); when(kafkaSinkConfig.getBootstrapServers()).thenReturn(Collections.singletonList(bootstrapServers)); props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers); @@ -145,8 +135,7 @@ public void TestPollRecordsJsonSASLPlainText() throws Exception { configureJasConfForSASLPlainText(); final int numRecords = 1; - when(topicConfig.getConsumerMaxPollRecords()).thenReturn(numRecords); - when(topicConfig.isCreate()).thenReturn(false); + when(topicConfig.isCreateTopic()).thenReturn(false); when(kafkaSinkConfig.getTopic()).thenReturn(topicConfig); when(kafkaSinkConfig.getAuthConfig()).thenReturn(authConfig); kafkaSink = createObjectUnderTest(); @@ -223,16 +212,8 @@ private void configureJasConfForSASLPlainText() { } private void consumeTestMessages(List> recList) { - - props.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, - topicConfig.getCommitInterval().toSecondsPart()); - props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, - topicConfig.getAutoOffsetReset()); - props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, - topicConfig.getAutoCommit()); - props.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, - topicConfig.getConsumerMaxPollRecords()); - props.put(ConsumerConfig.GROUP_ID_CONFIG, topicConfig.getGroupId()); + final String testGroup = "TestGroup_" + RandomStringUtils.randomAlphabetic(5); + props.put(ConsumerConfig.GROUP_ID_CONFIG, testGroup); props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class); props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkPlainTextTypeIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkPlainTextTypeIT.java index cd7ac9526f..c00304d4e8 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkPlainTextTypeIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkPlainTextTypeIT.java @@ -32,12 +32,10 @@ import org.opensearch.dataprepper.plugins.dlq.DlqProvider; import org.opensearch.dataprepper.plugins.dlq.DlqWriter; import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSinkConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.PlainTextAuthConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicProducerConfig; import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; -import java.time.Duration; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -61,7 +59,7 @@ public class KafkaSinkPlainTextTypeIT { private KafkaSinkConfig kafkaSinkConfig; @Mock - private TopicConfig topicConfig; + private TopicProducerConfig topicConfig; private KafkaSink kafkaSink; @@ -121,18 +119,10 @@ public void setup() { when(kafkaSinkConfig.getSerdeFormat()).thenReturn(MessageFormat.PLAINTEXT.toString()); when(kafkaSinkConfig.getPartitionKey()).thenReturn("test-${name}"); - final String testGroup = "TestGroup_" + RandomStringUtils.randomAlphabetic(5); testTopic = "TestTopic_" + RandomStringUtils.randomAlphabetic(5); - topicConfig = mock(TopicConfig.class); + topicConfig = mock(TopicProducerConfig.class); when(topicConfig.getName()).thenReturn(testTopic); - when(topicConfig.getGroupId()).thenReturn(testGroup); - when(topicConfig.getWorkers()).thenReturn(1); - when(topicConfig.getSessionTimeOut()).thenReturn(Duration.ofSeconds(45)); - when(topicConfig.getHeartBeatInterval()).thenReturn(Duration.ofSeconds(3)); - when(topicConfig.getAutoCommit()).thenReturn(false); - when(topicConfig.getAutoOffsetReset()).thenReturn("earliest"); - when(topicConfig.getThreadWaitingTime()).thenReturn(Duration.ofSeconds(1)); bootstrapServers = System.getProperty("tests.kafka.bootstrap_servers"); when(kafkaSinkConfig.getBootstrapServers()).thenReturn(Collections.singletonList(bootstrapServers)); props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers); @@ -144,8 +134,7 @@ public void TestPollRecordsPlainText() throws Exception { configureJasConfForSASLPlainText(); final int numRecords = 1; - when(topicConfig.getConsumerMaxPollRecords()).thenReturn(numRecords); - when(topicConfig.isCreate()).thenReturn(false); + when(topicConfig.isCreateTopic()).thenReturn(false); when(kafkaSinkConfig.getTopic()).thenReturn(topicConfig); when(kafkaSinkConfig.getAuthConfig()).thenReturn(authConfig); kafkaSink = createObjectUnderTest(); @@ -221,16 +210,9 @@ private void deleteTopic(AtomicBoolean created, String topicName) throws Interru } private void consumeTestMessages(List> recList) { + final String testGroup = "TestGroup_" + RandomStringUtils.randomAlphabetic(5); - props.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, - topicConfig.getCommitInterval().toSecondsPart()); - props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, - topicConfig.getAutoOffsetReset()); - props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, - topicConfig.getAutoCommit()); - props.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, - topicConfig.getConsumerMaxPollRecords()); - props.put(ConsumerConfig.GROUP_ID_CONFIG, topicConfig.getGroupId()); + props.put(ConsumerConfig.GROUP_ID_CONFIG, testGroup); props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class); props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/ConfluentKafkaProducerConsumerIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/ConfluentKafkaProducerConsumerIT.java index 0b0c8ef85e..ea11ec22b6 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/ConfluentKafkaProducerConsumerIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/ConfluentKafkaProducerConsumerIT.java @@ -5,49 +5,45 @@ package org.opensearch.dataprepper.plugins.kafka.source; +import io.micrometer.core.instrument.Counter; import org.apache.kafka.clients.producer.KafkaProducer; import org.apache.kafka.clients.producer.ProducerConfig; import org.apache.kafka.clients.producer.ProducerRecord; +import org.apache.kafka.common.errors.SerializationException; import org.apache.kafka.common.serialization.StringSerializer; -import org.junit.jupiter.api.Test; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; import org.mockito.Mock; 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.configuration.PipelineDescription; import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionType; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSourceConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.PlainTextAuthConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; -import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; -import org.opensearch.dataprepper.model.configuration.PipelineDescription; - -import static org.mockito.Mockito.when; -import org.mockito.Mock; -import static org.mockito.Mockito.mock; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.Mockito.doAnswer; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.CoreMatchers.equalTo; -import static org.awaitility.Awaitility.await; +import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; -import io.micrometer.core.instrument.Counter; -import java.util.List; +import java.time.Duration; +import java.time.Instant; import java.util.ArrayList; import java.util.Collection; +import java.util.Iterator; +import java.util.List; import java.util.Properties; import java.util.concurrent.atomic.AtomicInteger; -import java.util.Iterator; - -import java.time.Duration; -import java.time.Instant; -import org.apache.kafka.common.errors.SerializationException; -import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; +import static org.awaitility.Awaitility.await; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; public class ConfluentKafkaProducerConsumerIT { @Mock @@ -78,7 +74,7 @@ public class ConfluentKafkaProducerConsumerIT { private PlainTextAuthConfig plainTextAuthConfig; private KafkaSource kafkaSource; - private TopicConfig topicConfig; + private TopicConsumerConfig topicConfig; private Counter counter; private List receivedRecords; @@ -129,7 +125,7 @@ public void setup() { topicName = System.getProperty("tests.kafka.topic_name"); username = System.getProperty("tests.kafka.username"); password = System.getProperty("tests.kafka.password"); - topicConfig = mock(TopicConfig.class); + topicConfig = mock(TopicConsumerConfig.class); when(topicConfig.getName()).thenReturn(topicName); when(topicConfig.getGroupId()).thenReturn("testGroupConf"); when(topicConfig.getWorkers()).thenReturn(1); @@ -147,7 +143,7 @@ public void setup() { when(topicConfig.getFetchMaxWait()).thenReturn(500); when(topicConfig.getMaxPartitionFetchBytes()).thenReturn(1024L*1024); when(topicConfig.getReconnectBackoff()).thenReturn(Duration.ofSeconds(10)); - when(sourceConfig.getTopics()).thenReturn(List.of(topicConfig)); + when(sourceConfig.getTopics()).thenReturn((List) List.of(topicConfig)); when(sourceConfig.getBootstrapServers()).thenReturn(List.of(bootstrapServers)); encryptionConfig = mock(EncryptionConfig.class); when(sourceConfig.getEncryptionConfig()).thenReturn(encryptionConfig); diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/ConfluentKafkaProducerConsumerWithSchemaRegistryIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/ConfluentKafkaProducerConsumerWithSchemaRegistryIT.java index e37bf67799..4ffc153c4e 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/ConfluentKafkaProducerConsumerWithSchemaRegistryIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/ConfluentKafkaProducerConsumerWithSchemaRegistryIT.java @@ -5,56 +5,52 @@ package org.opensearch.dataprepper.plugins.kafka.source; +import com.fasterxml.jackson.annotation.JsonProperty; +import io.micrometer.core.instrument.Counter; +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.apache.commons.lang3.RandomStringUtils; import org.apache.kafka.clients.producer.KafkaProducer; import org.apache.kafka.clients.producer.ProducerConfig; import org.apache.kafka.clients.producer.ProducerRecord; +import org.apache.kafka.common.errors.SerializationException; import org.apache.kafka.common.serialization.StringSerializer; -import org.junit.jupiter.api.Test; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; import org.mockito.Mock; 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.configuration.PipelineDescription; import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionType; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSourceConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.PlainTextAuthConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaRegistryType; -import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; -import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; -import org.opensearch.dataprepper.model.configuration.PipelineDescription; -import org.apache.avro.Schema; - -import static org.mockito.Mockito.when; -import org.mockito.Mock; -import static org.mockito.Mockito.mock; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.Mockito.doAnswer; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.CoreMatchers.equalTo; -import static org.awaitility.Awaitility.await; -import org.apache.commons.lang3.RandomStringUtils; +import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; -import io.micrometer.core.instrument.Counter; -import java.util.List; +import java.time.Duration; import java.util.ArrayList; import java.util.Collection; import java.util.Iterator; +import java.util.List; import java.util.Properties; import java.util.Random; import java.util.concurrent.atomic.AtomicInteger; -import java.time.Duration; - -import org.apache.avro.generic.GenericRecord; -import org.apache.avro.generic.GenericData; -import org.apache.kafka.common.errors.SerializationException; -import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; -import com.fasterxml.jackson.annotation.JsonProperty; +import static org.awaitility.Awaitility.await; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; public class ConfluentKafkaProducerConsumerWithSchemaRegistryIT { public static class AvroRecord { @@ -131,8 +127,8 @@ public UserRecord(String name, Integer id, Number value) { private PlainTextAuthConfig plainTextAuthConfig; private KafkaSource kafkaSource; - private TopicConfig jsonTopicConfig; - private TopicConfig avroTopicConfig; + private TopicConsumerConfig jsonTopicConfig; + private TopicConsumerConfig avroTopicConfig; private Counter counter; private List receivedRecords; @@ -196,7 +192,7 @@ public void setup() { username = System.getProperty("tests.kafka.username"); password = System.getProperty("tests.kafka.password"); - jsonTopicConfig = mock(TopicConfig.class); + jsonTopicConfig = mock(TopicConsumerConfig.class); jsonTopicName = System.getProperty("tests.kafka.json_topic_name"); when(jsonTopicConfig.getName()).thenReturn(jsonTopicName); when(jsonTopicConfig.getGroupId()).thenReturn("testGroupConf"); @@ -210,7 +206,7 @@ public void setup() { when(jsonTopicConfig.getConsumerMaxPollRecords()).thenReturn(100); when(jsonTopicConfig.getMaxPollInterval()).thenReturn(Duration.ofSeconds(15)); - avroTopicConfig = mock(TopicConfig.class); + avroTopicConfig = mock(TopicConsumerConfig.class); avroTopicName = System.getProperty("tests.kafka.avro_topic_name"); when(avroTopicConfig.getName()).thenReturn(avroTopicName); when(avroTopicConfig.getGroupId()).thenReturn("testGroupConf"); @@ -229,14 +225,14 @@ public void setup() { when(jsonSourceConfig.getAuthConfig()).thenReturn(authConfig); when(jsonSourceConfig.getAcknowledgementsEnabled()).thenReturn(false); when(jsonSourceConfig.getSchemaConfig()).thenReturn(schemaConfig); - when(jsonSourceConfig.getTopics()).thenReturn(List.of(jsonTopicConfig)); + when(jsonSourceConfig.getTopics()).thenReturn((List) List.of(jsonTopicConfig)); when(jsonSourceConfig.getBootstrapServers()).thenReturn(List.of(bootstrapServers)); when(jsonSourceConfig.getEncryptionConfig()).thenReturn(encryptionConfig); when(avroSourceConfig.getAuthConfig()).thenReturn(authConfig); when(avroSourceConfig.getAcknowledgementsEnabled()).thenReturn(false); when(avroSourceConfig.getSchemaConfig()).thenReturn(schemaConfig); - when(avroSourceConfig.getTopics()).thenReturn(List.of(avroTopicConfig)); + when(avroSourceConfig.getTopics()).thenReturn((List) List.of(avroTopicConfig)); when(avroSourceConfig.getBootstrapServers()).thenReturn(List.of(bootstrapServers)); when(avroSourceConfig.getEncryptionConfig()).thenReturn(encryptionConfig); diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/JSONConsumerIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/JSONConsumerIT.java index cc777b25df..7f9c30f58f 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/JSONConsumerIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/JSONConsumerIT.java @@ -24,18 +24,15 @@ import org.mockito.junit.jupiter.MockitoExtension; 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.event.Event; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSourceConfig; +import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.yaml.snakeyaml.Yaml; import java.io.FileReader; import java.io.IOException; import java.io.Reader; import java.io.StringReader; -import java.util.List; import java.util.Map; import java.util.Properties; @@ -44,8 +41,6 @@ public class JSONConsumerIT { private PluginMetrics pluginMetrics; @Mock - TopicConfig topicConfig; - @Mock private SchemaConfig schemaConfig; private KafkaSourceConfig kafkaSourceConfig; @@ -75,8 +70,6 @@ public void configure() throws IOException { String json = mapper.writeValueAsString(kafkaConfigMap); Reader reader = new StringReader(json); kafkaSourceConfig = mapper.readValue(reader, KafkaSourceConfig.class); - List topicConfigList = kafkaSourceConfig.getTopics(); - topicConfig = topicConfigList.get(0); schemaConfig = kafkaSourceConfig.getSchemaConfig(); } } diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceJsonTypeIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceJsonTypeIT.java index 36fd90a100..17ab0d85bb 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceJsonTypeIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceJsonTypeIT.java @@ -25,10 +25,10 @@ import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.event.EventMetadata; import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionType; import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaKeyMode; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSourceConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.opensearch.dataprepper.plugins.kafka.extension.KafkaClusterConfigSupplier; import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; @@ -82,7 +82,7 @@ public class KafkaSourceJsonTypeIT { private KafkaClusterConfigSupplier kafkaClusterConfigSupplier; @Mock - private TopicConfig jsonTopic; + private TopicConsumerConfig jsonTopic; private KafkaSource kafkaSource; @@ -126,7 +126,7 @@ public void setup() throws Throwable { testKey = RandomStringUtils.randomAlphabetic(5); testGroup = "TestGroup_" + RandomStringUtils.randomAlphabetic(6); testTopic = "TestJsonTopic_" + RandomStringUtils.randomAlphabetic(5); - jsonTopic = mock(TopicConfig.class); + jsonTopic = mock(TopicConsumerConfig.class); when(jsonTopic.getName()).thenReturn(testTopic); when(jsonTopic.getGroupId()).thenReturn(testGroup); when(jsonTopic.getWorkers()).thenReturn(1); @@ -182,7 +182,7 @@ public void TestJsonRecordsWithNullKey() throws Exception { when(encryptionConfig.getType()).thenReturn(EncryptionType.NONE); when(jsonTopic.getConsumerMaxPollRecords()).thenReturn(numRecords); when(jsonTopic.getKafkaKeyMode()).thenReturn(KafkaKeyMode.INCLUDE_AS_FIELD); - when(sourceConfig.getTopics()).thenReturn(List.of(jsonTopic)); + when(sourceConfig.getTopics()).thenReturn((List) List.of(jsonTopic)); when(sourceConfig.getAuthConfig()).thenReturn(null); kafkaSource = createObjectUnderTest(); @@ -214,7 +214,7 @@ public void TestJsonRecordsWithNegativeAcknowledgements() throws Exception { when(encryptionConfig.getType()).thenReturn(EncryptionType.NONE); when(jsonTopic.getConsumerMaxPollRecords()).thenReturn(numRecords); when(jsonTopic.getKafkaKeyMode()).thenReturn(KafkaKeyMode.DISCARD); - when(sourceConfig.getTopics()).thenReturn(List.of(jsonTopic)); + when(sourceConfig.getTopics()).thenReturn((List) List.of(jsonTopic)); when(sourceConfig.getAuthConfig()).thenReturn(null); when(sourceConfig.getAcknowledgementsEnabled()).thenReturn(true); kafkaSource = createObjectUnderTest(); @@ -264,7 +264,7 @@ public void TestJsonRecordsWithKafkaKeyModeDiscard() throws Exception { when(encryptionConfig.getType()).thenReturn(EncryptionType.NONE); when(jsonTopic.getConsumerMaxPollRecords()).thenReturn(numRecords); when(jsonTopic.getKafkaKeyMode()).thenReturn(KafkaKeyMode.DISCARD); - when(sourceConfig.getTopics()).thenReturn(List.of(jsonTopic)); + when(sourceConfig.getTopics()).thenReturn((List) List.of(jsonTopic)); when(sourceConfig.getAuthConfig()).thenReturn(null); kafkaSource = createObjectUnderTest(); @@ -294,7 +294,7 @@ public void TestJsonRecordsWithKafkaKeyModeAsField() throws Exception { when(encryptionConfig.getType()).thenReturn(EncryptionType.NONE); when(jsonTopic.getConsumerMaxPollRecords()).thenReturn(numRecords); when(jsonTopic.getKafkaKeyMode()).thenReturn(KafkaKeyMode.INCLUDE_AS_FIELD); - when(sourceConfig.getTopics()).thenReturn(List.of(jsonTopic)); + when(sourceConfig.getTopics()).thenReturn((List) List.of(jsonTopic)); when(sourceConfig.getAuthConfig()).thenReturn(null); kafkaSource = createObjectUnderTest(); @@ -325,7 +325,7 @@ public void TestJsonRecordsWithKafkaKeyModeAsMetadata() throws Exception { when(encryptionConfig.getType()).thenReturn(EncryptionType.NONE); when(jsonTopic.getConsumerMaxPollRecords()).thenReturn(numRecords); when(jsonTopic.getKafkaKeyMode()).thenReturn(KafkaKeyMode.INCLUDE_AS_METADATA); - when(sourceConfig.getTopics()).thenReturn(List.of(jsonTopic)); + when(sourceConfig.getTopics()).thenReturn((List) List.of(jsonTopic)); when(sourceConfig.getAuthConfig()).thenReturn(null); kafkaSource = createObjectUnderTest(); diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceMultipleAuthTypeIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceMultipleAuthTypeIT.java index 693c7dc8af..acc4ca5a46 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceMultipleAuthTypeIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceMultipleAuthTypeIT.java @@ -23,9 +23,9 @@ import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionType; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSourceConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.PlainTextAuthConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.opensearch.dataprepper.plugins.kafka.extension.KafkaClusterConfigSupplier; @@ -67,7 +67,7 @@ public class KafkaSourceMultipleAuthTypeIT { private List topicList; @Mock - private TopicConfig plainTextTopic; + private TopicConsumerConfig plainTextTopic; @Mock private AuthConfig authConfig; @@ -128,7 +128,7 @@ public void setup() { final String testGroup = "TestGroup_"+RandomStringUtils.randomAlphabetic(6); final String testTopic = "TestTopic_"+RandomStringUtils.randomAlphabetic(5); - plainTextTopic = mock(TopicConfig.class); + plainTextTopic = mock(TopicConsumerConfig.class); when(plainTextTopic.getName()).thenReturn(testTopic); when(plainTextTopic.getGroupId()).thenReturn(testGroup); when(plainTextTopic.getWorkers()).thenReturn(1); @@ -154,7 +154,7 @@ public void TestPlainTextWithNoAuthKafkaNoEncryptionWithNoAuthSchemaRegistry() t final int numRecords = 1; when(encryptionConfig.getType()).thenReturn(EncryptionType.NONE); when(plainTextTopic.getConsumerMaxPollRecords()).thenReturn(numRecords); - when(sourceConfig.getTopics()).thenReturn(List.of(plainTextTopic)); + when(sourceConfig.getTopics()).thenReturn((List) List.of(plainTextTopic)); when(sourceConfig.getAuthConfig()).thenReturn(null); kafkaSource = createObjectUnderTest(); @@ -203,7 +203,7 @@ public void TestPlainTextWithAuthKafkaNoEncryptionWithNoAuthSchemaRegistry() thr saslAuthConfig = mock(AuthConfig.SaslAuthConfig.class); when(encryptionConfig.getType()).thenReturn(EncryptionType.NONE); when(plainTextTopic.getConsumerMaxPollRecords()).thenReturn(numRecords); - when(sourceConfig.getTopics()).thenReturn(List.of(plainTextTopic)); + when(sourceConfig.getTopics()).thenReturn((List) List.of(plainTextTopic)); plainTextAuthConfig = mock(PlainTextAuthConfig.class); when(plainTextAuthConfig.getUsername()).thenReturn(kafkaUsername); when(plainTextAuthConfig.getPassword()).thenReturn(kafkaPassword); @@ -262,7 +262,7 @@ public void TestPlainTextWithNoAuthKafkaEncryptionWithNoAuthSchemaRegistry() thr when(encryptionConfig.getType()).thenReturn(EncryptionType.SSL); when(plainTextTopic.getConsumerMaxPollRecords()).thenReturn(numRecords); when(sourceConfig.getBootstrapServers()).thenReturn(Collections.singletonList(sslBootstrapServers)); - when(sourceConfig.getTopics()).thenReturn(List.of(plainTextTopic)); + when(sourceConfig.getTopics()).thenReturn((List) List.of(plainTextTopic)); kafkaSource = createObjectUnderTest(); Properties props = new Properties(); @@ -318,7 +318,7 @@ public void TestPlainTextWithAuthKafkaEncryptionWithNoAuthSchemaRegistry() throw when(encryptionConfig.getType()).thenReturn(EncryptionType.SSL); when(plainTextTopic.getConsumerMaxPollRecords()).thenReturn(numRecords); when(sourceConfig.getBootstrapServers()).thenReturn(Collections.singletonList(saslsslBootstrapServers)); - when(sourceConfig.getTopics()).thenReturn(List.of(plainTextTopic)); + when(sourceConfig.getTopics()).thenReturn((List) List.of(plainTextTopic)); kafkaSource = createObjectUnderTest(); Properties props = new Properties(); diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/MskGlueRegistryMultiTypeIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/MskGlueRegistryMultiTypeIT.java index 4091a66966..b7e680daae 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/MskGlueRegistryMultiTypeIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/MskGlueRegistryMultiTypeIT.java @@ -36,7 +36,6 @@ import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionType; import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaKeyMode; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSourceConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.MskBrokerConnectionType; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaRegistryType; @@ -105,8 +104,8 @@ public class MskGlueRegistryMultiTypeIT { private KafkaClusterConfigSupplier kafkaClusterConfigSupplier; private KafkaSource kafkaSource; - private TopicConfig jsonTopic; - private TopicConfig avroTopic; + private SourceTopicConfig jsonTopic; + private SourceTopicConfig avroTopic; private Counter counter; @@ -162,8 +161,8 @@ public void setup() { final String testGroup = "TestGroup_"+RandomStringUtils.randomAlphabetic(6); final String testTopic = "TestTopic_"+RandomStringUtils.randomAlphabetic(5); - avroTopic = mock(TopicConfig.class); - jsonTopic = mock(TopicConfig.class); + avroTopic = mock(SourceTopicConfig.class); + jsonTopic = mock(SourceTopicConfig.class); when(avroTopic.getName()).thenReturn(testTopic); when(avroTopic.getGroupId()).thenReturn(testGroup); when(avroTopic.getWorkers()).thenReturn(1); @@ -201,7 +200,7 @@ public void TestJsonRecordConsumer() throws Exception { final int numRecords = 1; when(encryptionConfig.getType()).thenReturn(EncryptionType.SSL); when(jsonTopic.getConsumerMaxPollRecords()).thenReturn(numRecords); - when(sourceConfig.getTopics()).thenReturn(List.of(jsonTopic)); + when(sourceConfig.getTopics()).thenReturn((List) List.of(jsonTopic)); when(sourceConfig.getAuthConfig()).thenReturn(authConfig); when(authConfig.getSaslAuthConfig()).thenReturn(saslAuthConfig); when(saslAuthConfig.getAwsIamAuthConfig()).thenReturn(AwsIamAuthConfig.DEFAULT); @@ -268,7 +267,7 @@ public void TestAvroRecordConsumer() throws Exception { final int numRecords = 1; when(encryptionConfig.getType()).thenReturn(EncryptionType.SSL); when(avroTopic.getConsumerMaxPollRecords()).thenReturn(numRecords); - when(sourceConfig.getTopics()).thenReturn(List.of(avroTopic)); + when(sourceConfig.getTopics()).thenReturn((List) List.of(avroTopic)); when(sourceConfig.getAuthConfig()).thenReturn(authConfig); when(authConfig.getSaslAuthConfig()).thenReturn(saslAuthConfig); when(saslAuthConfig.getAwsIamAuthConfig()).thenReturn(AwsIamAuthConfig.DEFAULT); diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/PlainTextConsumerIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/PlainTextConsumerIT.java index a5118e64c5..7936cbadcf 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/PlainTextConsumerIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/PlainTextConsumerIT.java @@ -19,18 +19,15 @@ import org.mockito.Mock; 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.event.Event; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSourceConfig; +import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.yaml.snakeyaml.Yaml; import java.io.FileReader; import java.io.IOException; import java.io.Reader; import java.io.StringReader; -import java.util.List; import java.util.Map; import java.util.Properties; @@ -39,8 +36,6 @@ public class PlainTextConsumerIT { private PluginMetrics pluginMetrics; @Mock - TopicConfig topicConfig; - @Mock private SchemaConfig schemaConfig; private KafkaSourceConfig kafkaSourceConfig; @@ -70,8 +65,6 @@ public void configure() throws IOException { String json = mapper.writeValueAsString(kafkaConfigMap); Reader reader = new StringReader(json); kafkaSourceConfig = mapper.readValue(reader, KafkaSourceConfig.class); - List topicConfigList = kafkaSourceConfig.getTopics(); - topicConfig = topicConfigList.get(0); schemaConfig = kafkaSourceConfig.getSchemaConfig(); } } diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/BufferTopicConfig.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/BufferTopicConfig.java new file mode 100644 index 0000000000..87c92975b6 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/BufferTopicConfig.java @@ -0,0 +1,229 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.buffer; + +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.Valid; +import jakarta.validation.constraints.Size; +import org.opensearch.dataprepper.model.types.ByteCount; +import org.opensearch.dataprepper.plugins.kafka.configuration.CommonTopicConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaKeyMode; +import org.opensearch.dataprepper.plugins.kafka.configuration.KmsConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicProducerConfig; +import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; + +import java.time.Duration; + +class BufferTopicConfig extends CommonTopicConfig implements TopicProducerConfig, TopicConsumerConfig { + static final Duration DEFAULT_COMMIT_INTERVAL = Duration.ofSeconds(5); + private static final Integer DEFAULT_NUM_OF_PARTITIONS = 1; + private static final Short DEFAULT_REPLICATION_FACTOR = 1; + private static final Long DEFAULT_RETENTION_PERIOD = 604800000L; + static final boolean DEFAULT_AUTO_COMMIT = false; + static final ByteCount DEFAULT_FETCH_MAX_BYTES = ByteCount.parse("50mb"); + static final Duration DEFAULT_FETCH_MAX_WAIT = Duration.ofMillis(500); + static final ByteCount DEFAULT_FETCH_MIN_BYTES = ByteCount.parse("1b"); + static final ByteCount DEFAULT_MAX_PARTITION_FETCH_BYTES = ByteCount.parse("1mb"); + static final Duration DEFAULT_SESSION_TIMEOUT = Duration.ofSeconds(45); + static final String DEFAULT_AUTO_OFFSET_RESET = "earliest"; + static final Duration DEFAULT_THREAD_WAITING_TIME = Duration.ofSeconds(5); + static final Duration DEFAULT_MAX_POLL_INTERVAL = Duration.ofSeconds(300); + static final Integer DEFAULT_CONSUMER_MAX_POLL_RECORDS = 500; + static final Integer DEFAULT_NUM_OF_WORKERS = 2; + static final Duration DEFAULT_HEART_BEAT_INTERVAL_DURATION = Duration.ofSeconds(5); + + + @JsonProperty("encryption_key") + private String encryptionKey; + + @JsonProperty("kms") + private KmsConfig kmsConfig; + + @JsonProperty("commit_interval") + @Valid + @Size(min = 1) + private Duration commitInterval = DEFAULT_COMMIT_INTERVAL; + + @JsonProperty("number_of_partitions") + private Integer numberOfPartitions = DEFAULT_NUM_OF_PARTITIONS; + + @JsonProperty("replication_factor") + private Short replicationFactor = DEFAULT_REPLICATION_FACTOR; + + @JsonProperty("retention_period") + private Long retentionPeriod = DEFAULT_RETENTION_PERIOD; + + @JsonProperty("create_topic") + private boolean isCreateTopic = false; + + @JsonProperty("group_id") + @Valid + @Size(min = 1, max = 255, message = "size of group id should be between 1 and 255") + private String groupId; + + @JsonProperty("workers") + @Valid + @Size(min = 1, max = 200, message = "Number of worker threads should lies between 1 and 200") + private Integer workers = DEFAULT_NUM_OF_WORKERS; + + @JsonProperty("session_timeout") + @Valid + @Size(min = 1) + private Duration sessionTimeOut = DEFAULT_SESSION_TIMEOUT; + + @JsonProperty("auto_offset_reset") + private String autoOffsetReset = DEFAULT_AUTO_OFFSET_RESET; + + @JsonProperty("thread_waiting_time") + private Duration threadWaitingTime = DEFAULT_THREAD_WAITING_TIME; + + @JsonProperty("max_poll_interval") + private Duration maxPollInterval = DEFAULT_MAX_POLL_INTERVAL; + + @JsonProperty("consumer_max_poll_records") + private Integer consumerMaxPollRecords = DEFAULT_CONSUMER_MAX_POLL_RECORDS; + + @JsonProperty("heart_beat_interval") + @Valid + @Size(min = 1) + private Duration heartBeatInterval= DEFAULT_HEART_BEAT_INTERVAL_DURATION; + + @JsonProperty("auto_commit") + private Boolean autoCommit = DEFAULT_AUTO_COMMIT; + + @JsonProperty("max_partition_fetch_bytes") + private ByteCount maxPartitionFetchBytes = DEFAULT_MAX_PARTITION_FETCH_BYTES; + + @JsonProperty("fetch_max_bytes") + private ByteCount fetchMaxBytes = DEFAULT_FETCH_MAX_BYTES; + + @JsonProperty("fetch_max_wait") + @Valid + private Duration fetchMaxWait = DEFAULT_FETCH_MAX_WAIT; + + @JsonProperty("fetch_min_bytes") + private ByteCount fetchMinBytes = DEFAULT_FETCH_MIN_BYTES; + + @Override + public MessageFormat getSerdeFormat() { + return MessageFormat.BYTES; + } + + @Override + public String getEncryptionKey() { + return encryptionKey; + } + + @Override + public KmsConfig getKmsConfig() { + return kmsConfig; + } + + @Override + public KafkaKeyMode getKafkaKeyMode() { + return KafkaKeyMode.DISCARD; + } + + @Override + public String getGroupId() { + return groupId; + } + + @Override + public Duration getCommitInterval() { + return commitInterval; + } + + + @Override + public Integer getNumberOfPartitions() { + return numberOfPartitions; + } + + @Override + public Short getReplicationFactor() { + return replicationFactor; + } + + @Override + public Long getRetentionPeriod() { + return retentionPeriod; + } + + @Override + public boolean isCreateTopic() { + return isCreateTopic; + } + + @Override + public Boolean getAutoCommit() { + return autoCommit; + } + + @Override + public Duration getSessionTimeOut() { + return sessionTimeOut; + } + + @Override + public String getAutoOffsetReset() { + return autoOffsetReset; + } + + @Override + public Duration getThreadWaitingTime() { + return threadWaitingTime; + } + + @Override + public Duration getMaxPollInterval() { + return maxPollInterval; + } + + @Override + public Integer getConsumerMaxPollRecords() { + return consumerMaxPollRecords; + } + + @Override + public Integer getWorkers() { + return workers; + } + + @Override + public Duration getHeartBeatInterval() { + return heartBeatInterval; + } + + @Override + public long getFetchMaxBytes() { + long value = fetchMaxBytes.getBytes(); + if (value < 1 || value > 50*1024*1024) { + throw new RuntimeException("Invalid Fetch Max Bytes"); + } + return value; + } + + @Override + public Integer getFetchMaxWait() { + return Math.toIntExact(fetchMaxWait.toMillis()); + } + + @Override + public long getFetchMinBytes() { + long value = fetchMinBytes.getBytes(); + if (value < 1) { + throw new RuntimeException("Invalid Fetch Min Bytes"); + } + return value; + } + + @Override + public long getMaxPartitionFetchBytes() { + return maxPartitionFetchBytes.getBytes(); + } +} 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 b139fe5db9..86a934e972 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 @@ -1,3 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + package org.opensearch.dataprepper.plugins.kafka.buffer; import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; @@ -14,7 +19,6 @@ import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.buffer.blockingbuffer.BlockingBuffer; import org.opensearch.dataprepper.plugins.kafka.common.serialization.SerializationFactory; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaBufferConfig; import org.opensearch.dataprepper.plugins.kafka.consumer.KafkaCustomConsumer; import org.opensearch.dataprepper.plugins.kafka.consumer.KafkaCustomConsumerFactory; import org.opensearch.dataprepper.plugins.kafka.producer.KafkaCustomProducer; diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaBufferConfig.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferConfig.java similarity index 72% rename from data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaBufferConfig.java rename to data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferConfig.java index 1cab8d7133..0e0eea2837 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaBufferConfig.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferConfig.java @@ -1,4 +1,9 @@ -package org.opensearch.dataprepper.plugins.kafka.configuration; +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.buffer; import com.fasterxml.jackson.annotation.JsonProperty; import jakarta.validation.Valid; @@ -6,13 +11,21 @@ import jakarta.validation.constraints.Size; import org.opensearch.dataprepper.model.configuration.PluginModel; import org.opensearch.dataprepper.model.configuration.PluginSetting; +import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.AwsConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaConsumerConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerProperties; +import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; import java.time.Duration; import java.util.List; import java.util.Objects; import java.util.Optional; -public class KafkaBufferConfig implements KafkaProducerConfig, KafkaConsumerConfig { +class KafkaBufferConfig implements KafkaProducerConfig, KafkaConsumerConfig { private static final Duration DEFAULT_DRAIN_TIMEOUT = Duration.ofSeconds(30); @JsonProperty("bootstrap_servers") @@ -21,11 +34,7 @@ public class KafkaBufferConfig implements KafkaProducerConfig, KafkaConsumerConf @JsonProperty("topics") @NotNull @Size(min = 1, max = 1, message = "Only one topic currently supported for Kafka buffer") - private List topics; - - @JsonProperty("schema") - @Valid - private SchemaConfig schemaConfig; + private List topics; @Valid @JsonProperty("authentication") @@ -59,7 +68,7 @@ public AuthConfig getAuthConfig() { @Override public SchemaConfig getSchemaConfig() { - return schemaConfig; + return null; } @Override @@ -68,12 +77,12 @@ public String getSerdeFormat() { } @Override - public TopicConfig getTopic() { + public BufferTopicConfig getTopic() { return topics.get(0); } @Override - public List getTopics() { + public List getTopics() { return topics; } diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/common/KafkaDataConfigAdapter.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/common/KafkaDataConfigAdapter.java index e7d38bf8ba..c4be81fe26 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/common/KafkaDataConfigAdapter.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/common/KafkaDataConfigAdapter.java @@ -1,13 +1,19 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + package org.opensearch.dataprepper.plugins.kafka.common; import org.opensearch.dataprepper.plugins.kafka.common.key.KeyFactory; +import org.opensearch.dataprepper.plugins.kafka.configuration.CommonTopicConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; import java.util.function.Supplier; /** - * Adapts a {@link TopicConfig} to a {@link KafkaDataConfig}. + * Adapts a {@link CommonTopicConfig} to a {@link KafkaDataConfig}. */ public class KafkaDataConfigAdapter implements KafkaDataConfig { private final KeyFactory keyFactory; diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/CommonTopicConfig.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/CommonTopicConfig.java new file mode 100644 index 0000000000..c0ad938c80 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/CommonTopicConfig.java @@ -0,0 +1,49 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.configuration; + +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.Valid; +import jakarta.validation.constraints.NotNull; + +import java.time.Duration; + +/** + * This class has topic configurations which are common to all Kafka plugins - source, buffer, and sink. + *

+ * Be sure to only add to this configuration if the setting is applicable for all three types. + */ +public abstract class CommonTopicConfig implements TopicConfig { + static final Duration DEFAULT_RETRY_BACKOFF = Duration.ofSeconds(10); + static final Duration DEFAULT_RECONNECT_BACKOFF = Duration.ofSeconds(10); + + @JsonProperty("name") + @NotNull + @Valid + private String name; + + @JsonProperty("retry_backoff") + private Duration retryBackoff = DEFAULT_RETRY_BACKOFF; + + @JsonProperty("reconnect_backoff") + private Duration reconnectBackoff = DEFAULT_RECONNECT_BACKOFF; + + + @Override + public Duration getRetryBackoff() { + return retryBackoff; + } + + @Override + public Duration getReconnectBackoff() { + return reconnectBackoff; + } + + @Override + public String getName() { + return name; + } +} diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaConsumerConfig.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaConsumerConfig.java index 27b16feb53..c52573bb6b 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaConsumerConfig.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaConsumerConfig.java @@ -1,14 +1,18 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + package org.opensearch.dataprepper.plugins.kafka.configuration; import java.util.List; public interface KafkaConsumerConfig extends KafkaConnectionConfig { - String getClientDnsLookup(); boolean getAcknowledgementsEnabled(); SchemaConfig getSchemaConfig(); - List getTopics(); + List getTopics(); } diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaProducerConfig.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaProducerConfig.java index b08f97aca2..a46deb3412 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaProducerConfig.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaProducerConfig.java @@ -1,3 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + package org.opensearch.dataprepper.plugins.kafka.configuration; import org.opensearch.dataprepper.model.configuration.PluginModel; @@ -15,7 +20,7 @@ public interface KafkaProducerConfig extends KafkaConnectionConfig { String getSerdeFormat(); - TopicConfig getTopic(); + TopicProducerConfig getTopic(); KafkaProducerProperties getKafkaProducerProperties(); diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/TopicConfig.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/TopicConfig.java index 546587a15a..97358e3b1c 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/TopicConfig.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/TopicConfig.java @@ -5,281 +5,25 @@ package org.opensearch.dataprepper.plugins.kafka.configuration; -import com.fasterxml.jackson.annotation.JsonProperty; -import jakarta.validation.Valid; -import jakarta.validation.constraints.NotNull; -import jakarta.validation.constraints.Size; -import org.opensearch.dataprepper.model.types.ByteCount; import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; import java.time.Duration; /** - * * A helper class that helps to read consumer configuration values from - * pipelines.yaml + * Represents a topic configuration to use throughout the code. See the + * {@link TopicConsumerConfig} and {@link TopicProducerConfig} for configurations + * which are specific for those use-cases. */ -public class TopicConfig { - static final boolean DEFAULT_AUTO_COMMIT = false; - static final Duration DEFAULT_COMMIT_INTERVAL = Duration.ofSeconds(5); - static final Duration DEFAULT_SESSION_TIMEOUT = Duration.ofSeconds(45); - static final String DEFAULT_AUTO_OFFSET_RESET = "earliest"; - static final Duration DEFAULT_THREAD_WAITING_TIME = Duration.ofSeconds(5); - static final Duration DEFAULT_MAX_RECORD_FETCH_TIME = Duration.ofSeconds(4); - static final String DEFAULT_FETCH_MAX_BYTES = "50mb"; - static final Integer DEFAULT_FETCH_MAX_WAIT = 500; - static final String DEFAULT_FETCH_MIN_BYTES = "1b"; - static final Duration DEFAULT_RETRY_BACKOFF = Duration.ofSeconds(10); - static final Duration DEFAULT_RECONNECT_BACKOFF = Duration.ofSeconds(10); - static final String DEFAULT_MAX_PARTITION_FETCH_BYTES = "1mb"; - static final Duration DEFAULT_MAX_POLL_INTERVAL = Duration.ofSeconds(300); - static final Integer DEFAULT_CONSUMER_MAX_POLL_RECORDS = 500; - static final Integer DEFAULT_NUM_OF_WORKERS = 2; - static final Duration DEFAULT_HEART_BEAT_INTERVAL_DURATION = Duration.ofSeconds(5); +public interface TopicConfig { + String getName(); + MessageFormat getSerdeFormat(); - private static final Integer DEFAULT_NUM_OF_PARTITIONS = 1; - private static final Short DEFAULT_REPLICATION_FACTOR = 1; - private static final Long DEFAULT_RETENTION_PERIOD=604800000L; + String getEncryptionKey(); + KmsConfig getKmsConfig(); - @JsonProperty("name") - @NotNull - @Valid - private String name; + Duration getRetryBackoff(); - @JsonProperty("group_id") - @Valid - @Size(min = 1, max = 255, message = "size of group id should be between 1 and 255") - private String groupId; - - @JsonProperty("workers") - @Valid - @Size(min = 1, max = 200, message = "Number of worker threads should lies between 1 and 200") - private Integer workers = DEFAULT_NUM_OF_WORKERS; - - @JsonProperty("serde_format") - private MessageFormat serdeFormat= MessageFormat.PLAINTEXT; - - @JsonProperty("auto_commit") - private Boolean autoCommit = DEFAULT_AUTO_COMMIT; - - @JsonProperty("commit_interval") - @Valid - @Size(min = 1) - private Duration commitInterval = DEFAULT_COMMIT_INTERVAL; - - @JsonProperty("session_timeout") - @Valid - @Size(min = 1) - private Duration sessionTimeOut = DEFAULT_SESSION_TIMEOUT; - - @JsonProperty("auto_offset_reset") - private String autoOffsetReset = DEFAULT_AUTO_OFFSET_RESET; - - @JsonProperty("thread_waiting_time") - private Duration threadWaitingTime = DEFAULT_THREAD_WAITING_TIME; - - @JsonProperty("max_partition_fetch_bytes") - private String maxPartitionFetchBytes = DEFAULT_MAX_PARTITION_FETCH_BYTES; - - @JsonProperty("fetch_max_bytes") - private String fetchMaxBytes = DEFAULT_FETCH_MAX_BYTES; - - @JsonProperty("fetch_max_wait") - @Valid - @Size(min = 1) - private Integer fetchMaxWait = DEFAULT_FETCH_MAX_WAIT; - - @JsonProperty("fetch_min_bytes") - private String fetchMinBytes = DEFAULT_FETCH_MIN_BYTES; - - @JsonProperty("key_mode") - private KafkaKeyMode kafkaKeyMode = KafkaKeyMode.INCLUDE_AS_FIELD; - - @JsonProperty("retry_backoff") - private Duration retryBackoff = DEFAULT_RETRY_BACKOFF; - - @JsonProperty("reconnect_backoff") - private Duration reconnectBackoff = DEFAULT_RECONNECT_BACKOFF; - - @JsonProperty("max_poll_interval") - private Duration maxPollInterval = DEFAULT_MAX_POLL_INTERVAL; - - @JsonProperty("consumer_max_poll_records") - private Integer consumerMaxPollRecords = DEFAULT_CONSUMER_MAX_POLL_RECORDS; - - @JsonProperty("heart_beat_interval") - @Valid - @Size(min = 1) - private Duration heartBeatInterval= DEFAULT_HEART_BEAT_INTERVAL_DURATION; - - @JsonProperty("is_topic_create") - private Boolean isTopicCreate =Boolean.FALSE; - - @JsonProperty("number_of_partitions") - private Integer numberOfPartions = DEFAULT_NUM_OF_PARTITIONS; - - @JsonProperty("replication_factor") - private Short replicationFactor = DEFAULT_REPLICATION_FACTOR; - - @JsonProperty("retention_period") - private Long retentionPeriod=DEFAULT_RETENTION_PERIOD; - - @JsonProperty("encryption_key") - private String encryptionKey; - - @JsonProperty("kms") - private KmsConfig kmsConfig; - - public Long getRetentionPeriod() { - return retentionPeriod; - } - - public String getGroupId() { - return groupId; - } - - public void setGroupId(String groupId) { - this.groupId = groupId; - } - - public MessageFormat getSerdeFormat() { - return serdeFormat; - } - - public String getEncryptionKey() { - return encryptionKey; - } - - public KmsConfig getKmsConfig() { - return kmsConfig; - } - - public Boolean getAutoCommit() { - return autoCommit; - } - - public Duration getCommitInterval() { - return commitInterval; - } - - public void setCommitInterval(Duration commitInterval) { - this.commitInterval = commitInterval; - } - - public Duration getSessionTimeOut() { - return sessionTimeOut; - } - - public String getAutoOffsetReset() { - return autoOffsetReset; - } - - public void setAutoOffsetReset(String autoOffsetReset) { - this.autoOffsetReset = autoOffsetReset; - } - - public Duration getThreadWaitingTime() { - return threadWaitingTime; - } - - public void setThreadWaitingTime(Duration threadWaitingTime) { - this.threadWaitingTime = threadWaitingTime; - } - - public long getMaxPartitionFetchBytes() { - return ByteCount.parse(maxPartitionFetchBytes).getBytes(); - } - - public long getFetchMaxBytes() { - long value = ByteCount.parse(fetchMaxBytes).getBytes(); - if (value < 1 || value > 50*1024*1024) { - throw new RuntimeException("Invalid Fetch Max Bytes"); - } - return value; - } - - public void setAutoCommit(Boolean autoCommit) { - this.autoCommit = autoCommit; - } - - public Integer getFetchMaxWait() { - return fetchMaxWait; - } - - public long getFetchMinBytes() { - long value = ByteCount.parse(fetchMinBytes).getBytes(); - if (value < 1) { - throw new RuntimeException("Invalid Fetch Min Bytes"); - } - return value; - } - - public Duration getRetryBackoff() { - return retryBackoff; - } - - public Duration getReconnectBackoff() { - return reconnectBackoff; - } - - public void setRetryBackoff(Duration retryBackoff) { - this.retryBackoff = retryBackoff; - } - - public Duration getMaxPollInterval() { - return maxPollInterval; - } - - public void setMaxPollInterval(Duration maxPollInterval) { - this.maxPollInterval = maxPollInterval; - } - - public Integer getConsumerMaxPollRecords() { - return consumerMaxPollRecords; - } - - public void setConsumerMaxPollRecords(Integer consumerMaxPollRecords) { - this.consumerMaxPollRecords = consumerMaxPollRecords; - } - - public Integer getWorkers() { - return workers; - } - - public void setWorkers(Integer workers) { - this.workers = workers; - } - - public Duration getHeartBeatInterval() { - return heartBeatInterval; - } - - public void setHeartBeatInterval(Duration heartBeatInterval) { - this.heartBeatInterval = heartBeatInterval; - } - - public String getName() { - return name; - } - - public void setName(String name) { - this.name = name; - } - - public KafkaKeyMode getKafkaKeyMode() { - return kafkaKeyMode; - } - - public Boolean isCreate() { - return isTopicCreate; - } - - public Integer getNumberOfPartions() { - return numberOfPartions; - } - - public Short getReplicationFactor() { - return replicationFactor; - } + Duration getReconnectBackoff(); } diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/TopicConsumerConfig.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/TopicConsumerConfig.java new file mode 100644 index 0000000000..0ae2126cbe --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/TopicConsumerConfig.java @@ -0,0 +1,44 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.configuration; + +import java.time.Duration; + +/** + * An extension of the {@link TopicConfig} specifically for + * consumers from Kafka topics. + */ +public interface TopicConsumerConfig extends TopicConfig { + KafkaKeyMode getKafkaKeyMode(); + + String getGroupId(); + + Boolean getAutoCommit(); + + String getAutoOffsetReset(); + + Duration getCommitInterval(); + + long getMaxPartitionFetchBytes(); + + long getFetchMaxBytes(); + + Integer getFetchMaxWait(); + + long getFetchMinBytes(); + + Duration getThreadWaitingTime(); + + Duration getSessionTimeOut(); + + Duration getHeartBeatInterval(); + + Duration getMaxPollInterval(); + + Integer getConsumerMaxPollRecords(); + + Integer getWorkers(); +} diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/TopicProducerConfig.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/TopicProducerConfig.java new file mode 100644 index 0000000000..bceced39f0 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/TopicProducerConfig.java @@ -0,0 +1,17 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.configuration; + +/** + * An extension of the {@link TopicConfig} specifically for + * producers to Kafka topics. + */ +public interface TopicProducerConfig extends TopicConfig { + Integer getNumberOfPartitions(); + Short getReplicationFactor(); + Long getRetentionPeriod(); + boolean isCreateTopic(); +} 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 79e50f0647..e36e93e2d9 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 @@ -24,14 +24,14 @@ 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.codec.ByteDecoder; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.event.EventMetadata; import org.opensearch.dataprepper.model.log.JacksonLog; import org.opensearch.dataprepper.model.record.Record; -import org.opensearch.dataprepper.model.codec.ByteDecoder; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaConsumerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaKeyMode; -import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicMetrics; import org.opensearch.dataprepper.plugins.kafka.util.LogRateLimiter; import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; @@ -72,7 +72,7 @@ public class KafkaCustomConsumer implements Runnable, ConsumerRebalanceListener private KafkaConsumer consumer= null; private AtomicBoolean shutdownInProgress; private final String topicName; - private final TopicConfig topicConfig; + private final TopicConsumerConfig topicConfig; private MessageFormat schema; private final BufferAccumulator> bufferAccumulator; private final Buffer> buffer; @@ -97,7 +97,7 @@ public KafkaCustomConsumer(final KafkaConsumer consumer, final AtomicBoolean shutdownInProgress, final Buffer> buffer, final KafkaConsumerConfig consumerConfig, - final TopicConfig topicConfig, + final TopicConsumerConfig topicConfig, final String schemaType, final AcknowledgementSetManager acknowledgementSetManager, final ByteDecoder byteDecoder, 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 b686dcd113..de553afc99 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 @@ -1,3 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + package org.opensearch.dataprepper.plugins.kafka.consumer; import io.confluent.kafka.schemaregistry.client.CachedSchemaRegistryClient; @@ -26,6 +31,7 @@ import org.opensearch.dataprepper.plugins.kafka.common.key.KeyFactory; import org.opensearch.dataprepper.plugins.kafka.common.serialization.SerializationFactory; import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaConsumerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.OAuthConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.PlainTextAuthConfig; @@ -62,7 +68,7 @@ public KafkaCustomConsumerFactory(SerializationFactory serializationFactory, Aws this.awsCredentialsSupplier = awsCredentialsSupplier; } - public List createConsumersForTopic(final KafkaConsumerConfig kafkaConsumerConfig, final TopicConfig topic, + public List createConsumersForTopic(final KafkaConsumerConfig kafkaConsumerConfig, final TopicConsumerConfig topic, final Buffer> buffer, final PluginMetrics pluginMetrics, final AcknowledgementSetManager acknowledgementSetManager, final ByteDecoder byteDecoder, @@ -109,7 +115,7 @@ public List createConsumersForTopic(final KafkaConsumerConf return consumers; } - private Properties getConsumerProperties(final KafkaConsumerConfig sourceConfig, final TopicConfig topicConfig, final Properties authProperties) { + private Properties getConsumerProperties(final KafkaConsumerConfig sourceConfig, final TopicConsumerConfig topicConfig, final Properties authProperties) { Properties properties = (Properties)authProperties.clone(); if (StringUtils.isNotEmpty(sourceConfig.getClientDnsLookup())) { ClientDNSLookupType dnsLookupType = ClientDNSLookupType.getDnsLookupType(sourceConfig.getClientDnsLookup()); @@ -131,7 +137,7 @@ private Properties getConsumerProperties(final KafkaConsumerConfig sourceConfig, return properties; } - private void setConsumerTopicProperties(final Properties properties, final TopicConfig topicConfig) { + private void setConsumerTopicProperties(final Properties properties, final TopicConsumerConfig topicConfig) { properties.put(ConsumerConfig.GROUP_ID_CONFIG, topicConfig.getGroupId()); properties.put(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, (int)topicConfig.getMaxPartitionFetchBytes()); properties.put(ConsumerConfig.RETRY_BACKOFF_MS_CONFIG, ((Long)topicConfig.getRetryBackoff().toMillis()).intValue()); 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 97c1794658..ccd28eba35 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 @@ -1,3 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + package org.opensearch.dataprepper.plugins.kafka.producer; import org.apache.kafka.clients.producer.KafkaProducer; @@ -14,6 +19,7 @@ import org.opensearch.dataprepper.plugins.kafka.common.key.KeyFactory; import org.opensearch.dataprepper.plugins.kafka.common.serialization.SerializationFactory; import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicProducerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.opensearch.dataprepper.plugins.kafka.consumer.KafkaCustomConsumerFactory; @@ -75,10 +81,10 @@ private void prepareTopicAndSchema(final KafkaProducerConfig kafkaProducerConfig } private void checkTopicCreationCriteriaAndCreateTopic(final KafkaProducerConfig kafkaProducerConfig) { - final TopicConfig topic = kafkaProducerConfig.getTopic(); - if (!topic.isCreate()) { + final TopicProducerConfig topic = kafkaProducerConfig.getTopic(); + if (!topic.isCreateTopic()) { final TopicService topicService = new TopicService(kafkaProducerConfig); - topicService.createTopic(kafkaProducerConfig.getTopic().getName(), topic.getNumberOfPartions(), topic.getReplicationFactor()); + topicService.createTopic(kafkaProducerConfig.getTopic().getName(), topic.getNumberOfPartitions(), topic.getReplicationFactor()); topicService.closeAdminClient(); } 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 ea273c370d..5cf46ea71a 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 @@ -18,9 +18,8 @@ import org.opensearch.dataprepper.model.sink.Sink; import org.opensearch.dataprepper.model.sink.SinkContext; import org.opensearch.dataprepper.plugins.kafka.common.serialization.SerializationFactory; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSinkConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicProducerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.opensearch.dataprepper.plugins.kafka.producer.KafkaCustomProducer; import org.opensearch.dataprepper.plugins.kafka.producer.KafkaCustomProducerFactory; import org.opensearch.dataprepper.plugins.kafka.producer.ProducerWorker; @@ -145,10 +144,10 @@ private void prepareTopicAndSchema() { } private void checkTopicCreationCriteriaAndCreateTopic() { - final TopicConfig topic = kafkaSinkConfig.getTopic(); - if (topic.isCreate()) { + final TopicProducerConfig topic = kafkaSinkConfig.getTopic(); + if (topic.isCreateTopic()) { final TopicService topicService = new TopicService(kafkaSinkConfig); - topicService.createTopic(kafkaSinkConfig.getTopic().getName(), topic.getNumberOfPartions(), topic.getReplicationFactor()); + topicService.createTopic(kafkaSinkConfig.getTopic().getName(), topic.getNumberOfPartitions(), topic.getReplicationFactor()); topicService.closeAdminClient(); } diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaSinkConfig.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkConfig.java similarity index 82% rename from data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaSinkConfig.java rename to data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkConfig.java index b4eb00c2a1..cb573003c8 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaSinkConfig.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkConfig.java @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.dataprepper.plugins.kafka.configuration; +package org.opensearch.dataprepper.plugins.kafka.sink; import com.fasterxml.jackson.annotation.JsonProperty; import jakarta.validation.Valid; @@ -13,6 +13,13 @@ import org.apache.commons.lang3.ObjectUtils; import org.opensearch.dataprepper.model.configuration.PluginModel; import org.opensearch.dataprepper.model.configuration.PluginSetting; +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.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 java.util.LinkedHashMap; import java.util.List; @@ -24,8 +31,7 @@ * * A helper class that helps to read user configuration values from * pipelines.yaml */ - -public class KafkaSinkConfig implements KafkaProducerConfig{ +public class KafkaSinkConfig implements KafkaProducerConfig { public static final String DLQ = "dlq"; @@ -61,7 +67,7 @@ public void setDlqConfig(final PluginSetting pluginSetting) { @JsonProperty("topic") - TopicConfig topic; + SinkTopicConfig topic; @JsonProperty("authentication") private AuthConfig authConfig; @@ -140,11 +146,11 @@ public void setSchemaConfig(SchemaConfig schemaConfig) { this.schemaConfig = schemaConfig; } - public TopicConfig getTopic() { + public TopicProducerConfig getTopic() { return topic; } - public void setTopic(TopicConfig topic) { + public void setTopic(SinkTopicConfig topic) { this.topic = topic; } diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/sink/SinkTopicConfig.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/sink/SinkTopicConfig.java new file mode 100644 index 0000000000..adb7ced442 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/sink/SinkTopicConfig.java @@ -0,0 +1,68 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.sink; + +import com.fasterxml.jackson.annotation.JsonProperty; +import org.opensearch.dataprepper.plugins.kafka.configuration.CommonTopicConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.KmsConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicProducerConfig; +import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; + +public class SinkTopicConfig extends CommonTopicConfig implements TopicProducerConfig { + private static final Integer DEFAULT_NUM_OF_PARTITIONS = 1; + private static final Short DEFAULT_REPLICATION_FACTOR = 1; + private static final Long DEFAULT_RETENTION_PERIOD = 604800000L; + + @JsonProperty("serde_format") + private MessageFormat serdeFormat = MessageFormat.PLAINTEXT; + + @JsonProperty("number_of_partitions") + private Integer numberOfPartitions = DEFAULT_NUM_OF_PARTITIONS; + + @JsonProperty("replication_factor") + private Short replicationFactor = DEFAULT_REPLICATION_FACTOR; + + @JsonProperty("retention_period") + private Long retentionPeriod = DEFAULT_RETENTION_PERIOD; + + @JsonProperty("create_topic") + private boolean isCreateTopic = false; + + @Override + public MessageFormat getSerdeFormat() { + return serdeFormat; + } + + @Override + public String getEncryptionKey() { + return null; + } + + @Override + public KmsConfig getKmsConfig() { + return null; + } + + @Override + public Integer getNumberOfPartitions() { + return numberOfPartitions; + } + + @Override + public Short getReplicationFactor() { + return replicationFactor; + } + + @Override + public Long getRetentionPeriod() { + return retentionPeriod; + } + + @Override + public boolean isCreateTopic() { + return isCreateTopic; + } +} 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 3321e0d2c2..76219bf951 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 @@ -30,7 +30,7 @@ import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.source.Source; import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSourceConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.OAuthConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.PlainTextAuthConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; @@ -204,7 +204,7 @@ public void stopExecutor(final ExecutorService executorService, final long shutd } private long calculateLongestThreadWaitingTime() { - List topicsList = sourceConfig.getTopics(); + List topicsList = sourceConfig.getTopics(); return topicsList.stream(). map( topics -> topics.getThreadWaitingTime().toSeconds() @@ -217,7 +217,7 @@ KafkaConsumer getConsumer() { return kafkaConsumer; } - private Properties getConsumerProperties(final TopicConfig topicConfig, final Properties authProperties) { + private Properties getConsumerProperties(final TopicConsumerConfig topicConfig, final Properties authProperties) { Properties properties = (Properties) authProperties.clone(); if (StringUtils.isNotEmpty(sourceConfig.getClientDnsLookup())) { ClientDNSLookupType dnsLookupType = ClientDNSLookupType.getDnsLookupType(sourceConfig.getClientDnsLookup()); @@ -306,7 +306,7 @@ private void setPropertiesForSchemaType(Properties properties, TopicConfig topic } } - private void setConsumerTopicProperties(Properties properties, TopicConfig topicConfig) { + private void setConsumerTopicProperties(Properties properties, TopicConsumerConfig topicConfig) { properties.put(ConsumerConfig.GROUP_ID_CONFIG, consumerGroupID); properties.put(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, (int) topicConfig.getMaxPartitionFetchBytes()); properties.put(ConsumerConfig.RETRY_BACKOFF_MS_CONFIG, ((Long) topicConfig.getRetryBackoff().toMillis()).intValue()); diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaSourceConfig.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceConfig.java similarity index 79% rename from data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaSourceConfig.java rename to data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceConfig.java index ba95e75723..a7274d3987 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaSourceConfig.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceConfig.java @@ -3,12 +3,18 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.dataprepper.plugins.kafka.configuration; +package org.opensearch.dataprepper.plugins.kafka.source; import com.fasterxml.jackson.annotation.JsonProperty; import jakarta.validation.Valid; import jakarta.validation.constraints.NotNull; import jakarta.validation.constraints.Size; +import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.AwsConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaConsumerConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; import java.util.List; import java.util.Objects; @@ -26,7 +32,7 @@ public class KafkaSourceConfig implements KafkaConsumerConfig { @JsonProperty("topics") @NotNull @Size(min = 1, max = 10, message = "The number of Topics should be between 1 and 10") - private List topics; + private List topics; @JsonProperty("schema") @Valid @@ -57,11 +63,11 @@ public boolean getAcknowledgementsEnabled() { return acknowledgementsEnabled; } - public List getTopics() { + public List getTopics() { return topics; } - public void setTopics(List topics) { + public void setTopics(List topics) { this.topics = topics; } diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/source/SourceTopicConfig.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/source/SourceTopicConfig.java new file mode 100644 index 0000000000..adcf030f1f --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/source/SourceTopicConfig.java @@ -0,0 +1,200 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.source; + +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.Valid; +import jakarta.validation.constraints.Size; +import org.opensearch.dataprepper.model.types.ByteCount; +import org.opensearch.dataprepper.plugins.kafka.configuration.CommonTopicConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaKeyMode; +import org.opensearch.dataprepper.plugins.kafka.configuration.KmsConfig; +import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; + +import java.time.Duration; + +class SourceTopicConfig extends CommonTopicConfig implements TopicConsumerConfig { + static final boolean DEFAULT_AUTO_COMMIT = false; + static final Duration DEFAULT_COMMIT_INTERVAL = Duration.ofSeconds(5); + static final String DEFAULT_FETCH_MAX_BYTES = "50mb"; + static final Integer DEFAULT_FETCH_MAX_WAIT = 500; + static final String DEFAULT_FETCH_MIN_BYTES = "1b"; + static final String DEFAULT_MAX_PARTITION_FETCH_BYTES = "1mb"; + static final Duration DEFAULT_SESSION_TIMEOUT = Duration.ofSeconds(45); + static final String DEFAULT_AUTO_OFFSET_RESET = "earliest"; + static final Duration DEFAULT_THREAD_WAITING_TIME = Duration.ofSeconds(5); + static final Duration DEFAULT_MAX_POLL_INTERVAL = Duration.ofSeconds(300); + static final Integer DEFAULT_CONSUMER_MAX_POLL_RECORDS = 500; + static final Integer DEFAULT_NUM_OF_WORKERS = 2; + static final Duration DEFAULT_HEART_BEAT_INTERVAL_DURATION = Duration.ofSeconds(5); + + + @JsonProperty("serde_format") + private MessageFormat serdeFormat = MessageFormat.PLAINTEXT; + + @JsonProperty("commit_interval") + @Valid + @Size(min = 1) + private Duration commitInterval = DEFAULT_COMMIT_INTERVAL; + + @JsonProperty("key_mode") + private KafkaKeyMode kafkaKeyMode = KafkaKeyMode.INCLUDE_AS_FIELD; + + @JsonProperty("group_id") + @Valid + @Size(min = 1, max = 255, message = "size of group id should be between 1 and 255") + private String groupId; + + @JsonProperty("workers") + @Valid + @Size(min = 1, max = 200, message = "Number of worker threads should lies between 1 and 200") + private Integer workers = DEFAULT_NUM_OF_WORKERS; + + @JsonProperty("session_timeout") + @Valid + @Size(min = 1) + private Duration sessionTimeOut = DEFAULT_SESSION_TIMEOUT; + + @JsonProperty("auto_offset_reset") + private String autoOffsetReset = DEFAULT_AUTO_OFFSET_RESET; + + @JsonProperty("thread_waiting_time") + private Duration threadWaitingTime = DEFAULT_THREAD_WAITING_TIME; + + @JsonProperty("max_poll_interval") + private Duration maxPollInterval = DEFAULT_MAX_POLL_INTERVAL; + + @JsonProperty("consumer_max_poll_records") + private Integer consumerMaxPollRecords = DEFAULT_CONSUMER_MAX_POLL_RECORDS; + + @JsonProperty("heart_beat_interval") + @Valid + @Size(min = 1) + private Duration heartBeatInterval= DEFAULT_HEART_BEAT_INTERVAL_DURATION; + + @JsonProperty("auto_commit") + private Boolean autoCommit = DEFAULT_AUTO_COMMIT; + + @JsonProperty("max_partition_fetch_bytes") + private String maxPartitionFetchBytes = DEFAULT_MAX_PARTITION_FETCH_BYTES; + + @JsonProperty("fetch_max_bytes") + private String fetchMaxBytes = DEFAULT_FETCH_MAX_BYTES; + + @JsonProperty("fetch_max_wait") + @Valid + @Size(min = 1) + private Integer fetchMaxWait = DEFAULT_FETCH_MAX_WAIT; + + @JsonProperty("fetch_min_bytes") + private String fetchMinBytes = DEFAULT_FETCH_MIN_BYTES; + + + @Override + public String getEncryptionKey() { + return null; + } + + @Override + public KmsConfig getKmsConfig() { + return null; + } + + + @Override + public Duration getCommitInterval() { + return commitInterval; + } + + + @Override + public KafkaKeyMode getKafkaKeyMode() { + return kafkaKeyMode; + } + + @Override + public String getGroupId() { + return groupId; + } + + @Override + public MessageFormat getSerdeFormat() { + return serdeFormat; + } + + @Override + public Boolean getAutoCommit() { + return autoCommit; + } + + public void setAutoCommit(Boolean autoCommit) { + this.autoCommit = autoCommit; + } + + @Override + public long getFetchMaxBytes() { + long value = ByteCount.parse(fetchMaxBytes).getBytes(); + if (value < 1 || value > 50*1024*1024) { + throw new RuntimeException("Invalid Fetch Max Bytes"); + } + return value; + } + + @Override + public Integer getFetchMaxWait() { + return fetchMaxWait; + } + + @Override + public long getFetchMinBytes() { + long value = ByteCount.parse(fetchMinBytes).getBytes(); + if (value < 1) { + throw new RuntimeException("Invalid Fetch Min Bytes"); + } + return value; + } + + @Override + public long getMaxPartitionFetchBytes() { + return ByteCount.parse(maxPartitionFetchBytes).getBytes(); + } + + @Override + public Duration getSessionTimeOut() { + return sessionTimeOut; + } + + @Override + public String getAutoOffsetReset() { + return autoOffsetReset; + } + + @Override + public Duration getThreadWaitingTime() { + return threadWaitingTime; + } + + @Override + public Duration getMaxPollInterval() { + return maxPollInterval; + } + + @Override + public Integer getConsumerMaxPollRecords() { + return consumerMaxPollRecords; + } + + @Override + public Integer getWorkers() { + return workers; + } + + @Override + public Duration getHeartBeatInterval() { + return heartBeatInterval; + } +} 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 db01c919cf..779aefcab0 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 @@ -9,7 +9,7 @@ import org.opensearch.dataprepper.plugins.kafka.configuration.AwsIamAuthConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaConsumerConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSourceConfig; +import org.opensearch.dataprepper.plugins.kafka.source.KafkaSourceConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.OAuthConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionType; diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/buffer/BufferTopicConfigTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/buffer/BufferTopicConfigTest.java new file mode 100644 index 0000000000..60187fe41d --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/buffer/BufferTopicConfigTest.java @@ -0,0 +1,56 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.buffer; + +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.model.types.ByteCount; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.opensearch.dataprepper.test.helper.ReflectivelySetField.setField; + +class BufferTopicConfigTest { + private BufferTopicConfig createObjectUnderTest() { + return new BufferTopicConfig(); + } + + @Test + void verify_default_values() { + BufferTopicConfig objectUnderTest = createObjectUnderTest(); + + assertThat(objectUnderTest.getAutoCommit(), equalTo(BufferTopicConfig.DEFAULT_AUTO_COMMIT)); + assertThat(objectUnderTest.getCommitInterval(), equalTo(BufferTopicConfig.DEFAULT_COMMIT_INTERVAL)); + assertThat(objectUnderTest.getFetchMaxWait(), equalTo((int) BufferTopicConfig.DEFAULT_FETCH_MAX_WAIT.toMillis())); + assertThat(objectUnderTest.getFetchMinBytes(), equalTo(BufferTopicConfig.DEFAULT_FETCH_MIN_BYTES.getBytes())); + assertThat(objectUnderTest.getFetchMaxBytes(), equalTo(BufferTopicConfig.DEFAULT_FETCH_MAX_BYTES.getBytes())); + assertThat(objectUnderTest.getMaxPartitionFetchBytes(), equalTo(BufferTopicConfig.DEFAULT_MAX_PARTITION_FETCH_BYTES.getBytes())); + + assertThat(objectUnderTest.getSessionTimeOut(), equalTo(BufferTopicConfig.DEFAULT_SESSION_TIMEOUT)); + assertThat(objectUnderTest.getAutoOffsetReset(), equalTo(BufferTopicConfig.DEFAULT_AUTO_OFFSET_RESET)); + assertThat(objectUnderTest.getThreadWaitingTime(), equalTo(BufferTopicConfig.DEFAULT_THREAD_WAITING_TIME)); + assertThat(objectUnderTest.getMaxPollInterval(), equalTo(BufferTopicConfig.DEFAULT_MAX_POLL_INTERVAL)); + assertThat(objectUnderTest.getConsumerMaxPollRecords(), equalTo(BufferTopicConfig.DEFAULT_CONSUMER_MAX_POLL_RECORDS)); + assertThat(objectUnderTest.getWorkers(), equalTo(BufferTopicConfig.DEFAULT_NUM_OF_WORKERS)); + assertThat(objectUnderTest.getHeartBeatInterval(), equalTo(BufferTopicConfig.DEFAULT_HEART_BEAT_INTERVAL_DURATION)); + } + + @Test + void getFetchMaxBytes_on_large_value() throws NoSuchFieldException, IllegalAccessException { + BufferTopicConfig objectUnderTest = createObjectUnderTest(); + + setField(BufferTopicConfig.class, objectUnderTest, "fetchMaxBytes", ByteCount.parse("60mb")); + assertThrows(RuntimeException.class, () -> objectUnderTest.getFetchMaxBytes()); + } + + @Test + void invalid_getFetchMaxBytes_zero_bytes() throws NoSuchFieldException, IllegalAccessException { + BufferTopicConfig objectUnderTest = createObjectUnderTest(); + + setField(BufferTopicConfig.class, objectUnderTest, "fetchMaxBytes", ByteCount.zeroBytes()); + assertThrows(RuntimeException.class, () -> objectUnderTest.getFetchMaxBytes()); + } +} \ No newline at end of file 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 6747ab4894..06c44d3599 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 @@ -1,3 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + package org.opensearch.dataprepper.plugins.kafka.buffer; import org.junit.jupiter.api.BeforeEach; @@ -22,9 +27,7 @@ import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionType; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaBufferConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.PlainTextAuthConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.opensearch.dataprepper.plugins.kafka.producer.KafkaCustomProducer; import org.opensearch.dataprepper.plugins.kafka.producer.KafkaCustomProducerFactory; import org.opensearch.dataprepper.plugins.kafka.producer.ProducerWorker; @@ -83,7 +86,7 @@ class KafkaBufferTest { private PluginFactory pluginFactory; @Mock - TopicConfig topic1; + BufferTopicConfig topic1; @Mock AuthConfig authConfig; @Mock @@ -136,7 +139,7 @@ void setUp() { pluginMetrics = mock(PluginMetrics.class); acknowledgementSetManager = mock(AcknowledgementSetManager.class); when(topic1.getName()).thenReturn("topic1"); - when(topic1.isCreate()).thenReturn(true); + when(topic1.isCreateTopic()).thenReturn(true); when(topic1.getWorkers()).thenReturn(2); when(topic1.getCommitInterval()).thenReturn(Duration.ofSeconds(1)); diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/AuthConfigTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/AuthConfigTest.java index 8e5d0546b7..ed23995a64 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/AuthConfigTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/AuthConfigTest.java @@ -1,3 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + package org.opensearch.dataprepper.plugins.kafka.configuration; import static org.opensearch.dataprepper.test.helper.ReflectivelySetField.setField; @@ -11,6 +16,7 @@ import static org.hamcrest.Matchers.equalTo; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.plugins.kafka.source.KafkaSourceConfig; import org.yaml.snakeyaml.Yaml; import java.io.FileReader; diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/CommonTopicConfigTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/CommonTopicConfigTest.java new file mode 100644 index 0000000000..ab71db191b --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/CommonTopicConfigTest.java @@ -0,0 +1,84 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.kafka.configuration; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInfo; +import org.opensearch.dataprepper.plugins.kafka.source.KafkaSourceConfig; +import org.yaml.snakeyaml.Yaml; + +import java.io.FileReader; +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 static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.notNullValue; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; + +class CommonTopicConfigTest { + private TopicConfig topicConfig; + + private static final String YAML_FILE_WITH_CONSUMER_CONFIG = "sample-pipelines.yaml"; + + private static final String YAML_FILE_WITH_MISSING_CONSUMER_CONFIG = "sample-pipelines-1.yaml"; + + @BeforeEach + void setUp(TestInfo testInfo) throws IOException { + String fileName = testInfo.getTags().stream().findFirst().orElse(""); + Yaml yaml = new Yaml(); + FileReader fileReader = new FileReader(getClass().getClassLoader().getResource(fileName).getFile()); + Object data = yaml.load(fileReader); + ObjectMapper mapper = new ObjectMapper(); + if (data instanceof Map) { + Map propertyMap = (Map) data; + Map logPipelineMap = (Map) propertyMap.get("log-pipeline"); + Map sourceMap = (Map) logPipelineMap.get("source"); + Map kafkaConfigMap = (Map) sourceMap.get("kafka"); + mapper.registerModule(new JavaTimeModule()); + String json = mapper.writeValueAsString(kafkaConfigMap); + Reader reader = new StringReader(json); + KafkaSourceConfig kafkaSourceConfig = mapper.readValue(reader, KafkaSourceConfig.class); + List topicConfigList = kafkaSourceConfig.getTopics(); + topicConfig = topicConfigList.get(0); + } + } + + @Test + @Tag(YAML_FILE_WITH_CONSUMER_CONFIG) + void test_topicsConfig_not_null() { + assertThat(topicConfig, notNullValue()); + } + + @Test + @Tag(YAML_FILE_WITH_MISSING_CONSUMER_CONFIG) + void testConfigValues_default() { + assertEquals("my-topic-2", topicConfig.getName()); + assertEquals(CommonTopicConfig.DEFAULT_RETRY_BACKOFF, topicConfig.getRetryBackoff()); + assertEquals(CommonTopicConfig.DEFAULT_RECONNECT_BACKOFF, topicConfig.getReconnectBackoff()); + } + + @Test + @Tag(YAML_FILE_WITH_CONSUMER_CONFIG) + void testConfigValues_from_yaml() { + assertEquals("my-topic-1", topicConfig.getName()); + assertEquals(Duration.ofSeconds(100), topicConfig.getRetryBackoff()); + } + + @Test + @Tag(YAML_FILE_WITH_CONSUMER_CONFIG) + void testConfigValues_from_yaml_not_null() { + assertNotNull(topicConfig.getName()); + assertNotNull(topicConfig.getRetryBackoff()); + } +} diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/OAuthConfigTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/OAuthConfigTest.java index 25228da2ba..d33b67973b 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/OAuthConfigTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/OAuthConfigTest.java @@ -1,3 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + package org.opensearch.dataprepper.plugins.kafka.configuration; import com.fasterxml.jackson.databind.ObjectMapper; @@ -5,6 +10,7 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.plugins.kafka.source.KafkaSourceConfig; import org.yaml.snakeyaml.Yaml; import java.io.FileReader; diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/PlainTextAuthConfigTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/PlainTextAuthConfigTest.java index 560f0042eb..44a5e88554 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/PlainTextAuthConfigTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/PlainTextAuthConfigTest.java @@ -1,3 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + package org.opensearch.dataprepper.plugins.kafka.configuration; import com.fasterxml.jackson.databind.ObjectMapper; @@ -5,6 +10,7 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.Mock; +import org.opensearch.dataprepper.plugins.kafka.source.KafkaSourceConfig; import org.yaml.snakeyaml.Yaml; import java.io.FileReader; diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/SchemaConfigTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/SchemaConfigTest.java index cc9c3f69dd..e02604a19f 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/SchemaConfigTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/SchemaConfigTest.java @@ -1,3 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + package org.opensearch.dataprepper.plugins.kafka.configuration; import com.fasterxml.jackson.databind.ObjectMapper; @@ -5,6 +10,7 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.Mock; +import org.opensearch.dataprepper.plugins.kafka.source.KafkaSourceConfig; import org.yaml.snakeyaml.Yaml; import java.io.FileReader; diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/TopicConfigTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/TopicConfigTest.java deleted file mode 100644 index a85ed92727..0000000000 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/TopicConfigTest.java +++ /dev/null @@ -1,142 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.dataprepper.plugins.kafka.configuration; - -import com.fasterxml.jackson.databind.ObjectMapper; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Tag; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.TestInfo; -import org.mockito.Mock; -import org.yaml.snakeyaml.Yaml; -import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; - -import java.io.FileReader; -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 static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.notNullValue; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertThrows; - -import org.opensearch.dataprepper.model.types.ByteCount; -import static org.opensearch.dataprepper.test.helper.ReflectivelySetField.setField; - -class TopicConfigTest { - - @Mock - TopicConfig topicConfig; - - private static final String YAML_FILE_WITH_CONSUMER_CONFIG = "sample-pipelines.yaml"; - - private static final String YAML_FILE_WITH_MISSING_CONSUMER_CONFIG = "sample-pipelines-1.yaml"; - - @BeforeEach - void setUp(TestInfo testInfo) throws IOException { - String fileName = testInfo.getTags().stream().findFirst().orElse(""); - topicConfig = new TopicConfig(); - Yaml yaml = new Yaml(); - FileReader fileReader = new FileReader(getClass().getClassLoader().getResource(fileName).getFile()); - Object data = yaml.load(fileReader); - ObjectMapper mapper = new ObjectMapper(); - if (data instanceof Map) { - Map propertyMap = (Map) data; - Map logPipelineMap = (Map) propertyMap.get("log-pipeline"); - Map sourceMap = (Map) logPipelineMap.get("source"); - Map kafkaConfigMap = (Map) sourceMap.get("kafka"); - mapper.registerModule(new JavaTimeModule()); - String json = mapper.writeValueAsString(kafkaConfigMap); - Reader reader = new StringReader(json); - KafkaSourceConfig kafkaSourceConfig = mapper.readValue(reader, KafkaSourceConfig.class); - List topicConfigList = kafkaSourceConfig.getTopics(); - topicConfig = topicConfigList.get(0); - } - } - - @Test - @Tag(YAML_FILE_WITH_CONSUMER_CONFIG) - void test_topicsConfig_not_null() { - assertThat(topicConfig, notNullValue()); - } - - @Test - @Tag(YAML_FILE_WITH_MISSING_CONSUMER_CONFIG) - void testConfigValues_default() { - assertEquals("my-topic-2", topicConfig.getName()); - assertEquals("my-test-group", topicConfig.getGroupId()); - assertEquals(TopicConfig.DEFAULT_AUTO_COMMIT, topicConfig.getAutoCommit()); - assertEquals(TopicConfig.DEFAULT_COMMIT_INTERVAL, topicConfig.getCommitInterval()); - assertEquals(TopicConfig.DEFAULT_SESSION_TIMEOUT, topicConfig.getSessionTimeOut()); - assertEquals(TopicConfig.DEFAULT_AUTO_OFFSET_RESET, topicConfig.getAutoOffsetReset()); - assertEquals(TopicConfig.DEFAULT_THREAD_WAITING_TIME, topicConfig.getThreadWaitingTime()); - assertEquals(ByteCount.parse(TopicConfig.DEFAULT_FETCH_MAX_BYTES).getBytes(), topicConfig.getFetchMaxBytes()); - assertEquals(TopicConfig.DEFAULT_FETCH_MAX_WAIT, topicConfig.getFetchMaxWait()); - assertEquals(ByteCount.parse(TopicConfig.DEFAULT_FETCH_MIN_BYTES).getBytes(), topicConfig.getFetchMinBytes()); - assertEquals(TopicConfig.DEFAULT_RETRY_BACKOFF, topicConfig.getRetryBackoff()); - assertEquals(TopicConfig.DEFAULT_RECONNECT_BACKOFF, topicConfig.getReconnectBackoff()); - assertEquals(TopicConfig.DEFAULT_MAX_POLL_INTERVAL, topicConfig.getMaxPollInterval()); - assertEquals(TopicConfig.DEFAULT_CONSUMER_MAX_POLL_RECORDS, topicConfig.getConsumerMaxPollRecords()); - assertEquals(TopicConfig.DEFAULT_NUM_OF_WORKERS, topicConfig.getWorkers()); - assertEquals(TopicConfig.DEFAULT_HEART_BEAT_INTERVAL_DURATION, topicConfig.getHeartBeatInterval()); - assertEquals(ByteCount.parse(TopicConfig.DEFAULT_MAX_PARTITION_FETCH_BYTES).getBytes(), topicConfig.getMaxPartitionFetchBytes()); - } - - @Test - @Tag(YAML_FILE_WITH_CONSUMER_CONFIG) - void testConfigValues_from_yaml() { - assertEquals("my-topic-1", topicConfig.getName()); - assertEquals(false, topicConfig.getAutoCommit()); - assertEquals(Duration.ofSeconds(5), topicConfig.getCommitInterval()); - assertEquals(45000, topicConfig.getSessionTimeOut().toMillis()); - assertEquals("earliest", topicConfig.getAutoOffsetReset()); - assertEquals(Duration.ofSeconds(1), topicConfig.getThreadWaitingTime()); - assertEquals(52428800L, topicConfig.getFetchMaxBytes()); - assertEquals(500L, topicConfig.getFetchMaxWait().longValue()); - assertEquals(1L, topicConfig.getFetchMinBytes()); - assertEquals(Duration.ofSeconds(100), topicConfig.getRetryBackoff()); - assertEquals(Duration.ofSeconds(300), topicConfig.getMaxPollInterval()); - assertEquals(500L, topicConfig.getConsumerMaxPollRecords().longValue()); - assertEquals(5, topicConfig.getWorkers().intValue()); - assertEquals(Duration.ofSeconds(3), topicConfig.getHeartBeatInterval()); - assertEquals(10*ByteCount.parse(TopicConfig.DEFAULT_MAX_PARTITION_FETCH_BYTES).getBytes(), topicConfig.getMaxPartitionFetchBytes()); - } - - @Test - @Tag(YAML_FILE_WITH_CONSUMER_CONFIG) - void testConfigValues_from_yaml_not_null() { - assertNotNull(topicConfig.getName()); - assertNotNull(topicConfig.getAutoCommit()); - assertNotNull(topicConfig.getCommitInterval()); - assertNotNull(topicConfig.getSessionTimeOut()); - assertNotNull(topicConfig.getAutoOffsetReset()); - assertNotNull(topicConfig.getThreadWaitingTime()); - assertNotNull(topicConfig.getFetchMaxBytes()); - assertNotNull(topicConfig.getFetchMaxWait()); - assertNotNull(topicConfig.getFetchMinBytes()); - assertNotNull(topicConfig.getRetryBackoff()); - assertNotNull(topicConfig.getMaxPollInterval()); - assertNotNull(topicConfig.getConsumerMaxPollRecords()); - assertNotNull(topicConfig.getWorkers()); - assertNotNull(topicConfig.getHeartBeatInterval()); - } - - @Test - @Tag(YAML_FILE_WITH_CONSUMER_CONFIG) - void TestInvalidConfigValues() throws NoSuchFieldException, IllegalAccessException { - setField(TopicConfig.class, topicConfig, "fetchMaxBytes", "60mb"); - assertThrows(RuntimeException.class, () -> topicConfig.getFetchMaxBytes()); - setField(TopicConfig.class, topicConfig, "fetchMaxBytes", "0b"); - assertThrows(RuntimeException.class, () -> topicConfig.getFetchMaxBytes()); - setField(TopicConfig.class, topicConfig, "fetchMinBytes", "0b"); - assertThrows(RuntimeException.class, () -> topicConfig.getFetchMinBytes()); - } - -} 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 0d443e4413..e32ce36836 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 @@ -32,9 +32,9 @@ import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.buffer.blockingbuffer.BlockingBuffer; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaConsumerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaKeyMode; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSourceConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicMetrics; import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; @@ -50,7 +50,6 @@ import java.util.concurrent.atomic.AtomicBoolean; import static org.awaitility.Awaitility.await; - import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.hasEntry; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -71,13 +70,13 @@ public class KafkaCustomConsumerTest { private Buffer> buffer; @Mock - private KafkaSourceConfig sourceConfig; + private KafkaConsumerConfig sourceConfig; private ExecutorService callbackExecutor; private AcknowledgementSetManager acknowledgementSetManager; @Mock - private TopicConfig topicConfig; + private TopicConsumerConfig topicConfig; @Mock private KafkaTopicMetrics topicMetrics; @@ -115,7 +114,7 @@ public void setUp() { counter = mock(Counter.class); posCounter = mock(Counter.class); negCounter = mock(Counter.class); - topicConfig = mock(TopicConfig.class); + topicConfig = mock(TopicConsumerConfig.class); when(topicMetrics.getNumberOfPositiveAcknowledgements()).thenReturn(posCounter); when(topicMetrics.getNumberOfNegativeAcknowledgements()).thenReturn(negCounter); when(topicMetrics.getNumberOfRecordsCommitted()).thenReturn(counter); @@ -138,7 +137,7 @@ public void setUp() { callbackExecutor = Executors.newFixedThreadPool(2); acknowledgementSetManager = new DefaultAcknowledgementSetManager(callbackExecutor, Duration.ofMillis(2000)); - sourceConfig = mock(KafkaSourceConfig.class); + sourceConfig = mock(KafkaConsumerConfig.class); buffer = getBuffer(); shutdownInProgress = new AtomicBoolean(false); when(topicConfig.getName()).thenReturn("topic1"); diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerTest.java index 5216537f9f..13f9879660 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerTest.java @@ -26,9 +26,9 @@ 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.kafka.configuration.KafkaSinkConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.TopicProducerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.opensearch.dataprepper.plugins.kafka.sink.DLQSink; import java.io.IOException; @@ -47,13 +47,12 @@ @ExtendWith(MockitoExtension.class) @MockitoSettings(strictness = Strictness.LENIENT) - public class KafkaCustomProducerTest { private KafkaCustomProducer producer; @Mock - private KafkaSinkConfig kafkaSinkConfig; + private KafkaProducerConfig kafkaSinkConfig; private Record record; @@ -70,8 +69,8 @@ public void setUp() { DefaultEventHandle defaultEventHandle = mock(DefaultEventHandle.class); event.setEventHandle(defaultEventHandle); record = new Record<>(event); - final TopicConfig topicConfig = new TopicConfig(); - topicConfig.setName("test-topic"); + final TopicProducerConfig topicConfig = mock(TopicProducerConfig.class); + when(topicConfig.getName()).thenReturn("test-topic"); when(kafkaSinkConfig.getTopic()).thenReturn(topicConfig); when(kafkaSinkConfig.getSchemaConfig()).thenReturn(mock(SchemaConfig.class)); diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/sink/DLQSinkTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/sink/DLQSinkTest.java index 09f697cfda..d01c62be11 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/sink/DLQSinkTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/sink/DLQSinkTest.java @@ -16,7 +16,6 @@ import org.opensearch.dataprepper.model.plugin.PluginFactory; import org.opensearch.dataprepper.plugins.dlq.DlqProvider; import org.opensearch.dataprepper.plugins.dlq.DlqWriter; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSinkConfig; import org.springframework.test.util.ReflectionTestUtils; import org.yaml.snakeyaml.Yaml; diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaSinkConfigTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkConfigTest.java similarity index 92% rename from data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaSinkConfigTest.java rename to data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkConfigTest.java index e3c4d6fde8..1a297b5175 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaSinkConfigTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkConfigTest.java @@ -3,11 +3,13 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.dataprepper.plugins.kafka.configuration; +package org.opensearch.dataprepper.plugins.kafka.sink; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.model.configuration.PluginSetting; +import org.opensearch.dataprepper.plugins.kafka.configuration.AuthConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; import java.util.Arrays; import java.util.HashMap; @@ -35,7 +37,7 @@ void setUp() { kafkaSinkConfig = new KafkaSinkConfig(); kafkaSinkConfig.setBootStrapServers(Arrays.asList("127.0.0.1:9093")); kafkaSinkConfig.setAuthConfig(mock(AuthConfig.class)); - kafkaSinkConfig.setTopic(mock(TopicConfig.class)); + kafkaSinkConfig.setTopic(mock(SinkTopicConfig.class)); kafkaSinkConfig.setSchemaConfig((mock(SchemaConfig.class))); kafkaSinkConfig.setThreadWaitTime(10L); // kafkaSinkConfig.setSerdeFormat("JSON"); 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 8758b13881..d61eb22df0 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 @@ -26,9 +26,7 @@ import org.opensearch.dataprepper.model.plugin.PluginFactory; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.sink.SinkContext; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSinkConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.opensearch.dataprepper.plugins.kafka.producer.ProducerWorker; import org.springframework.test.util.ReflectionTestUtils; import org.yaml.snakeyaml.Yaml; @@ -198,8 +196,8 @@ public void isReadyTest() { @Test public void doOutputTestForAutoTopicCreate() { - TopicConfig topicConfig = mock(TopicConfig.class); - when(topicConfig.isCreate()).thenReturn(true); + SinkTopicConfig topicConfig = mock(SinkTopicConfig.class); + when(topicConfig.isCreateTopic()).thenReturn(true); SchemaConfig schemaConfig = mock(SchemaConfig.class); when(schemaConfig.isCreate()).thenReturn(true); diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaSourceConfigTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceConfigTest.java similarity index 86% rename from data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaSourceConfigTest.java rename to data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceConfigTest.java index b2634659cc..5cb86c507d 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/configuration/KafkaSourceConfigTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceConfigTest.java @@ -1,10 +1,18 @@ -package org.opensearch.dataprepper.plugins.kafka.configuration; +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.source; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; +import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.Mock; +import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; +import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionType; import org.yaml.snakeyaml.Yaml; import java.io.FileReader; @@ -27,7 +35,7 @@ class KafkaSourceConfigTest { @Mock - KafkaSourceConfig kafkaSourceConfig; + KafkaSourceConfig kafkaSourceConfig; List bootstrapServers; @@ -73,14 +81,14 @@ void test_setters() throws NoSuchFieldException, IllegalAccessException { kafkaSourceConfig = new KafkaSourceConfig(); EncryptionConfig encryptionConfig = kafkaSourceConfig.getEncryptionConfig(); kafkaSourceConfig.setBootStrapServers(new ArrayList<>(Arrays.asList("127.0.0.1:9092"))); - TopicConfig topicConfig = mock(TopicConfig.class); + SourceTopicConfig topicConfig = mock(SourceTopicConfig.class); kafkaSourceConfig.setTopics(Collections.singletonList(topicConfig)); assertEquals(Collections.singletonList("127.0.0.1:9092"), kafkaSourceConfig.getBootstrapServers()); assertEquals(Collections.singletonList(topicConfig), kafkaSourceConfig.getTopics()); setField(KafkaSourceConfig.class, kafkaSourceConfig, "acknowledgementsEnabled", true); assertEquals(true, kafkaSourceConfig.getAcknowledgementsEnabled()); - assertEquals(EncryptionType.SSL, kafkaSourceConfig.getEncryptionConfig().getType()); + Assertions.assertEquals(EncryptionType.SSL, kafkaSourceConfig.getEncryptionConfig().getType()); setField(EncryptionConfig.class, encryptionConfig, "type", EncryptionType.NONE); assertEquals(EncryptionType.NONE, encryptionConfig.getType()); setField(EncryptionConfig.class, encryptionConfig, "type", EncryptionType.SSL); diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceTest.java index d9c78ab0af..a851b7209c 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceTest.java @@ -6,48 +6,44 @@ package org.opensearch.dataprepper.plugins.kafka.source; import org.apache.kafka.common.config.ConfigException; -import org.opensearch.dataprepper.model.buffer.Buffer; -import org.opensearch.dataprepper.model.record.Record; -import org.opensearch.dataprepper.model.event.Event; +import org.junit.jupiter.api.Assertions; +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.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; 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.configuration.PipelineDescription; -import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; +import org.opensearch.dataprepper.model.event.Event; +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.TopicConsumerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSourceConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; -import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionType; +import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; import org.opensearch.dataprepper.plugins.kafka.extension.KafkaClusterConfigSupplier; import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; -import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; -import org.junit.jupiter.api.Assertions; -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.mockito.junit.jupiter.MockitoSettings; -import org.mockito.quality.Strictness; +import java.time.Duration; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.spy; import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import java.util.Arrays; -import java.util.List; -import java.util.Collections; -import java.util.Objects; -import java.time.Duration; @ExtendWith(MockitoExtension.class) @MockitoSettings(strictness = Strictness.LENIENT) @@ -73,7 +69,7 @@ class KafkaSourceTest { private PipelineDescription pipelineDescription; @Mock - TopicConfig topic1, topic2; + TopicConsumerConfig topic1, topic2; @Mock private Buffer> buffer; @@ -117,7 +113,7 @@ void setUp() throws Exception { when(topic1.getThreadWaitingTime()).thenReturn(Duration.ofSeconds(10)); when(topic2.getThreadWaitingTime()).thenReturn(Duration.ofSeconds(10)); when(sourceConfig.getBootstrapServers()).thenReturn(Collections.singletonList("http://localhost:1234")); - when(sourceConfig.getTopics()).thenReturn(Arrays.asList(topic1, topic2)); + when(sourceConfig.getTopics()).thenReturn((List) List.of(topic1, topic2)); when(sourceConfig.getSchemaConfig()).thenReturn(null); when(sourceConfig.getEncryptionConfig()).thenReturn(encryptionConfig); when(encryptionConfig.getType()).thenReturn(EncryptionType.NONE); diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/source/SourceTopicConfigTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/source/SourceTopicConfigTest.java new file mode 100644 index 0000000000..b5ad09b853 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/source/SourceTopicConfigTest.java @@ -0,0 +1,56 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.source; + +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.model.types.ByteCount; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.opensearch.dataprepper.test.helper.ReflectivelySetField.setField; + +class SourceTopicConfigTest { + private SourceTopicConfig createObjectUnderTest() { + return new SourceTopicConfig(); + } + + @Test + void verify_default_values() { + SourceTopicConfig objectUnderTest = createObjectUnderTest(); + + assertThat(objectUnderTest.getAutoCommit(), equalTo(SourceTopicConfig.DEFAULT_AUTO_COMMIT)); + assertThat(objectUnderTest.getCommitInterval(), equalTo(SourceTopicConfig.DEFAULT_COMMIT_INTERVAL)); + assertThat(objectUnderTest.getFetchMaxWait(), equalTo(SourceTopicConfig.DEFAULT_FETCH_MAX_WAIT)); + assertThat(objectUnderTest.getFetchMinBytes(), equalTo(ByteCount.parse(SourceTopicConfig.DEFAULT_FETCH_MIN_BYTES).getBytes())); + assertThat(objectUnderTest.getFetchMaxBytes(), equalTo(ByteCount.parse(SourceTopicConfig.DEFAULT_FETCH_MAX_BYTES).getBytes())); + assertThat(objectUnderTest.getMaxPartitionFetchBytes(), equalTo(ByteCount.parse(SourceTopicConfig.DEFAULT_MAX_PARTITION_FETCH_BYTES).getBytes())); + + assertThat(objectUnderTest.getSessionTimeOut(), equalTo(SourceTopicConfig.DEFAULT_SESSION_TIMEOUT)); + assertThat(objectUnderTest.getAutoOffsetReset(), equalTo(SourceTopicConfig.DEFAULT_AUTO_OFFSET_RESET)); + assertThat(objectUnderTest.getThreadWaitingTime(), equalTo(SourceTopicConfig.DEFAULT_THREAD_WAITING_TIME)); + assertThat(objectUnderTest.getMaxPollInterval(), equalTo(SourceTopicConfig.DEFAULT_MAX_POLL_INTERVAL)); + assertThat(objectUnderTest.getConsumerMaxPollRecords(), equalTo(SourceTopicConfig.DEFAULT_CONSUMER_MAX_POLL_RECORDS)); + assertThat(objectUnderTest.getWorkers(), equalTo(SourceTopicConfig.DEFAULT_NUM_OF_WORKERS)); + assertThat(objectUnderTest.getHeartBeatInterval(), equalTo(SourceTopicConfig.DEFAULT_HEART_BEAT_INTERVAL_DURATION)); + } + + @Test + void getFetchMaxBytes_on_large_value() throws NoSuchFieldException, IllegalAccessException { + SourceTopicConfig objectUnderTest = createObjectUnderTest(); + + setField(SourceTopicConfig.class, objectUnderTest, "fetchMaxBytes", "60mb"); + assertThrows(RuntimeException.class, () -> objectUnderTest.getFetchMaxBytes()); + } + + @Test + void invalid_getFetchMaxBytes_zero_bytes() throws NoSuchFieldException, IllegalAccessException { + SourceTopicConfig objectUnderTest = createObjectUnderTest(); + + setField(SourceTopicConfig.class, objectUnderTest, "fetchMaxBytes", "0b"); + assertThrows(RuntimeException.class, () -> objectUnderTest.getFetchMaxBytes()); + } +} \ No newline at end of file diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/AuthenticationPropertyConfigurerTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/AuthenticationPropertyConfigurerTest.java index 0656fe96ee..ac6f78bb87 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/AuthenticationPropertyConfigurerTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/AuthenticationPropertyConfigurerTest.java @@ -11,7 +11,7 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSinkConfig; +import org.opensearch.dataprepper.plugins.kafka.sink.KafkaSinkConfig; import org.yaml.snakeyaml.Yaml; import java.io.FileReader; diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/SinkPropertyConfigurerTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/SinkPropertyConfigurerTest.java index 4319085f97..36f32736c5 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/SinkPropertyConfigurerTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/SinkPropertyConfigurerTest.java @@ -1,3 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + package org.opensearch.dataprepper.plugins.kafka.util; import com.fasterxml.jackson.databind.ObjectMapper; @@ -9,8 +14,8 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.MockedStatic; import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaSinkConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; +import org.opensearch.dataprepper.plugins.kafka.sink.KafkaSinkConfig; import org.springframework.test.util.ReflectionTestUtils; import org.yaml.snakeyaml.Yaml; @@ -25,8 +30,6 @@ @ExtendWith(MockitoExtension.class) public class SinkPropertyConfigurerTest { - - KafkaSinkConfig kafkaSinkConfig; MockedStatic authenticationPropertyConfigurerMockedStatic; diff --git a/data-prepper-plugins/kafka-plugins/src/test/resources/sample-pipelines-sink.yaml b/data-prepper-plugins/kafka-plugins/src/test/resources/sample-pipelines-sink.yaml index e417522854..47e5c7d3d2 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/resources/sample-pipelines-sink.yaml +++ b/data-prepper-plugins/kafka-plugins/src/test/resources/sample-pipelines-sink.yaml @@ -15,7 +15,7 @@ log-pipeline : serde_format: plaintext topic: name: plaintext_test_topic - is_topic_create: false + create_topic: false producer_properties: buffer_memory: 10mb compression_type: gzip From c560e1f85a431d8c4fc7dddf425504f385861b5d Mon Sep 17 00:00:00 2001 From: Aiden Dai <68811299+daixba@users.noreply.github.com> Date: Wed, 1 Nov 2023 03:48:32 +0800 Subject: [PATCH 16/53] Add bug fixes and improvements to DDB source (#3559) Signed-off-by: Aiden Dai --- .../dynamodb-source/build.gradle | 1 + .../dynamodb/export/DataFileLoader.java | 60 ++++++--- .../export/DataFileLoaderFactory.java | 13 +- .../dynamodb/export/DataFileScheduler.java | 15 ++- .../source/dynamodb/model/ExportSummary.java | 4 +- .../source/dynamodb/model/LoadStatus.java | 20 +-- .../source/dynamodb/stream/ShardConsumer.java | 54 ++++++-- .../dynamodb/stream/ShardConsumerFactory.java | 12 +- .../dynamodb/stream/StreamCheckpointer.java | 5 +- .../dynamodb/export/DataFileLoaderTest.java | 117 ++++++++++++------ .../export/ManifestFileReaderTest.java | 2 +- .../dynamodb/stream/ShardConsumerTest.java | 72 ++++++++--- .../dynamodb/stream/StreamSchedulerTest.java | 5 +- 13 files changed, 255 insertions(+), 125 deletions(-) diff --git a/data-prepper-plugins/dynamodb-source/build.gradle b/data-prepper-plugins/dynamodb-source/build.gradle index d3a57de0e6..1fc14af6ea 100644 --- a/data-prepper-plugins/dynamodb-source/build.gradle +++ b/data-prepper-plugins/dynamodb-source/build.gradle @@ -31,6 +31,7 @@ dependencies { testImplementation platform('org.junit:junit-bom:5.9.1') testImplementation 'org.junit.jupiter:junit-jupiter' + testImplementation testLibs.mockito.inline testImplementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml' } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoader.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoader.java index 6d1d0aedd9..fcede4dba9 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoader.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoader.java @@ -5,12 +5,21 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.export; +import org.opensearch.dataprepper.buffer.common.BufferAccumulator; +import org.opensearch.dataprepper.metrics.PluginMetrics; +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.plugins.source.dynamodb.converter.ExportRecordConverter; +import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.BufferedReader; +import java.io.IOException; +import java.io.InputStream; import java.io.InputStreamReader; +import java.time.Duration; import java.util.ArrayList; import java.util.List; import java.util.zip.GZIPInputStream; @@ -37,30 +46,37 @@ public class DataFileLoader implements Runnable { */ private static final int DEFAULT_CHECKPOINT_INTERVAL_MILLS = 2 * 60_000; + static final Duration BUFFER_TIMEOUT = Duration.ofSeconds(60); + static final int DEFAULT_BUFFER_BATCH_SIZE = 1_000; + private final String bucketName; private final String key; private final ExportRecordConverter recordConverter; - private final S3ObjectReader s3ObjectReader; + private final S3ObjectReader objectReader; private final DataFileCheckpointer checkpointer; - // Start Line is the checkpoint + /** + * Start Line is the checkpoint + */ private final int startLine; private DataFileLoader(Builder builder) { - this.s3ObjectReader = builder.s3ObjectReader; - this.recordConverter = builder.recordConverter; + this.objectReader = builder.objectReader; this.bucketName = builder.bucketName; this.key = builder.key; this.checkpointer = builder.checkpointer; this.startLine = builder.startLine; + + final BufferAccumulator> bufferAccumulator = BufferAccumulator.create(builder.buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT); + recordConverter = new ExportRecordConverter(bufferAccumulator, builder.tableInfo, builder.pluginMetrics); } - public static Builder builder() { - return new Builder(); + public static Builder builder(final S3ObjectReader s3ObjectReader, final PluginMetrics pluginMetrics, final Buffer> buffer) { + return new Builder(s3ObjectReader, pluginMetrics, buffer); } @@ -69,9 +85,14 @@ public static Builder builder() { */ static class Builder { - private S3ObjectReader s3ObjectReader; + private final S3ObjectReader objectReader; + + private final PluginMetrics pluginMetrics; + + private final Buffer> buffer; + + private TableInfo tableInfo; - private ExportRecordConverter recordConverter; private DataFileCheckpointer checkpointer; @@ -81,13 +102,14 @@ static class Builder { private int startLine; - public Builder s3ObjectReader(S3ObjectReader s3ObjectReader) { - this.s3ObjectReader = s3ObjectReader; - return this; + public Builder(final S3ObjectReader objectReader, final PluginMetrics pluginMetrics, final Buffer> buffer) { + this.objectReader = objectReader; + this.pluginMetrics = pluginMetrics; + this.buffer = buffer; } - public Builder recordConverter(ExportRecordConverter recordConverter) { - this.recordConverter = recordConverter; + public Builder tableInfo(TableInfo tableInfo) { + this.tableInfo = tableInfo; return this; } @@ -128,7 +150,9 @@ public void run() { int lineCount = 0; int lastLineProcessed = 0; - try (GZIPInputStream gzipInputStream = new GZIPInputStream(s3ObjectReader.readFile(bucketName, key))) { + try { + InputStream inputStream = objectReader.readFile(bucketName, key); + GZIPInputStream gzipInputStream = new GZIPInputStream(inputStream); BufferedReader reader = new BufferedReader(new InputStreamReader(gzipInputStream)); String line; @@ -170,11 +194,15 @@ public void run() { } lines.clear(); + reader.close(); + gzipInputStream.close(); + inputStream.close(); LOG.info("Complete loading s3://{}/{}", bucketName, key); - } catch (Exception e) { + } catch (IOException e) { checkpointer.checkpoint(lineCount); - String errorMessage = String.format("Loading of s3://{}/{} completed with Exception: {}", bucketName, key, e.getMessage()); + + String errorMessage = String.format("Loading of s3://%s/%s completed with Exception: %S", bucketName, key, e.getMessage()); throw new RuntimeException(errorMessage); } } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactory.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactory.java index 9e2378e3df..c8e5318f19 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactory.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactory.java @@ -5,25 +5,19 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.export; -import org.opensearch.dataprepper.buffer.common.BufferAccumulator; import org.opensearch.dataprepper.metrics.PluginMetrics; 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.enhanced.EnhancedSourceCoordinator; -import org.opensearch.dataprepper.plugins.source.dynamodb.converter.ExportRecordConverter; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.DataFilePartition; import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo; import software.amazon.awssdk.services.s3.S3Client; -import java.time.Duration; - /** * Factory class for DataFileLoader thread. */ public class DataFileLoaderFactory { - static final Duration BUFFER_TIMEOUT = Duration.ofSeconds(60); - static final int DEFAULT_BUFFER_BATCH_SIZE = 1_000; private final EnhancedSourceCoordinator coordinator; @@ -39,17 +33,14 @@ public DataFileLoaderFactory(EnhancedSourceCoordinator coordinator, S3Client s3C } public Runnable createDataFileLoader(DataFilePartition dataFilePartition, TableInfo tableInfo) { - final BufferAccumulator> bufferAccumulator = BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT); - ExportRecordConverter recordProcessor = new ExportRecordConverter(bufferAccumulator, tableInfo, pluginMetrics); DataFileCheckpointer checkpointer = new DataFileCheckpointer(coordinator, dataFilePartition); // Start a data loader thread. - DataFileLoader loader = DataFileLoader.builder() - .s3ObjectReader(objectReader) + DataFileLoader loader = DataFileLoader.builder(objectReader, pluginMetrics, buffer) .bucketName(dataFilePartition.getBucket()) .key(dataFilePartition.getKey()) - .recordConverter(recordProcessor) + .tableInfo(tableInfo) .checkpointer(checkpointer) .startLine(dataFilePartition.getProgressState().get().getLoaded()) .build(); diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java index 91432f3dad..79b99bf817 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java @@ -63,8 +63,6 @@ public DataFileScheduler(EnhancedSourceCoordinator coordinator, DataFileLoaderFa this.coordinator = coordinator; this.pluginMetrics = pluginMetrics; this.loaderFactory = loaderFactory; - - executor = Executors.newFixedThreadPool(MAX_JOB_COUNT); this.exportFileSuccessCounter = pluginMetrics.counter(EXPORT_S3_OBJECTS_PROCESSED_COUNT); @@ -76,7 +74,7 @@ private void processDataFilePartition(DataFilePartition dataFilePartition) { String tableArn = getTableArn(exportArn); TableInfo tableInfo = getTableInfo(tableArn); - + Runnable loader = loaderFactory.createDataFileLoader(dataFilePartition, tableInfo); CompletableFuture runLoader = CompletableFuture.runAsync(loader, executor); runLoader.whenComplete(completeDataLoader(dataFilePartition)); @@ -166,6 +164,17 @@ private BiConsumer completeDataLoader(DataFilePartition dataFilePartition) { }; } + /** + * There is a global state with sourcePartitionKey the export Arn, + * to track the number of files are processed.
+ * Each time, load of a data file is completed, + * The state must be updated.
+ * Note that the state may be updated since multiple threads are updating the same state. + * Retry is required. + * + * @param exportArn Export Arn. + * @param loaded Number records Loaded. + */ private void updateState(String exportArn, int loaded) { String streamArn = getStreamArn(exportArn); diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/ExportSummary.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/ExportSummary.java index 881f59f605..578c5c7625 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/ExportSummary.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/ExportSummary.java @@ -57,7 +57,7 @@ public class ExportSummary { private long billedSizeBytes; @JsonProperty("itemCount") - private int itemCount; + private long itemCount; @JsonProperty("outputFormat") private String outputFormat; @@ -115,7 +115,7 @@ public long getBilledSizeBytes() { return billedSizeBytes; } - public int getItemCount() { + public long getItemCount() { return itemCount; } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/LoadStatus.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/LoadStatus.java index fd84c87e98..6493c70b51 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/LoadStatus.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/LoadStatus.java @@ -18,11 +18,11 @@ public class LoadStatus { private int loadedFiles; - private int totalRecords; + private long totalRecords; - private int loadedRecords; + private long loadedRecords; - public LoadStatus(int totalFiles, int loadedFiles, int totalRecords, int loadedRecords) { + public LoadStatus(int totalFiles, int loadedFiles, long totalRecords, long loadedRecords) { this.totalFiles = totalFiles; this.loadedFiles = loadedFiles; this.totalRecords = totalRecords; @@ -45,7 +45,7 @@ public void setLoadedFiles(int loadedFiles) { this.loadedFiles = loadedFiles; } - public int getTotalRecords() { + public long getTotalRecords() { return totalRecords; } @@ -53,11 +53,11 @@ public void setTotalRecords(int totalRecords) { this.totalRecords = totalRecords; } - public int getLoadedRecords() { + public long getLoadedRecords() { return loadedRecords; } - public void setLoadedRecords(int loadedRecords) { + public void setLoadedRecords(long loadedRecords) { this.loadedRecords = loadedRecords; } @@ -72,10 +72,10 @@ public Map toMap() { public static LoadStatus fromMap(Map map) { return new LoadStatus( - (int) map.get(TOTAL_FILES), - (int) map.get(LOADED_FILES), - (int) map.get(TOTAL_RECORDS), - (int) map.get(LOADED_RECORDS) + ((Number) map.get(TOTAL_FILES)).intValue(), + ((Number) map.get(LOADED_FILES)).intValue(), + ((Number) map.get(TOTAL_RECORDS)).longValue(), + ((Number) map.get(LOADED_RECORDS)).longValue() ); } } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java index ef47eb39ef..f80b145b8f 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java @@ -6,15 +6,21 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.stream; +import org.opensearch.dataprepper.buffer.common.BufferAccumulator; +import org.opensearch.dataprepper.metrics.PluginMetrics; +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.plugins.source.dynamodb.converter.StreamRecordConverter; +import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.awssdk.core.exception.SdkException; import software.amazon.awssdk.services.dynamodb.model.GetRecordsRequest; import software.amazon.awssdk.services.dynamodb.model.GetRecordsResponse; -import software.amazon.awssdk.services.dynamodb.model.Record; import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsClient; +import java.time.Duration; import java.time.Instant; import java.util.List; import java.util.stream.Collectors; @@ -46,11 +52,21 @@ public class ShardConsumer implements Runnable { */ private static final int DEFAULT_WAIT_FOR_EXPORT_INTERVAL_MILLS = 60_000; + + /** + * Default number of times in the wait for export to do regular checkpoint. + */ + private static final int DEFAULT_WAIT_COUNT_TO_CHECKPOINT = 5; + /** * Default regular checkpoint interval */ private static final int DEFAULT_CHECKPOINT_INTERVAL_MILLS = 2 * 60_000; + static final Duration BUFFER_TIMEOUT = Duration.ofSeconds(60); + static final int DEFAULT_BUFFER_BATCH_SIZE = 1_000; + + private final DynamoDbStreamsClient dynamoDbStreamsClient; private final StreamRecordConverter recordConverter; @@ -65,15 +81,17 @@ public class ShardConsumer implements Runnable { private ShardConsumer(Builder builder) { this.dynamoDbStreamsClient = builder.dynamoDbStreamsClient; - this.recordConverter = builder.recordConverter; this.checkpointer = builder.checkpointer; this.shardIterator = builder.shardIterator; this.startTime = builder.startTime; this.waitForExport = builder.waitForExport; + + final BufferAccumulator> bufferAccumulator = BufferAccumulator.create(builder.buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT); + recordConverter = new StreamRecordConverter(bufferAccumulator, builder.tableInfo, builder.pluginMetrics); } - public static Builder builder(DynamoDbStreamsClient streamsClient) { - return new Builder(streamsClient); + public static Builder builder(final DynamoDbStreamsClient dynamoDbStreamsClient, final PluginMetrics pluginMetrics, final Buffer> buffer) { + return new Builder(dynamoDbStreamsClient, pluginMetrics, buffer); } @@ -81,8 +99,11 @@ static class Builder { private final DynamoDbStreamsClient dynamoDbStreamsClient; + private final PluginMetrics pluginMetrics; + + private final Buffer> buffer; - private StreamRecordConverter recordConverter; + private TableInfo tableInfo; private StreamCheckpointer checkpointer; @@ -93,12 +114,14 @@ static class Builder { private boolean waitForExport; - public Builder(DynamoDbStreamsClient dynamoDbStreamsClient) { + public Builder(final DynamoDbStreamsClient dynamoDbStreamsClient, final PluginMetrics pluginMetrics, final Buffer> buffer) { this.dynamoDbStreamsClient = dynamoDbStreamsClient; + this.pluginMetrics = pluginMetrics; + this.buffer = buffer; } - public Builder recordConverter(StreamRecordConverter recordConverter) { - this.recordConverter = recordConverter; + public Builder tableInfo(TableInfo tableInfo) { + this.tableInfo = tableInfo; return this; } @@ -157,7 +180,7 @@ public void run() { .build(); - List records; + List records; GetRecordsResponse response; try { response = dynamoDbStreamsClient.getRecords(req); @@ -205,14 +228,27 @@ public void run() { checkpointer.checkpoint(sequenceNumber); throw new RuntimeException("Shard Consumer is interrupted"); } + + if (waitForExport) { + waitForExport(); + } } private void waitForExport() { LOG.debug("Start waiting for export to be done and loaded"); + int numberOfWaits = 0; while (!checkpointer.isExportDone()) { LOG.debug("Export is in progress, wait..."); try { Thread.sleep(DEFAULT_WAIT_FOR_EXPORT_INTERVAL_MILLS); + // The wait for export may take a long time + // Need to extend the timeout of the ownership in the coordination store. + // Otherwise, the lease will expire. + numberOfWaits++; + if (numberOfWaits % DEFAULT_WAIT_COUNT_TO_CHECKPOINT == 0) { + // To extend the timeout of lease + checkpointer.checkpoint(null); + } } catch (InterruptedException e) { LOG.error("Wait for export is interrupted ({})", e.getMessage()); // Directly quit the process diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java index ba311456b1..df605941e1 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java @@ -5,13 +5,11 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.stream; -import org.opensearch.dataprepper.buffer.common.BufferAccumulator; import org.opensearch.dataprepper.metrics.PluginMetrics; 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.enhanced.EnhancedSourceCoordinator; -import org.opensearch.dataprepper.plugins.source.dynamodb.converter.StreamRecordConverter; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.GlobalState; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.StreamPartition; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.state.StreamProgressState; @@ -21,7 +19,6 @@ import org.slf4j.LoggerFactory; import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsClient; -import java.time.Duration; import java.time.Instant; import java.util.Optional; @@ -33,8 +30,6 @@ public class ShardConsumerFactory { private static final int STREAM_TO_TABLE_OFFSET = "stream/".length(); - static final Duration BUFFER_TIMEOUT = Duration.ofSeconds(60); - static final int DEFAULT_BUFFER_BATCH_SIZE = 1_000; private final DynamoDbStreamsClient streamsClient; @@ -83,12 +78,9 @@ public Runnable createConsumer(StreamPartition streamPartition) { StreamCheckpointer checkpointer = new StreamCheckpointer(enhancedSourceCoordinator, streamPartition); String tableArn = getTableArn(streamPartition.getStreamArn()); TableInfo tableInfo = getTableInfo(tableArn); - final BufferAccumulator> bufferAccumulator = BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT); - StreamRecordConverter recordConverter = new StreamRecordConverter(bufferAccumulator, tableInfo, pluginMetrics); - - ShardConsumer shardConsumer = ShardConsumer.builder(streamsClient) - .recordConverter(recordConverter) + ShardConsumer shardConsumer = ShardConsumer.builder(streamsClient, pluginMetrics, buffer) + .tableInfo(tableInfo) .checkpointer(checkpointer) .shardIterator(shardIter) .startTime(startTime) diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamCheckpointer.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamCheckpointer.java index ccbfd268f3..ae807c2ca1 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamCheckpointer.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamCheckpointer.java @@ -32,15 +32,12 @@ public StreamCheckpointer(EnhancedSourceCoordinator coordinator, StreamPartition private void setSequenceNumber(String sequenceNumber) { // Must only update progress if sequence number is not empty - // A blank sequence number means the current sequence number in the progress state has not changed + // A blank sequence number means the current sequence number in the progress state has not changed, do nothing if (sequenceNumber != null && !sequenceNumber.isEmpty()) { Optional progressState = streamPartition.getProgressState(); if (progressState.isPresent()) { progressState.get().setSequenceNumber(sequenceNumber); - } else { - } - } } diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java index 3f3d15d8cd..443e68f2fb 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java @@ -5,17 +5,24 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.export; +import io.micrometer.core.instrument.Counter; 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.MockedStatic; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.buffer.common.BufferAccumulator; +import org.opensearch.dataprepper.metrics.PluginMetrics; +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.enhanced.EnhancedSourceCoordinator; import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourcePartition; -import org.opensearch.dataprepper.plugins.source.dynamodb.converter.ExportRecordConverter; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.DataFilePartition; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.state.DataFileProgressState; +import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo; +import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableMetadata; import software.amazon.awssdk.core.ResponseInputStream; import software.amazon.awssdk.http.AbortableInputStream; import software.amazon.awssdk.services.s3.S3Client; @@ -26,26 +33,24 @@ import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.charset.StandardCharsets; -import java.util.List; import java.util.Optional; import java.util.Random; -import java.util.StringJoiner; import java.util.UUID; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.Future; -import java.util.concurrent.TimeUnit; import java.util.zip.GZIPOutputStream; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.equalTo; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.BDDMockito.given; +import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.plugins.source.dynamodb.export.DataFileLoader.BUFFER_TIMEOUT; +import static org.opensearch.dataprepper.plugins.source.dynamodb.export.DataFileLoader.DEFAULT_BUFFER_BATCH_SIZE; @ExtendWith(MockitoExtension.class) -@Disabled class DataFileLoaderTest { @Mock @@ -54,19 +59,35 @@ class DataFileLoaderTest { @Mock private S3Client s3Client; + @Mock + private PluginMetrics pluginMetrics; + + @Mock + private Buffer> buffer; + + @Mock + private BufferAccumulator> bufferAccumulator; - private S3ObjectReader s3ObjectReader; @Mock - private ExportRecordConverter recordConverter; + private Counter testCounter; + + + private S3ObjectReader objectReader; private DataFileCheckpointer checkpointer; private DataFilePartition dataFilePartition; + private TableInfo tableInfo; + + private final String tableName = UUID.randomUUID().toString(); private final String tableArn = "arn:aws:dynamodb:us-west-2:123456789012:table/" + tableName; + private final String partitionKeyAttrName = "PK"; + private final String sortKeyAttrName = "SK"; + private final String manifestKey = UUID.randomUUID().toString(); private final String bucketName = UUID.randomUUID().toString(); private final String prefix = UUID.randomUUID().toString(); @@ -78,7 +99,7 @@ class DataFileLoaderTest { private final int total = random.nextInt(10); @BeforeEach - void setup() throws IOException { + void setup() throws Exception { DataFileProgressState state = new DataFileProgressState(); state.setLoaded(0); @@ -86,34 +107,50 @@ void setup() throws IOException { dataFilePartition = new DataFilePartition(exportArn, bucketName, manifestKey, Optional.of(state)); + TableMetadata metadata = TableMetadata.builder() + .exportRequired(true) + .streamRequired(true) + .partitionKeyAttributeName(partitionKeyAttrName) + .sortKeyAttributeName(sortKeyAttrName) + .build(); + tableInfo = new TableInfo(tableArn, metadata); + when(s3Client.getObject(any(GetObjectRequest.class))).thenReturn(generateGzipInputStream(total)); - s3ObjectReader = new S3ObjectReader(s3Client); + objectReader = new S3ObjectReader(s3Client); lenient().when(coordinator.createPartition(any(EnhancedSourcePartition.class))).thenReturn(true); lenient().doNothing().when(coordinator).completePartition(any(EnhancedSourcePartition.class)); lenient().doNothing().when(coordinator).saveProgressStateForPartition(any(EnhancedSourcePartition.class)); lenient().doNothing().when(coordinator).giveUpPartition(any(EnhancedSourcePartition.class)); - lenient().doNothing().when(recordConverter).writeToBuffer(any(List.class)); - + doNothing().when(bufferAccumulator).add(any(Record.class)); + doNothing().when(bufferAccumulator).flush(); checkpointer = new DataFileCheckpointer(coordinator, dataFilePartition); + given(pluginMetrics.counter(anyString())).willReturn(testCounter); + } - private ResponseInputStream generateGzipInputStream(int numberOfRecords) throws IOException { + private ResponseInputStream generateGzipInputStream(int numberOfRecords) { - StringJoiner stringJoiner = new StringJoiner("\\n"); + StringBuilder sb = new StringBuilder(); for (int i = 0; i < numberOfRecords; i++) { - stringJoiner.add(UUID.randomUUID().toString()); + final String pk = UUID.randomUUID().toString(); + final String sk = UUID.randomUUID().toString(); + String line = " $ion_1_0 {Item:{PK:\"" + pk + "\",SK:\"" + sk + "\"}}"; + sb.append(line + "\n"); } - final String data = stringJoiner.toString(); + final String data = sb.toString(); final byte[] dataBytes = data.getBytes(StandardCharsets.UTF_8); final ByteArrayOutputStream byteOut = new ByteArrayOutputStream(); - final GZIPOutputStream gzipOut = new GZIPOutputStream(byteOut); - gzipOut.write(dataBytes, 0, dataBytes.length); - gzipOut.close(); + try (final GZIPOutputStream gzipOut = new GZIPOutputStream(byteOut)) { + gzipOut.write(dataBytes, 0, dataBytes.length); + } catch (IOException e) { + e.printStackTrace(); + } + final byte[] bites = byteOut.toByteArray(); final ByteArrayInputStream byteInStream = new ByteArrayInputStream(bites); @@ -126,28 +163,28 @@ private ResponseInputStream generateGzipInputStream(int numbe } @Test - void test_run_loadFile_correctly() throws InterruptedException { - - DataFileLoader loader = DataFileLoader.builder() - .bucketName(bucketName) - .key(manifestKey) - .s3ObjectReader(s3ObjectReader) - .recordConverter(recordConverter) - .checkpointer(checkpointer) - .build(); + void test_run_loadFile_correctly() throws Exception { + DataFileLoader loader; + try ( + final MockedStatic bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class) + ) { + bufferAccumulatorMockedStatic.when(() -> BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT)).thenReturn(bufferAccumulator); + loader = DataFileLoader.builder(objectReader, pluginMetrics, buffer) + .bucketName(bucketName) + .key(manifestKey) + .checkpointer(checkpointer) + .tableInfo(tableInfo) + .build(); + } - ExecutorService executorService = Executors.newSingleThreadExecutor(); - final Future future = executorService.submit(loader); - Thread.sleep(100); - executorService.shutdown(); - future.cancel(true); - assertThat(executorService.awaitTermination(1000, TimeUnit.MILLISECONDS), equalTo(true)); + loader.run(); // Should call s3 getObject verify(s3Client).getObject(any(GetObjectRequest.class)); // Should write to buffer - verify(recordConverter).writeToBuffer(any(List.class)); + verify(bufferAccumulator, times(total)).add(any(Record.class)); + verify(bufferAccumulator).flush(); // Should do one last checkpoint when done. verify(coordinator).saveProgressStateForPartition(any(DataFilePartition.class)); diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ManifestFileReaderTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ManifestFileReaderTest.java index bfd3f2369d..4b93e2ef73 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ManifestFileReaderTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ManifestFileReaderTest.java @@ -50,7 +50,7 @@ void parseSummaryFile() { final String manifestFilesS3Key = UUID.randomUUID().toString(); final String outputFormat = "DYNAMODB_JSON"; long billedSizeBytes = random.nextLong(); - int itemCount = random.nextInt(10000); + long itemCount = random.nextLong(); String summaryData = String.format("{\"version\":\"%s\",\"exportArn\": \"%s\",\"startTime\":\"%s\",\"endTime\":\"%s\",\"tableArn\":\"%s\",\"tableId\":\"%s\",\"exportTime\":\"%s\",\"s3Bucket\":\"%s\",\"s3Prefix\":\"%s\",\"s3SseAlgorithm\":\"%s\",\"s3SseKmsKeyId\":null,\"manifestFilesS3Key\":\"%s\",\"billedSizeBytes\":%d,\"itemCount\":%d,\"outputFormat\":\"%s\"}", diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerTest.java index 87bc68ce7d..417787becc 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerTest.java @@ -5,17 +5,23 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.stream; +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.MockedStatic; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.buffer.common.BufferAccumulator; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.buffer.Buffer; +import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator; import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourcePartition; -import org.opensearch.dataprepper.plugins.source.dynamodb.converter.StreamRecordConverter; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.GlobalState; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.StreamPartition; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.state.StreamProgressState; +import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo; import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableMetadata; import software.amazon.awssdk.services.dynamodb.model.AttributeValue; import software.amazon.awssdk.services.dynamodb.model.GetRecordsRequest; @@ -30,12 +36,20 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Random; import java.util.UUID; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.BDDMockito.given; +import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.plugins.source.dynamodb.stream.ShardConsumer.BUFFER_TIMEOUT; +import static org.opensearch.dataprepper.plugins.source.dynamodb.stream.ShardConsumer.DEFAULT_BUFFER_BATCH_SIZE; @ExtendWith(MockitoExtension.class) class ShardConsumerTest { @@ -46,15 +60,27 @@ class ShardConsumerTest { private DynamoDbStreamsClient dynamoDbStreamsClient; @Mock - private StreamRecordConverter recordConverter; + private PluginMetrics pluginMetrics; + + @Mock + private Buffer> buffer; + + @Mock + private BufferAccumulator> bufferAccumulator; @Mock private GlobalState tableInfoGlobalState; + @Mock + private Counter testCounter; + + private StreamCheckpointer checkpointer; private StreamPartition streamPartition; + private TableInfo tableInfo; + private final String tableName = UUID.randomUUID().toString(); private final String tableArn = "arn:aws:dynamodb:us-west-2:123456789012:table/" + tableName; @@ -62,16 +88,19 @@ class ShardConsumerTest { private final String partitionKeyAttrName = "PK"; private final String sortKeyAttrName = "SK"; - private final String exportArn = tableArn + "/export/01693291918297-bfeccbea"; private final String streamArn = tableArn + "/stream/2023-09-14T05:46:45.367"; private final String shardId = "shardId-" + UUID.randomUUID(); private final String shardIterator = UUID.randomUUID().toString(); + private final Random random = new Random(); + + private final int total = random.nextInt(10); + @BeforeEach - void setup() { + void setup() throws Exception { StreamProgressState state = new StreamProgressState(); state.setWaitForExport(false); @@ -87,6 +116,7 @@ void setup() { .partitionKeyAttributeName(partitionKeyAttrName) .sortKeyAttributeName(sortKeyAttrName) .build(); + tableInfo = new TableInfo(tableArn, metadata); lenient().when(tableInfoGlobalState.getProgressState()).thenReturn(Optional.of(metadata.toMap())); lenient().when(coordinator.createPartition(any(EnhancedSourcePartition.class))).thenReturn(true); @@ -94,34 +124,46 @@ void setup() { lenient().doNothing().when(coordinator).saveProgressStateForPartition(any(EnhancedSourcePartition.class)); lenient().doNothing().when(coordinator).giveUpPartition(any(EnhancedSourcePartition.class)); + doNothing().when(bufferAccumulator).add(any(org.opensearch.dataprepper.model.record.Record.class)); + doNothing().when(bufferAccumulator).flush(); + checkpointer = new StreamCheckpointer(coordinator, streamPartition); - List records = buildRecords(10); + List records = buildRecords(total); GetRecordsResponse response = GetRecordsResponse.builder() .records(records) .nextShardIterator(null) .build(); when(dynamoDbStreamsClient.getRecords(any(GetRecordsRequest.class))).thenReturn(response); + + given(pluginMetrics.counter(anyString())).willReturn(testCounter); } @Test - void test_run_shardConsumer_correctly() { - - ShardConsumer shardConsumer = ShardConsumer.builder(dynamoDbStreamsClient) - .shardIterator(shardIterator) - .checkpointer(checkpointer) - .recordConverter(recordConverter) - .startTime(null) - .waitForExport(false) - .build(); + void test_run_shardConsumer_correctly() throws Exception { + ShardConsumer shardConsumer; + try ( + final MockedStatic bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class) + ) { + bufferAccumulatorMockedStatic.when(() -> BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT)).thenReturn(bufferAccumulator); + shardConsumer = ShardConsumer.builder(dynamoDbStreamsClient, pluginMetrics, buffer) + .shardIterator(shardIterator) + .checkpointer(checkpointer) + .tableInfo(tableInfo) + .startTime(null) + .waitForExport(false) + .build(); + } + shardConsumer.run(); // Should call GetRecords verify(dynamoDbStreamsClient).getRecords(any(GetRecordsRequest.class)); // Should write to buffer - verify(recordConverter).writeToBuffer(any(List.class)); + verify(bufferAccumulator, times(total)).add(any(org.opensearch.dataprepper.model.record.Record.class)); + verify(bufferAccumulator).flush(); // Should complete the consumer as reach to end of shard verify(coordinator).saveProgressStateForPartition(any(StreamPartition.class)); diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamSchedulerTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamSchedulerTest.java index 586a3c188a..a551a6052f 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamSchedulerTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamSchedulerTest.java @@ -6,7 +6,6 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.stream; 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; @@ -103,7 +102,6 @@ void setup() { @Test - @Disabled public void test_normal_run() throws InterruptedException { given(coordinator.acquireAvailablePartition(StreamPartition.PARTITION_TYPE)).willReturn(Optional.of(streamPartition)).willReturn(Optional.empty()); @@ -112,10 +110,9 @@ public void test_normal_run() throws InterruptedException { ExecutorService executorService = Executors.newSingleThreadExecutor(); final Future future = executorService.submit(() -> scheduler.run()); - Thread.sleep(100); + Thread.sleep(2000); executorService.shutdown(); future.cancel(true); - assertThat(executorService.awaitTermination(1000, TimeUnit.MILLISECONDS), equalTo(true)); // Should acquire the stream partition verify(coordinator).acquireAvailablePartition(StreamPartition.PARTITION_TYPE); From a30fdb5bbf5d0e34b301ee1591ea6e56a375d651 Mon Sep 17 00:00:00 2001 From: kkondaka <41027584+kkondaka@users.noreply.github.com> Date: Tue, 31 Oct 2023 13:23:55 -0700 Subject: [PATCH 17/53] OpenSearch Sink: Add log messages when there is no exception (#3532) Add log messages when there is exception Signed-off-by: Kondaka --- .../sink/opensearch/BulkRetryStrategy.java | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java index b075fe7730..451d153bfa 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java @@ -238,6 +238,13 @@ private BulkOperationRequestResponse handleRetriesAndFailures(final Accumulating if (doRetry) { if (retryCount % 5 == 0) { LOG.warn("Bulk Operation Failed. Number of retries {}. Retrying... ", retryCount, e); + if (e == null) { + for (final BulkResponseItem bulkItemResponse : bulkResponse.items()) { + if (Objects.nonNull(bulkItemResponse.error())) { + LOG.warn("operation = {}, error = {}", bulkItemResponse.operationType(), bulkItemResponse.error()); + } + } + } } bulkRequestNumberOfRetries.increment(); return new BulkOperationRequestResponse(bulkRequestForRetry, bulkResponse); @@ -248,7 +255,13 @@ private BulkOperationRequestResponse handleRetriesAndFailures(final Accumulating } private void handleFailures(final AccumulatingBulkRequest bulkRequest, final BulkResponse bulkResponse, final Throwable failure) { - if (Objects.isNull(failure)) { + LOG.warn("Bulk Operation Failed.", failure); + if (failure == null) { + for (final BulkResponseItem bulkItemResponse : bulkResponse.items()) { + if (Objects.nonNull(bulkItemResponse.error())) { + LOG.warn("operation = {}, error = {}", bulkItemResponse.operationType(), bulkItemResponse.error()); + } + } handleFailures(bulkRequest, bulkResponse.items()); } else { handleFailures(bulkRequest, failure); From 24aa5dae3b106a8094af536dde723df1a8f9870e Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 1 Nov 2023 10:01:53 -0700 Subject: [PATCH 18/53] Bump org.apache.logging.log4j:log4j-bom in /data-prepper-core (#3573) Bumps [org.apache.logging.log4j:log4j-bom](https://github.com/apache/logging-log4j2) from 2.20.0 to 2.21.1. - [Release notes](https://github.com/apache/logging-log4j2/releases) - [Commits](https://github.com/apache/logging-log4j2/compare/rel/2.20.0...rel/2.21.1) --- updated-dependencies: - dependency-name: org.apache.logging.log4j:log4j-bom dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- data-prepper-core/build.gradle | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/data-prepper-core/build.gradle b/data-prepper-core/build.gradle index 72d9c47c54..4c7aecf32f 100644 --- a/data-prepper-core/build.gradle +++ b/data-prepper-core/build.gradle @@ -27,7 +27,7 @@ dependencies { implementation 'javax.ws.rs:javax.ws.rs-api:2.1.1' implementation 'software.amazon.awssdk:cloudwatch' implementation 'org.hibernate.validator:hibernate-validator:8.0.1.Final' - implementation platform('org.apache.logging.log4j:log4j-bom:2.20.0') + implementation platform('org.apache.logging.log4j:log4j-bom:2.21.1') implementation 'org.apache.logging.log4j:log4j-core' implementation 'org.apache.logging.log4j:log4j-slf4j2-impl' implementation 'javax.inject:javax.inject:1' From 9748acd821f28e699f65d85707892c9b5ef949b9 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 1 Nov 2023 10:02:27 -0700 Subject: [PATCH 19/53] Bump org.apache.logging.log4j:log4j-jpl in /data-prepper-core (#3574) Bumps org.apache.logging.log4j:log4j-jpl from 2.20.0 to 2.21.1. --- updated-dependencies: - dependency-name: org.apache.logging.log4j:log4j-jpl dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- data-prepper-core/build.gradle | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/data-prepper-core/build.gradle b/data-prepper-core/build.gradle index 4c7aecf32f..1400d72ad2 100644 --- a/data-prepper-core/build.gradle +++ b/data-prepper-core/build.gradle @@ -39,7 +39,7 @@ dependencies { exclude group: 'commons-logging', module: 'commons-logging' } implementation 'software.amazon.cloudwatchlogs:aws-embedded-metrics:2.0.0-beta-1' - testImplementation 'org.apache.logging.log4j:log4j-jpl:2.20.0' + testImplementation 'org.apache.logging.log4j:log4j-jpl:2.21.1' testImplementation testLibs.spring.test implementation libs.armeria.core implementation libs.armeria.grpc From d8682290ae31e43922647930800b8f18f51c5886 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 1 Nov 2023 11:46:48 -0700 Subject: [PATCH 20/53] Bump org.apache.logging.log4j:log4j-bom in /data-prepper-expression (#3569) Bumps [org.apache.logging.log4j:log4j-bom](https://github.com/apache/logging-log4j2) from 2.20.0 to 2.21.1. - [Release notes](https://github.com/apache/logging-log4j2/releases) - [Commits](https://github.com/apache/logging-log4j2/compare/rel/2.20.0...rel/2.21.1) --- updated-dependencies: - dependency-name: org.apache.logging.log4j:log4j-bom dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- data-prepper-expression/build.gradle | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/data-prepper-expression/build.gradle b/data-prepper-expression/build.gradle index 41a5a44263..70d3c1d78d 100644 --- a/data-prepper-expression/build.gradle +++ b/data-prepper-expression/build.gradle @@ -24,7 +24,7 @@ dependencies { implementation(libs.spring.context) { exclude group: 'commons-logging', module: 'commons-logging' } - implementation platform('org.apache.logging.log4j:log4j-bom:2.20.0') + implementation platform('org.apache.logging.log4j:log4j-bom:2.21.1') implementation 'org.apache.logging.log4j:log4j-core' implementation 'org.apache.logging.log4j:log4j-slf4j2-impl' implementation 'com.github.seancfoley:ipaddress:5.4.0' From 7869eb7aac41965adc9a3eab4f85d4121f7b41b9 Mon Sep 17 00:00:00 2001 From: kkondaka <41027584+kkondaka@users.noreply.github.com> Date: Wed, 1 Nov 2023 13:52:58 -0700 Subject: [PATCH 21/53] Modify EventHandle to be created for every event and support internal and external origination times (#3546) * Modify EventHandle to be created for every event and support internal and external origination times Signed-off-by: Krishna Kondaka * Fixed build failures Signed-off-by: Krishna Kondaka * Fixed build failures Signed-off-by: Krishna Kondaka * fixed failing checkstyle error Signed-off-by: Krishna Kondaka * Fixed build errors Signed-off-by: Krishna Kondaka * Addressed review comments by adding InternalEventHandle Signed-off-by: Krishna Kondaka * Fixed checkstyle errors Signed-off-by: Krishna Kondaka * Fixed build errors Signed-off-by: Krishna Kondaka --------- Signed-off-by: Krishna Kondaka Co-authored-by: Krishna Kondaka --- .../model/event/DefaultEventHandle.java | 59 ++++++++++ .../model/event/DefaultEventMetadata.java | 14 +++ .../dataprepper/model/event/EventHandle.java | 28 +++++ .../model/event/EventMetadata.java | 14 +++ .../model/event/InternalEventHandle.java | 28 +++++ .../dataprepper/model/event/JacksonEvent.java | 6 +- .../model/event/DefaultEventHandleTests.java | 59 ++++++++++ .../model/event/DefaultEventMetadataTest.java | 11 ++ .../model/event/JacksonEventTest.java | 27 ++--- .../AcknowledgementSetMonitor.java | 10 +- .../DefaultAcknowledgementSet.java | 12 +- .../dataprepper/event/DefaultEventHandle.java | 29 ----- .../dataprepper/pipeline/ProcessWorker.java | 16 ++- .../router/RouterCopyRecordStrategy.java | 6 +- .../AcknowledgementSetMonitorTests.java | 2 +- ...DefaultAcknowledgementSetManagerTests.java | 28 ++--- .../DefaultAcknowledgementSetTests.java | 38 ++++-- .../event/DefaultEventHandleTests.java | 32 ----- .../dataprepper/pipeline/PipelineTests.java | 14 ++- .../router/RouterCopyRecordStrategyTests.java | 33 +++--- .../client/CloudWatchLogsService.java | 4 +- .../client/CloudWatchLogsServiceTest.java | 9 -- .../plugins/processor/GeoIPProcessorTest.java | 4 +- .../plugins/sink/http/HttpSinkServiceIT.java | 2 - .../sink/http/service/HttpSinkService.java | 6 +- .../kafka/producer/KafkaCustomProducer.java | 4 +- .../producer/KafkaCustomProducerTest.java | 3 - .../sink/opensearch/OpenSearchSinkIT.java | 15 --- .../dataprepper/plugins/sink/s3/S3SinkIT.java | 2 - .../plugins/sink/s3/S3SinkServiceIT.java | 5 +- .../plugins/sink/s3/S3SinkService.java | 6 +- .../plugins/sink/s3/S3SinkServiceTest.java | 109 ++++++++++++------ .../plugins/sink/sns/SnsSinkServiceIT.java | 4 +- .../plugins/sink/sns/SnsSinkService.java | 4 +- 34 files changed, 399 insertions(+), 244 deletions(-) create mode 100644 data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java create mode 100644 data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/InternalEventHandle.java create mode 100644 data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventHandleTests.java delete mode 100644 data-prepper-core/src/main/java/org/opensearch/dataprepper/event/DefaultEventHandle.java delete mode 100644 data-prepper-core/src/test/java/org/opensearch/dataprepper/event/DefaultEventHandleTests.java diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java new file mode 100644 index 0000000000..86278bea74 --- /dev/null +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java @@ -0,0 +1,59 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.model.event; + +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; +import java.lang.ref.WeakReference; + +import java.time.Instant; +import java.io.Serializable; + +public class DefaultEventHandle implements EventHandle, InternalEventHandle, Serializable { + private Instant externalOriginationTime; + private final Instant internalOriginationTime; + private WeakReference acknowledgementSetRef; + + public DefaultEventHandle(final Instant internalOriginationTime) { + this.acknowledgementSetRef = null; + this.externalOriginationTime = null; + this.internalOriginationTime = internalOriginationTime; + } + + @Override + public void setAcknowledgementSet(final AcknowledgementSet acknowledgementSet) { + this.acknowledgementSetRef = new WeakReference<>(acknowledgementSet); + } + + @Override + public void setExternalOriginationTime(final Instant externalOriginationTime) { + this.externalOriginationTime = externalOriginationTime; + } + + public AcknowledgementSet getAcknowledgementSet() { + if (acknowledgementSetRef == null) { + return null; + } + return acknowledgementSetRef.get(); + } + + @Override + public Instant getInternalOriginationTime() { + return this.internalOriginationTime; + } + + @Override + public Instant getExternalOriginationTime() { + return this.externalOriginationTime; + } + + @Override + public void release(boolean result) { + AcknowledgementSet acknowledgementSet = getAcknowledgementSet(); + if (acknowledgementSet != null) { + acknowledgementSet.release(this, result); + } + } +} diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventMetadata.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventMetadata.java index e2ce55caa2..883297d567 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventMetadata.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventMetadata.java @@ -27,6 +27,8 @@ public class DefaultEventMetadata implements EventMetadata { private final Instant timeReceived; + private Instant externalOriginationTime; + private Map attributes; private Set tags; @@ -43,6 +45,7 @@ private DefaultEventMetadata(final Builder builder) { this.attributes = builder.attributes == null ? new HashMap<>() : new HashMap<>(builder.attributes); this.tags = builder.tags == null ? new HashSet<>() : new HashSet(builder.tags); + this.externalOriginationTime = null; } private DefaultEventMetadata(final EventMetadata eventMetadata) { @@ -50,6 +53,7 @@ private DefaultEventMetadata(final EventMetadata eventMetadata) { this.timeReceived = eventMetadata.getTimeReceived(); this.attributes = new HashMap<>(eventMetadata.getAttributes()); this.tags = new HashSet<>(eventMetadata.getTags()); + this.externalOriginationTime = null; } @Override @@ -62,6 +66,16 @@ public Instant getTimeReceived() { return timeReceived; } + @Override + public Instant getExternalOriginationTime() { + return externalOriginationTime; + } + + @Override + public void setExternalOriginationTime(Instant externalOriginationTime) { + this.externalOriginationTime = externalOriginationTime; + } + @Override public Map getAttributes() { return attributes; diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java index 26f2d29e5f..64ef3be574 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java @@ -5,6 +5,8 @@ package org.opensearch.dataprepper.model.event; +import java.time.Instant; + public interface EventHandle { /** * releases event handle @@ -14,4 +16,30 @@ public interface EventHandle { * @since 2.2 */ void release(boolean result); + + /** + * sets external origination time + * + * @param externalOriginationTime externalOriginationTime to be set in the event handle + * @since 2.6 + */ + void setExternalOriginationTime(final Instant externalOriginationTime); + + /** + * gets external origination time + * + * @return returns externalOriginationTime from the event handle. This can be null if it is never set. + * @since 2.6 + */ + Instant getExternalOriginationTime(); + + /** + * gets internal origination time + * + * @return returns internalOriginationTime from the event handle. + * @since 2.6 + */ + Instant getInternalOriginationTime(); + + } diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventMetadata.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventMetadata.java index 511a87c1fa..5db56ba85c 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventMetadata.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventMetadata.java @@ -31,6 +31,20 @@ public interface EventMetadata extends Serializable { */ Instant getTimeReceived(); + /** + * Returns the external origination time of the event + * @return the external origination time + * @since 2.6 + */ + Instant getExternalOriginationTime(); + + /** + * Sets the external origination time of the event + * @param externalOriginationTime the external origination time + * @since 2.6 + */ + void setExternalOriginationTime(Instant externalOriginationTime); + /** * Returns the attributes * @return a map of attributes diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/InternalEventHandle.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/InternalEventHandle.java new file mode 100644 index 0000000000..3817365f17 --- /dev/null +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/InternalEventHandle.java @@ -0,0 +1,28 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.model.event; + +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; + +public interface InternalEventHandle { + /** + * sets acknowledgement set + * + * @param acknowledgementSet acknowledgementSet to be set in the event handle + * @since 2.6 + */ + void setAcknowledgementSet(final AcknowledgementSet acknowledgementSet); + + /** + * gets acknowledgement set + * + * @return returns acknowledgementSet from the event handle + * @since 2.6 + */ + AcknowledgementSet getAcknowledgementSet(); + +} + 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 686ee1d59f..846c4c3011 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 @@ -91,11 +91,13 @@ protected JacksonEvent(final Builder builder) { } this.jsonNode = getInitialJsonNode(builder.data); + this.eventHandle = new DefaultEventHandle(eventMetadata.getTimeReceived()); } protected JacksonEvent(final JacksonEvent otherEvent) { this.jsonNode = otherEvent.jsonNode.deepCopy(); this.eventMetadata = DefaultEventMetadata.fromEventMetadata(otherEvent.eventMetadata); + this.eventHandle = new DefaultEventHandle(eventMetadata.getTimeReceived()); } public static Event fromMessage(String message) { @@ -152,10 +154,6 @@ public void put(final String key, final Object value) { } } - public void setEventHandle(EventHandle handle) { - this.eventHandle = handle; - } - @Override public EventHandle getEventHandle() { return eventHandle; diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventHandleTests.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventHandleTests.java new file mode 100644 index 0000000000..3f1af4fb26 --- /dev/null +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventHandleTests.java @@ -0,0 +1,59 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.model.event; + +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import org.junit.jupiter.api.Test; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.mockito.Mockito.verify; +import static org.mockito.ArgumentMatchers.any; +import org.mockito.Mock; + +import java.time.Instant; + +class DefaultEventHandleTests { + @Mock + private AcknowledgementSet acknowledgementSet; + + @Test + void testBasic() { + Instant now = Instant.now(); + DefaultEventHandle eventHandle = new DefaultEventHandle(now); + assertThat(eventHandle.getAcknowledgementSet(), equalTo(null)); + assertThat(eventHandle.getInternalOriginationTime(), equalTo(now)); + assertThat(eventHandle.getExternalOriginationTime(), equalTo(null)); + eventHandle.release(true); + } + + @Test + void testWithAcknowledgementSet() { + acknowledgementSet = mock(AcknowledgementSet.class); + when(acknowledgementSet.release(any(EventHandle.class), any(Boolean.class))).thenReturn(true); + Instant now = Instant.now(); + DefaultEventHandle eventHandle = new DefaultEventHandle(now); + assertThat(eventHandle.getAcknowledgementSet(), equalTo(null)); + assertThat(eventHandle.getInternalOriginationTime(), equalTo(now)); + assertThat(eventHandle.getExternalOriginationTime(), equalTo(null)); + eventHandle.setAcknowledgementSet(acknowledgementSet); + eventHandle.release(true); + verify(acknowledgementSet).release(eventHandle, true); + } + + @Test + void testWithExternalOriginationTime() { + Instant now = Instant.now(); + DefaultEventHandle eventHandle = new DefaultEventHandle(now); + assertThat(eventHandle.getAcknowledgementSet(), equalTo(null)); + assertThat(eventHandle.getInternalOriginationTime(), equalTo(now)); + assertThat(eventHandle.getExternalOriginationTime(), equalTo(null)); + eventHandle.setExternalOriginationTime(now.minusSeconds(60)); + assertThat(eventHandle.getExternalOriginationTime(), equalTo(now.minusSeconds(60))); + eventHandle.release(true); + } +} diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventMetadataTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventMetadataTest.java index fa624a9e2b..479e7be0c2 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventMetadataTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventMetadataTest.java @@ -24,6 +24,7 @@ import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.not; import static org.hamcrest.CoreMatchers.notNullValue; +import static org.hamcrest.CoreMatchers.nullValue; import static org.hamcrest.CoreMatchers.sameInstance; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.anEmptyMap; @@ -80,6 +81,16 @@ public void testGetTimeReceived() { assertThat(timeReceived, is(equalTo(testTimeReceived))); } + @Test + public void testExternalOriginationTime() { + Instant externalOriginationTime = eventMetadata.getExternalOriginationTime(); + assertThat(externalOriginationTime, is(nullValue())); + Instant now = Instant.now(); + eventMetadata.setExternalOriginationTime(now); + externalOriginationTime = eventMetadata.getExternalOriginationTime(); + assertThat(externalOriginationTime, is(equalTo(now))); + } + @Test public void testGetAttributes() { final Map attributes = eventMetadata.getAttributes(); 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 62ce3dc48d..deb68f83bd 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 @@ -42,12 +42,6 @@ public class JacksonEventTest { - class TestEventHandle implements EventHandle { - @Override - public void release(boolean result) { - } - } - private Event event; private String eventType; @@ -398,6 +392,8 @@ public void testBuild_withEventType() { .build(); assertThat(event.getMetadata().getEventType(), is(equalTo(eventType))); + assertThat(event.getEventHandle(), is(notNullValue())); + assertThat(event.getEventHandle().getInternalOriginationTime(), is(notNullValue())); } @Test @@ -411,6 +407,8 @@ public void testBuild_withTimeReceived() { .build(); assertThat(event.getMetadata().getTimeReceived(), is(equalTo(now))); + assertThat(event.getEventHandle(), is(notNullValue())); + assertThat(event.getEventHandle().getInternalOriginationTime(), is(equalTo(now))); } @Test @@ -422,6 +420,8 @@ public void testBuild_withMessageValue() { assertThat(event, is(notNullValue())); assertThat(event.get("message", String.class), is(equalTo(message))); + assertThat(event.getEventHandle(), is(notNullValue())); + assertThat(event.getEventHandle().getInternalOriginationTime(), is(notNullValue())); } @Test @@ -678,6 +678,8 @@ void fromEvent_with_a_JacksonEvent() { assertThat(createdEvent, notNullValue()); assertThat(createdEvent, not(sameInstance(originalEvent))); + assertThat(event.getEventHandle(), is(notNullValue())); + assertThat(event.getEventHandle().getInternalOriginationTime(), is(notNullValue())); assertThat(createdEvent.toMap(), equalTo(dataObject)); assertThat(createdEvent.getJsonNode(), not(sameInstance(originalEvent.getJsonNode()))); @@ -707,19 +709,6 @@ void fromEvent_with_a_non_JacksonEvent() { assertThat(createdEvent.getMetadata(), equalTo(eventMetadata)); } - @Test - void testEventHandleGetAndSet() { - EventHandle testEventHandle = new TestEventHandle(); - final String jsonString = "{\"foo\": \"bar\"}"; - - final JacksonEvent event = JacksonEvent.builder() - .withEventType(eventType) - .withData(jsonString) - .build(); - event.setEventHandle(testEventHandle); - assertThat(event.getEventHandle(), equalTo(testEventHandle)); - } - @Test void testJsonStringBuilder() { final String jsonString = "{\"foo\":\"bar\"}"; diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitor.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitor.java index 1057418876..af9860cc9a 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitor.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitor.java @@ -6,7 +6,8 @@ package org.opensearch.dataprepper.acknowledgements; import org.opensearch.dataprepper.model.event.EventHandle; -import org.opensearch.dataprepper.event.DefaultEventHandle; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; +import org.opensearch.dataprepper.model.event.InternalEventHandle; import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; import java.util.concurrent.locks.ReentrantLock; @@ -33,7 +34,12 @@ class AcknowledgementSetMonitor implements Runnable { private final AtomicInteger numNullHandles; private DefaultAcknowledgementSet getAcknowledgementSet(final EventHandle eventHandle) { - return (DefaultAcknowledgementSet)((DefaultEventHandle)eventHandle).getAcknowledgementSet(); + if (eventHandle instanceof DefaultEventHandle) { + InternalEventHandle internalEventHandle = (InternalEventHandle)(DefaultEventHandle)eventHandle; + return (DefaultAcknowledgementSet)internalEventHandle.getAcknowledgementSet(); + } else { + throw new RuntimeException("Unsupported event handle"); + } } public AcknowledgementSetMonitor() { diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java index 3c8fe12159..741f08939d 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java @@ -5,10 +5,11 @@ package org.opensearch.dataprepper.acknowledgements; -import org.opensearch.dataprepper.event.DefaultEventHandle; import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.event.EventHandle; +import org.opensearch.dataprepper.model.event.InternalEventHandle; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; import org.opensearch.dataprepper.model.event.JacksonEvent; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -53,9 +54,12 @@ public void add(Event event) { lock.lock(); try { if (event instanceof JacksonEvent) { - EventHandle eventHandle = new DefaultEventHandle(this); - ((JacksonEvent) event).setEventHandle(eventHandle); - pendingAcknowledgments.put(eventHandle, new AtomicInteger(1)); + EventHandle eventHandle = event.getEventHandle(); + if (eventHandle instanceof DefaultEventHandle) { + InternalEventHandle internalEventHandle = (InternalEventHandle)(DefaultEventHandle)eventHandle; + internalEventHandle.setAcknowledgementSet(this); + pendingAcknowledgments.put(eventHandle, new AtomicInteger(1)); + } } } finally { lock.unlock(); diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/event/DefaultEventHandle.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/event/DefaultEventHandle.java deleted file mode 100644 index 025b130bf6..0000000000 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/event/DefaultEventHandle.java +++ /dev/null @@ -1,29 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.event; - -import org.opensearch.dataprepper.model.event.EventHandle; -import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; -import java.lang.ref.WeakReference; - -public class DefaultEventHandle implements EventHandle { - private final WeakReference acknowledgementSetRef; - public DefaultEventHandle(AcknowledgementSet acknowledgementSet) { - this.acknowledgementSetRef = new WeakReference<>(acknowledgementSet); - } - - public AcknowledgementSet getAcknowledgementSet() { - return acknowledgementSetRef.get(); - } - - @Override - public void release(boolean result) { - AcknowledgementSet acknowledgementSet = getAcknowledgementSet(); - if (acknowledgementSet != null) { - acknowledgementSet.release(this, result); - } - } -} 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 56d81ba68b..fb4effb413 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 @@ -13,6 +13,8 @@ import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.event.EventHandle; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; +import org.opensearch.dataprepper.model.event.InternalEventHandle; import org.opensearch.dataprepper.pipeline.common.FutureHelper; import org.opensearch.dataprepper.pipeline.common.FutureHelperResult; import org.slf4j.Logger; @@ -22,7 +24,6 @@ import java.util.List; import java.util.ArrayList; import java.util.Map; -import java.util.Objects; import java.util.Set; import java.util.concurrent.Future; import java.util.stream.Collectors; @@ -97,10 +98,15 @@ private void processAcknowledgements(List inputEvents, Collection outputR // For each event in the input events list that is not present in the output events, send positive acknowledgement, if acknowledgements are enabled for it inputEvents.forEach(event -> { EventHandle eventHandle = event.getEventHandle(); - if (Objects.nonNull(eventHandle) && !outputEventsSet.contains(event)) { - eventHandle.release(true); - } else if (acknowledgementsEnabled && Objects.isNull(eventHandle)) { - invalidEventHandlesCounter.increment(); + if (eventHandle != null && eventHandle instanceof DefaultEventHandle) { + InternalEventHandle internalEventHandle = (InternalEventHandle)(DefaultEventHandle)eventHandle; + if (internalEventHandle.getAcknowledgementSet() != null && !outputEventsSet.contains(event)) { + eventHandle.release(true); + } else if (acknowledgementsEnabled) { + invalidEventHandlesCounter.increment(); + } + } else if (eventHandle != null) { + throw new RuntimeException("Unexpected EventHandle"); } }); } diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategy.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategy.java index df7316b981..1bd2944c2e 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategy.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategy.java @@ -18,7 +18,7 @@ import org.opensearch.dataprepper.model.event.EventHandle; import org.opensearch.dataprepper.model.event.EventBuilder; import org.opensearch.dataprepper.model.event.EventMetadata; -import org.opensearch.dataprepper.event.DefaultEventHandle; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; import org.opensearch.dataprepper.acknowledgements.InactiveAcknowledgementSetManager; @@ -65,7 +65,7 @@ private void acquireEventReference(final Record record) { } if (referencedRecords.contains(record) || ((routedRecords != null) && routedRecords.contains(record))) { EventHandle eventHandle = ((JacksonEvent)record.getData()).getEventHandle(); - if (eventHandle != null) { + if (eventHandle != null && eventHandle instanceof DefaultEventHandle) { acknowledgementSetManager.acquireEventReference(eventHandle); } } else if (!referencedRecords.contains(record)) { @@ -97,7 +97,7 @@ public Record getRecord(final Record record) { JacksonEvent newRecordEvent; Record newRecord; DefaultEventHandle eventHandle = (DefaultEventHandle)recordEvent.getEventHandle(); - if (eventHandle != null) { + if (eventHandle != null && eventHandle.getAcknowledgementSet() != null) { final EventMetadata eventMetadata = recordEvent.getMetadata(); final EventBuilder eventBuilder = (EventBuilder) eventFactory.eventBuilder(EventBuilder.class).withEventMetadata(eventMetadata).withData(recordEvent.toMap()); newRecordEvent = (JacksonEvent) eventBuilder.build(); diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitorTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitorTests.java index 8c9d065704..6c85b5c4de 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitorTests.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementSetMonitorTests.java @@ -5,7 +5,7 @@ package org.opensearch.dataprepper.acknowledgements; -import org.opensearch.dataprepper.event.DefaultEventHandle; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.Mock; diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManagerTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManagerTests.java index 486617e9a0..c9fd556214 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManagerTests.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManagerTests.java @@ -6,7 +6,7 @@ package org.opensearch.dataprepper.acknowledgements; import org.opensearch.dataprepper.model.event.JacksonEvent; -import org.opensearch.dataprepper.model.event.EventHandle; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; import org.junit.jupiter.api.BeforeEach; @@ -18,8 +18,6 @@ import static org.awaitility.Awaitility.await; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.doAnswer; -import static org.mockito.ArgumentMatchers.any; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; @@ -40,32 +38,28 @@ class DefaultAcknowledgementSetManagerTests { @Mock JacksonEvent event3; - EventHandle eventHandle1; - EventHandle eventHandle2; - EventHandle eventHandle3; + DefaultEventHandle eventHandle1; + DefaultEventHandle eventHandle2; + DefaultEventHandle eventHandle3; Boolean result; @BeforeEach void setup() { callbackExecutor = Executors.newFixedThreadPool(2); event1 = mock(JacksonEvent.class); - doAnswer((i) -> { - eventHandle1 = i.getArgument(0); - return null; - }).when(event1).setEventHandle(any()); + eventHandle1 = mock(DefaultEventHandle.class); lenient().when(event1.getEventHandle()).thenReturn(eventHandle1); event2 = mock(JacksonEvent.class); - doAnswer((i) -> { - eventHandle2 = i.getArgument(0); - return null; - }).when(event2).setEventHandle(any()); + eventHandle2 = mock(DefaultEventHandle.class); lenient().when(event2.getEventHandle()).thenReturn(eventHandle2); acknowledgementSetManager = createObjectUnderTest(); AcknowledgementSet acknowledgementSet1 = acknowledgementSetManager.create((flag) -> { result = flag; }, TEST_TIMEOUT); acknowledgementSet1.add(event1); acknowledgementSet1.add(event2); + lenient().when(eventHandle1.getAcknowledgementSet()).thenReturn(acknowledgementSet1); + lenient().when(eventHandle2.getAcknowledgementSet()).thenReturn(acknowledgementSet1); acknowledgementSet1.complete(); } @@ -97,14 +91,12 @@ void testExpirations() throws InterruptedException { @Test void testMultipleAcknowledgementSets() { event3 = mock(JacksonEvent.class); - doAnswer((i) -> { - eventHandle3 = i.getArgument(0); - return null; - }).when(event3).setEventHandle(any()); + eventHandle3 = mock(DefaultEventHandle.class); lenient().when(event3.getEventHandle()).thenReturn(eventHandle3); AcknowledgementSet acknowledgementSet2 = acknowledgementSetManager.create((flag) -> { result = flag; }, TEST_TIMEOUT); acknowledgementSet2.add(event3); + lenient().when(eventHandle3.getAcknowledgementSet()).thenReturn(acknowledgementSet2); acknowledgementSet2.complete(); acknowledgementSetManager.releaseEventReference(eventHandle2, true); diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetTests.java index c4403e4b2f..8a4aa1485a 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetTests.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetTests.java @@ -11,8 +11,9 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.event.DefaultEventHandle; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; import org.opensearch.dataprepper.model.event.JacksonEvent; +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; import java.time.Duration; import java.util.concurrent.ExecutorService; @@ -24,7 +25,6 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; @@ -34,7 +34,12 @@ class DefaultAcknowledgementSetTests { private DefaultAcknowledgementSet defaultAcknowledgementSet; @Mock private JacksonEvent event; + @Mock + private JacksonEvent event2; + @Mock private DefaultEventHandle handle; + @Mock + private DefaultEventHandle handle2; private ExecutorService executor; private Boolean acknowledgementSetResult; @@ -78,12 +83,21 @@ void setupEvent() { callbackInterrupted = new AtomicBoolean(false); event = mock(JacksonEvent.class); - - doAnswer((i) -> { - handle = i.getArgument(0); + handle = mock(DefaultEventHandle.class); + lenient().doAnswer(a -> { + AcknowledgementSet acknowledgementSet = a.getArgument(0); + lenient().when(handle.getAcknowledgementSet()).thenReturn(acknowledgementSet); return null; - }).when(event).setEventHandle(any()); + }).when(handle).setAcknowledgementSet(any(AcknowledgementSet.class)); lenient().when(event.getEventHandle()).thenReturn(handle); + event2 = mock(JacksonEvent.class); + lenient().doAnswer(a -> { + AcknowledgementSet acknowledgementSet = a.getArgument(0); + lenient().when(handle2.getAcknowledgementSet()).thenReturn(acknowledgementSet); + return null; + }).when(handle2).setAcknowledgementSet(any(AcknowledgementSet.class)); + handle2 = mock(DefaultEventHandle.class); + lenient().when(event2.getEventHandle()).thenReturn(handle2); } @Test @@ -115,7 +129,6 @@ void testDefaultAcknowledgementInvalidAcquire() { defaultAcknowledgementSet.add(event); defaultAcknowledgementSet.complete(); DefaultAcknowledgementSet secondAcknowledgementSet = createObjectUnderTest(); - DefaultEventHandle handle2 = new DefaultEventHandle(secondAcknowledgementSet); defaultAcknowledgementSet.acquire(handle2); assertThat(invalidAcquiresCounter, equalTo(1)); } @@ -125,7 +138,6 @@ void testDefaultAcknowledgementInvalidRelease() { defaultAcknowledgementSet.add(event); defaultAcknowledgementSet.complete(); DefaultAcknowledgementSet secondAcknowledgementSet = createObjectUnderTest(); - DefaultEventHandle handle2 = new DefaultEventHandle(secondAcknowledgementSet); assertThat(defaultAcknowledgementSet.release(handle2, true), equalTo(false)); assertThat(invalidReleasesCounter, equalTo(1)); } @@ -170,6 +182,11 @@ void testDefaultAcknowledgementSetNegativeAcknowledgements() throws Exception { defaultAcknowledgementSet.add(event); defaultAcknowledgementSet.complete(); assertThat(handle, not(equalTo(null))); + lenient().doAnswer(a -> { + AcknowledgementSet acknowledgementSet = a.getArgument(0); + lenient().when(handle.getAcknowledgementSet()).thenReturn(acknowledgementSet); + return null; + }).when(handle).setAcknowledgementSet(any(AcknowledgementSet.class)); assertThat(handle.getAcknowledgementSet(), equalTo(defaultAcknowledgementSet)); defaultAcknowledgementSet.acquire(handle); assertThat(defaultAcknowledgementSet.release(handle, true), equalTo(false)); @@ -198,6 +215,11 @@ void testDefaultAcknowledgementSetExpirations() throws Exception { ); defaultAcknowledgementSet.add(event); defaultAcknowledgementSet.complete(); + lenient().doAnswer(a -> { + AcknowledgementSet acknowledgementSet = a.getArgument(0); + lenient().when(handle.getAcknowledgementSet()).thenReturn(acknowledgementSet); + return null; + }).when(handle).setAcknowledgementSet(any(AcknowledgementSet.class)); assertThat(handle, not(equalTo(null))); assertThat(handle.getAcknowledgementSet(), equalTo(defaultAcknowledgementSet)); assertThat(defaultAcknowledgementSet.release(handle, true), equalTo(true)); diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/event/DefaultEventHandleTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/event/DefaultEventHandleTests.java deleted file mode 100644 index 2b0895ad6c..0000000000 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/event/DefaultEventHandleTests.java +++ /dev/null @@ -1,32 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.event; - -import org.opensearch.dataprepper.model.event.EventHandle; -import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.equalTo; -import org.junit.jupiter.api.Test; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; -import static org.mockito.Mockito.verify; -import static org.mockito.ArgumentMatchers.any; -import org.mockito.Mock; - -class DefaultEventHandleTests { - @Mock - private AcknowledgementSet acknowledgementSet; - - @Test - void testBasic() { - acknowledgementSet = mock(AcknowledgementSet.class); - when(acknowledgementSet.release(any(EventHandle.class), any(Boolean.class))).thenReturn(true); - DefaultEventHandle eventHandle = new DefaultEventHandle(acknowledgementSet); - assertThat(eventHandle.getAcknowledgementSet(), equalTo(acknowledgementSet)); - eventHandle.release(true); - verify(acknowledgementSet).release(eventHandle, true); - } -} diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/PipelineTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/PipelineTests.java index 2fc86b80eb..75c1154baa 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/PipelineTests.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/PipelineTests.java @@ -10,10 +10,11 @@ import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; import org.opensearch.dataprepper.model.buffer.Buffer; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.event.EventFactory; -import org.opensearch.dataprepper.model.event.EventHandle; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; import org.opensearch.dataprepper.model.event.JacksonEvent; import org.opensearch.dataprepper.model.processor.Processor; import org.opensearch.dataprepper.model.record.Record; @@ -83,7 +84,7 @@ class PipelineTests { private Duration peerForwarderDrainTimeout; private EventFactory eventFactory; private JacksonEvent event; - private EventHandle eventHandle; + private DefaultEventHandle eventHandle; private AcknowledgementSetManager acknowledgementSetManager; @BeforeEach @@ -391,6 +392,7 @@ class PublishToSink { private List records; private List> dataFlowComponents; private Source mockSource; + private AcknowledgementSet acknowledgementSet; @BeforeEach void setUp() { @@ -426,8 +428,10 @@ void publishToSinks_calls_route_with_Events_and_Sinks_verify_AcknowledgementSetM RouterCopyRecordStrategy routerCopyRecordStrategy = (RouterCopyRecordStrategy)a.getArgument(2); Record rec = records.get(0); event = mock(JacksonEvent.class); - eventHandle = mock(EventHandle.class); + eventHandle = mock(DefaultEventHandle.class); + acknowledgementSet = mock(AcknowledgementSet.class); when(event.getEventHandle()).thenReturn(eventHandle); + when(eventHandle.getAcknowledgementSet()).thenReturn(acknowledgementSet); when(rec.getData()).thenReturn(event); routerCopyRecordStrategy.getRecord(rec); routerCopyRecordStrategy.getRecord(rec); @@ -437,7 +441,7 @@ void publishToSinks_calls_route_with_Events_and_Sinks_verify_AcknowledgementSetM Pipeline pipeline = createObjectUnderTest(); when(mockSource.areAcknowledgementsEnabled()).thenReturn(true); pipeline.publishToSinks(records); - verify(acknowledgementSetManager).acquireEventReference(any(EventHandle.class)); + verify(acknowledgementSetManager).acquireEventReference(any(DefaultEventHandle.class)); verify(router) .route(anyCollection(), eq(dataFlowComponents), any(RouterGetRecordStrategy.class), any(BiConsumer.class)); @@ -451,7 +455,7 @@ void publishToSinks_calls_route_with_Events_and_Sinks_verify_InactiveAcknowledge RouterCopyRecordStrategy routerCopyRecordStrategy = (RouterCopyRecordStrategy)a.getArgument(2); Record rec = records.get(0); event = mock(JacksonEvent.class); - eventHandle = mock(EventHandle.class); + eventHandle = mock(DefaultEventHandle.class); when(event.getEventHandle()).thenReturn(null); when(rec.getData()).thenReturn(event); routerCopyRecordStrategy.getRecord(rec); diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategyTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategyTests.java index a9c030793b..4c56113323 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategyTests.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/pipeline/router/RouterCopyRecordStrategyTests.java @@ -42,10 +42,9 @@ import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.event.JacksonEvent; import org.opensearch.dataprepper.model.event.EventFactory; -import org.opensearch.dataprepper.model.event.EventHandle; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; import org.opensearch.dataprepper.model.event.EventBuilder; import org.opensearch.dataprepper.model.event.EventMetadata; -import org.opensearch.dataprepper.event.DefaultEventHandle; import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; @@ -67,7 +66,7 @@ public class RouterCopyRecordStrategyTests { private JacksonEvent event; - private Map handleRefCount; + private Map handleRefCount; private static class TestComponent { } @@ -80,11 +79,11 @@ void setUp() { acknowledgementSet1 = mock(AcknowledgementSet.class); try { lenient().doAnswer((i) -> { - EventHandle handle = (EventHandle) i.getArgument(0); + DefaultEventHandle handle = (DefaultEventHandle) i.getArgument(0); int v = handleRefCount.getOrDefault(handle, 0); handleRefCount.put(handle, v+1); return null; - }).when(acknowledgementSetManager).acquireEventReference(any(EventHandle.class)); + }).when(acknowledgementSetManager).acquireEventReference(any(DefaultEventHandle.class)); } catch (Exception e){} mockRecordsIn = IntStream.range(0, 10) .mapToObj(i -> mock(Record.class)) @@ -99,12 +98,12 @@ void setUp() { .collect(Collectors.toList()); } - private void attachEventHandlesToRecordsIn(List eventHandles) { + private void attachEventHandlesToRecordsIn(List eventHandles) { Iterator iter = recordsIn.iterator(); while (iter.hasNext()) { Record r = (Record) iter.next(); - DefaultEventHandle handle = new DefaultEventHandle(acknowledgementSet1); - ((JacksonEvent)r.getData()).setEventHandle(handle); + DefaultEventHandle handle = (DefaultEventHandle)((JacksonEvent)r.getData()).getEventHandle(); + handle.setAcknowledgementSet(acknowledgementSet1); eventHandles.add(handle); } } @@ -187,10 +186,10 @@ void test_one_record_with_acknowledgements() { = Collections.singletonList(dataFlowComponent); final RouterCopyRecordStrategy getRecordStrategy = createObjectUnderTest(dataFlowComponents); - List eventHandles = new ArrayList<>(); + List eventHandles = new ArrayList<>(); attachEventHandlesToRecordsIn(eventHandles); Record firstRecord = recordsIn.iterator().next(); - EventHandle firstHandle = ((Event)firstRecord.getData()).getEventHandle(); + DefaultEventHandle firstHandle = (DefaultEventHandle)((Event)firstRecord.getData()).getEventHandle(); Record recordOut = getRecordStrategy.getRecord(firstRecord); assertThat(recordOut, sameInstance(firstRecord)); assertTrue(getRecordStrategy.getReferencedRecords().contains(firstRecord)); @@ -209,7 +208,7 @@ void test_multiple_records_with_acknowledgements() { = Collections.singletonList(dataFlowComponent); final RouterCopyRecordStrategy getRecordStrategy = createObjectUnderTest(dataFlowComponents); - List eventHandles = new ArrayList<>(); + List eventHandles = new ArrayList<>(); attachEventHandlesToRecordsIn(eventHandles); Collection recordsOut = getRecordStrategy.getAllRecords(recordsIn); assertThat(recordsOut.size(), equalTo(recordsIn.size())); @@ -238,12 +237,12 @@ void test_one_record_with_acknowledgements_and_multi_components() { } final RouterCopyRecordStrategy getRecordStrategy = createObjectUnderTest(dataFlowComponents); - List eventHandles = new ArrayList<>(); + List eventHandles = new ArrayList<>(); attachEventHandlesToRecordsIn(eventHandles); try { doAnswer((i) -> { JacksonEvent e1 = (JacksonEvent) i.getArgument(0); - e1.setEventHandle(new DefaultEventHandle(acknowledgementSet1)); + ((DefaultEventHandle)e1.getEventHandle()).setAcknowledgementSet(acknowledgementSet1); return null; }).when(acknowledgementSet1).add(any(JacksonEvent.class)); } catch (Exception e){} @@ -254,14 +253,14 @@ void test_one_record_with_acknowledgements_and_multi_components() { when(eventBuilder.build()).thenReturn(JacksonEvent.fromEvent(event)); when(eventFactory.eventBuilder(EventBuilder.class)).thenReturn(eventBuilder); Record firstRecord = recordsIn.iterator().next(); - EventHandle firstHandle = ((Event)firstRecord.getData()).getEventHandle(); + DefaultEventHandle firstHandle = (DefaultEventHandle)((Event)firstRecord.getData()).getEventHandle(); Record recordOut = getRecordStrategy.getRecord(firstRecord); assertThat(recordOut, sameInstance(firstRecord)); assertTrue(getRecordStrategy.getReferencedRecords().contains(firstRecord)); recordOut = getRecordStrategy.getRecord(firstRecord); assertThat(recordOut, not(sameInstance(firstRecord))); assertFalse(handleRefCount.containsKey(firstHandle)); - EventHandle newHandle = ((JacksonEvent)recordOut.getData()).getEventHandle(); + DefaultEventHandle newHandle = (DefaultEventHandle)((JacksonEvent)recordOut.getData()).getEventHandle(); assertTrue(getRecordStrategy.getReferencedRecords().contains(recordOut)); assertThat(newHandle, not(equalTo(null))); assertFalse(handleRefCount.containsKey(newHandle)); @@ -276,12 +275,12 @@ void test_multiple_records_with_acknowledgements_and_multi_components() { } final RouterCopyRecordStrategy getRecordStrategy = createObjectUnderTest(dataFlowComponents); - List eventHandles = new ArrayList<>(); + List eventHandles = new ArrayList<>(); attachEventHandlesToRecordsIn(eventHandles); try { doAnswer((i) -> { JacksonEvent e1 = (JacksonEvent) i.getArgument(0); - e1.setEventHandle(new DefaultEventHandle(acknowledgementSet1)); + ((DefaultEventHandle)e1.getEventHandle()).setAcknowledgementSet(acknowledgementSet1); return null; }).when(acknowledgementSet1).add(any(JacksonEvent.class)); } catch (Exception e){} diff --git a/data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsService.java b/data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsService.java index a27d9ea4cb..fc9963ab46 100644 --- a/data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsService.java +++ b/data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsService.java @@ -104,9 +104,7 @@ private void stageLogEvents() { } private void addToBuffer(final Record log, final String logString) { - if (log.getData().getEventHandle() != null) { - bufferedEventHandles.add(log.getData().getEventHandle()); - } + bufferedEventHandles.add(log.getData().getEventHandle()); buffer.writeEvent(logString.getBytes(StandardCharsets.UTF_8)); } } diff --git a/data-prepper-plugins/cloudwatch-logs/src/test/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsServiceTest.java b/data-prepper-plugins/cloudwatch-logs/src/test/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsServiceTest.java index e8e416d53e..8e7fcebc4e 100644 --- a/data-prepper-plugins/cloudwatch-logs/src/test/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsServiceTest.java +++ b/data-prepper-plugins/cloudwatch-logs/src/test/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsServiceTest.java @@ -8,7 +8,6 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.model.event.Event; -import org.opensearch.dataprepper.model.event.EventHandle; import org.opensearch.dataprepper.model.event.JacksonEvent; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.sink.cloudwatch_logs.buffer.Buffer; @@ -62,8 +61,6 @@ Collection> getSampleRecordsCollectionSmall() { final ArrayList> returnCollection = new ArrayList<>(); for (int i = 0; i < 5; i++) { JacksonEvent mockJacksonEvent = (JacksonEvent) JacksonEvent.fromMessage("testMessage"); - final EventHandle mockEventHandle = mock(EventHandle.class); - mockJacksonEvent.setEventHandle(mockEventHandle); returnCollection.add(new Record<>(mockJacksonEvent)); } @@ -74,8 +71,6 @@ Collection> getSampleRecordsCollection() { final ArrayList> returnCollection = new ArrayList<>(); for (int i = 0; i < thresholdConfig.getBatchSize(); i++) { JacksonEvent mockJacksonEvent = (JacksonEvent) JacksonEvent.fromMessage("testMessage"); - final EventHandle mockEventHandle = mock(EventHandle.class); - mockJacksonEvent.setEventHandle(mockEventHandle); returnCollection.add(new Record<>(mockJacksonEvent)); } @@ -86,8 +81,6 @@ Collection> getSampleRecordsOfLargerSize() { final ArrayList> returnCollection = new ArrayList<>(); for (int i = 0; i < thresholdConfig.getBatchSize() * 2; i++) { JacksonEvent mockJacksonEvent = (JacksonEvent) JacksonEvent.fromMessage("a".repeat((int) (thresholdConfig.getMaxRequestSizeBytes()/24))); - final EventHandle mockEventHandle = mock(EventHandle.class); - mockJacksonEvent.setEventHandle(mockEventHandle); returnCollection.add(new Record<>(mockJacksonEvent)); } @@ -98,8 +91,6 @@ Collection> getSampleRecordsOfLimitSize() { final ArrayList> returnCollection = new ArrayList<>(); for (int i = 0; i < thresholdConfig.getBatchSize(); i++) { JacksonEvent mockJacksonEvent = (JacksonEvent) JacksonEvent.fromMessage("testMessage".repeat((int) thresholdConfig.getMaxEventSizeBytes())); - final EventHandle mockEventHandle = mock(EventHandle.class); - mockJacksonEvent.setEventHandle(mockEventHandle); returnCollection.add(new Record<>(mockJacksonEvent)); } 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 64e7ad23a4..77243d7c12 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 @@ -13,7 +13,6 @@ import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.event.Event; -import org.opensearch.dataprepper.model.event.EventHandle; import org.opensearch.dataprepper.model.event.JacksonEvent; import org.opensearch.dataprepper.model.log.JacksonLog; import org.opensearch.dataprepper.model.record.Record; @@ -185,7 +184,6 @@ private Collection> setEventQueue() { private static Record createRecord() { String json = "{\"peer\": {\"ip\": \"136.226.242.205\", \"host\": \"example.org\" }, \"status\": \"success\"}"; final JacksonEvent event = JacksonLog.builder().withData(json).build(); - event.setEventHandle(mock(EventHandle.class)); return new Record<>(event); } -} \ No newline at end of file +} diff --git a/data-prepper-plugins/http-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/http/HttpSinkServiceIT.java b/data-prepper-plugins/http-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/http/HttpSinkServiceIT.java index 47a9a18f85..2095edcb48 100644 --- a/data-prepper-plugins/http-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/http/HttpSinkServiceIT.java +++ b/data-prepper-plugins/http-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/http/HttpSinkServiceIT.java @@ -20,7 +20,6 @@ import org.opensearch.dataprepper.model.configuration.PluginModel; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.event.Event; -import org.opensearch.dataprepper.model.event.EventHandle; import org.opensearch.dataprepper.model.event.JacksonEvent; import org.opensearch.dataprepper.model.log.JacksonLog; import org.opensearch.dataprepper.model.plugin.PluginFactory; @@ -139,7 +138,6 @@ private Collection> setEventQueue(final int records) { private static Record createRecord() { final JacksonEvent event = JacksonLog.builder().withData("{\"name\":\""+ UUID.randomUUID() +"\"}").build(); - event.setEventHandle(mock(EventHandle.class)); return new Record<>(event); } diff --git a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkService.java b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkService.java index 480886d52b..c392bf4ae0 100644 --- a/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkService.java +++ b/data-prepper-plugins/http-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/http/service/HttpSinkService.java @@ -179,9 +179,7 @@ public void output(Collection> records) { int count = currentBuffer.getEventCount() +1; currentBuffer.setEventCount(count); - if(event.getEventHandle() != null) { - bufferedEventHandles.add(event.getEventHandle()); - } + bufferedEventHandles.add(event.getEventHandle()); if (ThresholdValidator.checkThresholdExceed(currentBuffer, maxEvents, maxBytes, maxCollectionDuration)) { codec.complete(outputStream); final HttpEndPointResponse failedHttpEndPointResponses = pushToEndPoint(getCurrentBufferData(currentBuffer)); @@ -379,4 +377,4 @@ private void accessTokenIfExpired(final Integer tokenExpired,final String url){ } } -} \ No newline at end of file +} 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 ea0d6f6d59..ddca58ae2e 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 @@ -96,9 +96,7 @@ public void produceRawData(final byte[] bytes, final String key) { } public void produceRecords(final Record record) { - if (record.getData().getEventHandle() != null) { - bufferedEventHandles.add(record.getData().getEventHandle()); - } + bufferedEventHandles.add(record.getData().getEventHandle()); Event event = getEvent(record); final String key = event.formatString(kafkaProducerConfig.getPartitionKey(), expressionEvaluator); try { diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerTest.java index 13f9879660..8b23975401 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerTest.java @@ -21,7 +21,6 @@ import org.mockito.junit.jupiter.MockitoExtension; import org.mockito.junit.jupiter.MockitoSettings; import org.mockito.quality.Strictness; -import org.opensearch.dataprepper.event.DefaultEventHandle; import org.opensearch.dataprepper.expression.ExpressionEvaluator; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.event.JacksonEvent; @@ -66,8 +65,6 @@ public class KafkaCustomProducerTest { @BeforeEach public void setUp() { event = (JacksonEvent) JacksonEvent.fromMessage(UUID.randomUUID().toString()); - DefaultEventHandle defaultEventHandle = mock(DefaultEventHandle.class); - event.setEventHandle(defaultEventHandle); record = new Record<>(event); final TopicProducerConfig topicConfig = mock(TopicProducerConfig.class); when(topicConfig.getName()).thenReturn("test-topic"); 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 dff10c608c..755ccb5283 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 @@ -39,7 +39,6 @@ import org.opensearch.dataprepper.metrics.MetricsTestUtil; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.event.Event; -import org.opensearch.dataprepper.model.event.EventHandle; import org.opensearch.dataprepper.model.event.EventType; import org.opensearch.dataprepper.model.event.JacksonEvent; import org.opensearch.dataprepper.model.opensearch.OpenSearchBulkActions; @@ -89,7 +88,6 @@ import static org.hamcrest.Matchers.closeTo; 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.mock; import static org.mockito.Mockito.when; import static org.opensearch.dataprepper.plugins.sink.opensearch.OpenSearchIntegrationHelper.createContentParser; @@ -115,7 +113,6 @@ public class OpenSearchSinkIT { private static final String TRACE_INGESTION_TEST_DISABLED_REASON = "Trace ingestion is not supported for ES 6"; private RestClient client; - private EventHandle eventHandle; private SinkContext sinkContext; private String testTagsTargetKey; @@ -152,10 +149,6 @@ public void setup() { expressionEvaluator = mock(ExpressionEvaluator.class); when(expressionEvaluator.isValidExpressionStatement(any(String.class))).thenReturn(false); - eventHandle = mock(EventHandle.class); - lenient().doAnswer(a -> { - return null; - }).when(eventHandle).release(any(Boolean.class)); } @BeforeEach @@ -890,7 +883,6 @@ public void testEventOutputWithTags() throws IOException, InterruptedException { .withData("{\"log\": \"foobar\"}") .withEventType("event") .build(); - ((JacksonEvent)testEvent).setEventHandle(eventHandle); List tagsList = List.of("tag1", "tag2"); testEvent.getMetadata().addTags(tagsList); @@ -920,7 +912,6 @@ public void testEventOutput() throws IOException, InterruptedException { .withData("{\"log\": \"foobar\"}") .withEventType("event") .build(); - ((JacksonEvent) testEvent).setEventHandle(eventHandle); final List> testRecords = Collections.singletonList(new Record<>(testEvent)); @@ -948,7 +939,6 @@ public void testOpenSearchDocumentId(final String testDocumentIdField) throws IO .withData(Map.of("arbitrary_data", UUID.randomUUID().toString())) .withEventType("event") .build(); - ((JacksonEvent) testEvent).setEventHandle(eventHandle); testEvent.put(testDocumentIdField, expectedId); final List> testRecords = Collections.singletonList(new Record<>(testEvent)); @@ -974,7 +964,6 @@ public void testOpenSearchRoutingField(final String testRoutingField) throws IOE .withData(Map.of("arbitrary_data", UUID.randomUUID().toString())) .withEventType("event") .build(); - ((JacksonEvent) testEvent).setEventHandle(eventHandle); testEvent.put(testRoutingField, expectedRoutingField); final List> testRecords = Collections.singletonList(new Record<>(testEvent)); @@ -1003,7 +992,6 @@ public void testOpenSearchDynamicIndex(final String testIndex) throws IOExceptio .withData(dataMap) .withEventType("event") .build(); - ((JacksonEvent) testEvent).setEventHandle(eventHandle); testEvent.put(testIndex, testIndexName); Map expectedMap = testEvent.toMap(); @@ -1037,7 +1025,6 @@ public void testOpenSearchDynamicIndexWithDate(final String testIndex, final Str .withData(dataMap) .withEventType("event") .build(); - ((JacksonEvent) testEvent).setEventHandle(eventHandle); testEvent.put(testIndex, testIndexName); Map expectedMap = testEvent.toMap(); @@ -1067,7 +1054,6 @@ public void testOpenSearchIndexWithDate(final String testDatePattern) throws IOE .withData(dataMap) .withEventType("event") .build(); - ((JacksonEvent) testEvent).setEventHandle(eventHandle); Map expectedMap = testEvent.toMap(); @@ -1370,7 +1356,6 @@ private Record jsonStringToRecord(final String jsonString) { .withEventType(EventType.TRACE.toString()) .withData(objectMapper.readValue(jsonString, Map.class)).build()); JacksonEvent event = (JacksonEvent) record.getData(); - event.setEventHandle(eventHandle); return record; } catch (final JsonProcessingException e) { throw new RuntimeException(e); diff --git a/data-prepper-plugins/s3-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkIT.java b/data-prepper-plugins/s3-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkIT.java index a83a5be9db..7032104bb3 100644 --- a/data-prepper-plugins/s3-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkIT.java +++ b/data-prepper-plugins/s3-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkIT.java @@ -23,7 +23,6 @@ import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.event.DefaultEventMetadata; import org.opensearch.dataprepper.model.event.Event; -import org.opensearch.dataprepper.model.event.EventHandle; import org.opensearch.dataprepper.model.event.EventMetadata; import org.opensearch.dataprepper.model.event.EventType; import org.opensearch.dataprepper.model.event.JacksonEvent; @@ -262,7 +261,6 @@ private Event generateTestEvent(final Map eventData) { .withEventType(EventType.LOG.toString()) .build(); final JacksonEvent event = JacksonLog.builder().withData(eventData).withEventMetadata(defaultEventMetadata).build(); - event.setEventHandle(mock(EventHandle.class)); return JacksonEvent.builder() .withData(eventData) .withEventMetadata(defaultEventMetadata) diff --git a/data-prepper-plugins/s3-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkServiceIT.java b/data-prepper-plugins/s3-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkServiceIT.java index 9d67a4bf5b..951225937d 100644 --- a/data-prepper-plugins/s3-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkServiceIT.java +++ b/data-prepper-plugins/s3-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkServiceIT.java @@ -38,7 +38,6 @@ import org.opensearch.dataprepper.model.codec.OutputCodec; import org.opensearch.dataprepper.model.event.DefaultEventMetadata; import org.opensearch.dataprepper.model.event.Event; -import org.opensearch.dataprepper.model.event.EventHandle; import org.opensearch.dataprepper.model.event.EventMetadata; import org.opensearch.dataprepper.model.event.EventType; import org.opensearch.dataprepper.model.event.JacksonEvent; @@ -92,7 +91,6 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; import static org.mockito.Mockito.lenient; -import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @ExtendWith(MockitoExtension.class) @@ -296,7 +294,6 @@ private static Record createRecord() { withTags(testTags).build(); Map json = generateJson(); final JacksonEvent event = JacksonLog.builder().withData(json).withEventMetadata(defaultEventMetadata).build(); - event.setEventHandle(mock(EventHandle.class)); return new Record<>(event); } @@ -413,4 +410,4 @@ private List> createParquetRecordsList(final InputStream private MessageType createdParquetSchema(ParquetMetadata parquetMetadata) { return parquetMetadata.getFileMetaData().getSchema(); } -} \ No newline at end of file +} diff --git a/data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkService.java b/data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkService.java index 65de8225b8..879854b546 100644 --- a/data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkService.java +++ b/data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkService.java @@ -28,7 +28,6 @@ import java.util.Collection; import java.util.LinkedList; import java.util.List; -import java.util.Objects; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; @@ -127,9 +126,7 @@ void output(Collection> records) { int count = currentBuffer.getEventCount() + 1; currentBuffer.setEventCount(count); - if (event.getEventHandle() != null) { - bufferedEventHandles.add(event.getEventHandle()); - } + bufferedEventHandles.add(event.getEventHandle()); } catch (Exception ex) { if(sampleException == null) { sampleException = ex; @@ -149,7 +146,6 @@ void output(Collection> records) { failedEvents .stream() .map(Event::getEventHandle) - .filter(Objects::nonNull) .forEach(eventHandle -> eventHandle.release(false)); LOG.error("Unable to add {} events to buffer. Dropping these events. Sample exception provided.", failedEvents.size(), sampleException); } diff --git a/data-prepper-plugins/s3-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkServiceTest.java b/data-prepper-plugins/s3-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkServiceTest.java index b154d30526..7160660137 100644 --- a/data-prepper-plugins/s3-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkServiceTest.java +++ b/data-prepper-plugins/s3-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkServiceTest.java @@ -17,11 +17,13 @@ import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.event.EventHandle; +import org.opensearch.dataprepper.model.event.DefaultEventHandle; import org.opensearch.dataprepper.model.event.JacksonEvent; import org.opensearch.dataprepper.model.plugin.PluginFactory; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.sink.OutputCodecContext; import org.opensearch.dataprepper.model.types.ByteCount; +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; import org.opensearch.dataprepper.plugins.sink.s3.accumulator.Buffer; import org.opensearch.dataprepper.plugins.sink.s3.accumulator.BufferFactory; import org.opensearch.dataprepper.plugins.sink.s3.accumulator.BufferTypeOptions; @@ -84,6 +86,7 @@ class S3SinkServiceTest { private DistributionSummary s3ObjectSizeSummary; private Random random; private String tagsTargetKey; + private AcknowledgementSet acknowledgementSet; @BeforeEach void setUp() { @@ -91,6 +94,7 @@ void setUp() { random = new Random(); tagsTargetKey = RandomStringUtils.randomAlphabetic(5); s3SinkConfig = mock(S3SinkConfig.class); + acknowledgementSet = mock(AcknowledgementSet.class); codecContext = new OutputCodecContext(tagsTargetKey, Collections.emptyList(), Collections.emptyList()); s3Client = mock(S3Client.class); ThresholdOptions thresholdOptions = mock(ThresholdOptions.class); @@ -134,6 +138,10 @@ void setUp() { lenient().when(pluginMetrics.summary(S3SinkService.S3_OBJECTS_SIZE)).thenReturn(s3ObjectSizeSummary); } + private DefaultEventHandle castToDefaultHandle(EventHandle eventHandle) { + return (DefaultEventHandle)eventHandle; + } + private S3SinkService createObjectUnderTest() { return new S3SinkService(s3SinkConfig, bufferFactory, codec, codecContext, s3Client, keyGenerator, Duration.ofMillis(100), pluginMetrics); } @@ -377,13 +385,16 @@ void output_will_release_all_handles_since_a_flush() throws IOException { doNothing().when(codec).writeEvent(event, outputStream); final S3SinkService s3SinkService = createObjectUnderTest(); final Collection> records = generateRandomStringEventRecord(); + final List eventHandles = records.stream().map(Record::getData).map(Event::getEventHandle).map(this::castToDefaultHandle).collect(Collectors.toList()); + for (DefaultEventHandle eventHandle : eventHandles) { + eventHandle.setAcknowledgementSet(acknowledgementSet); + } s3SinkService.output(records); - final List eventHandles = records.stream().map(Record::getData).map(Event::getEventHandle).collect(Collectors.toList()); - for (EventHandle eventHandle : eventHandles) { - verify(eventHandle).release(true); + verify(acknowledgementSet).release(eventHandle, true); } + } @Test @@ -400,21 +411,29 @@ void output_will_skip_releasing_events_without_EventHandle_objects() throws IOEx doNothing().when(codec).writeEvent(event1, outputStream); final S3SinkService s3SinkService = createObjectUnderTest(); final Collection> records = generateRandomStringEventRecord(); - records.stream() - .map(Record::getData) - .map(event -> (JacksonEvent) event) - .forEach(event -> event.setEventHandle(null)); + final List eventHandles = records.stream().map(Record::getData).map(Event::getEventHandle).map(this::castToDefaultHandle).collect(Collectors.toList()); + for (DefaultEventHandle eventHandle : eventHandles) { + eventHandle.setAcknowledgementSet(acknowledgementSet); + } s3SinkService.output(records); + for (EventHandle eventHandle : eventHandles) { + verify(acknowledgementSet).release(eventHandle, true); + } final Collection> records2 = generateRandomStringEventRecord(); - s3SinkService.output(records2); + final List eventHandles2 = records2.stream().map(Record::getData).map(Event::getEventHandle).map(this::castToDefaultHandle).collect(Collectors.toList()); - final List eventHandles2 = records2.stream().map(Record::getData).map(Event::getEventHandle).collect(Collectors.toList()); + for (DefaultEventHandle eventHandle : eventHandles2) { + eventHandle.setAcknowledgementSet(acknowledgementSet); + } + + s3SinkService.output(records2); for (EventHandle eventHandle : eventHandles2) { - verify(eventHandle).release(true); + verify(acknowledgementSet).release(eventHandle, true); } + } @Test @@ -433,12 +452,15 @@ void output_will_release_all_handles_since_a_flush_when_S3_fails() throws IOExce doNothing().when(codec).writeEvent(event, outputStream); final S3SinkService s3SinkService = createObjectUnderTest(); final List> records = generateEventRecords(1); - s3SinkService.output(records); + final List eventHandles = records.stream().map(Record::getData).map(Event::getEventHandle).map(this::castToDefaultHandle).collect(Collectors.toList()); - final List eventHandles = records.stream().map(Record::getData).map(Event::getEventHandle).collect(Collectors.toList()); + for (DefaultEventHandle eventHandle : eventHandles) { + eventHandle.setAcknowledgementSet(acknowledgementSet); + } + s3SinkService.output(records); for (EventHandle eventHandle : eventHandles) { - verify(eventHandle).release(false); + verify(acknowledgementSet).release(eventHandle, false); } } @@ -456,21 +478,24 @@ void output_will_release_only_new_handles_since_a_flush() throws IOException { doNothing().when(codec).writeEvent(event, outputStream); final S3SinkService s3SinkService = createObjectUnderTest(); final Collection> records = generateRandomStringEventRecord(); + final List eventHandles = records.stream().map(Record::getData).map(Event::getEventHandle).map(this::castToDefaultHandle).collect(Collectors.toList()); + for (DefaultEventHandle eventHandle : eventHandles) { + eventHandle.setAcknowledgementSet(acknowledgementSet); + } s3SinkService.output(records); + for (EventHandle eventHandle : eventHandles) { + verify(acknowledgementSet).release(eventHandle, true); + } final Collection> records2 = generateRandomStringEventRecord(); - s3SinkService.output(records2); - - final List eventHandles1 = records.stream().map(Record::getData).map(Event::getEventHandle).collect(Collectors.toList()); - - for (EventHandle eventHandle : eventHandles1) { - verify(eventHandle).release(true); + final List eventHandles2 = records2.stream().map(Record::getData).map(Event::getEventHandle).map(this::castToDefaultHandle).collect(Collectors.toList()); + for (DefaultEventHandle eventHandle : eventHandles2) { + eventHandle.setAcknowledgementSet(acknowledgementSet); } - - final List eventHandles2 = records2.stream().map(Record::getData).map(Event::getEventHandle).collect(Collectors.toList()); - + s3SinkService.output(records2); for (EventHandle eventHandle : eventHandles2) { - verify(eventHandle).release(true); + verify(acknowledgementSet).release(eventHandle, true); } + } @Test @@ -489,6 +514,10 @@ void output_will_skip_and_drop_failed_records() throws IOException { List> records = generateEventRecords(2); Event event1 = records.get(0).getData(); Event event2 = records.get(1).getData(); + DefaultEventHandle eventHandle1 = (DefaultEventHandle)event1.getEventHandle(); + DefaultEventHandle eventHandle2 = (DefaultEventHandle)event2.getEventHandle(); + eventHandle1.setAcknowledgementSet(acknowledgementSet); + eventHandle2.setAcknowledgementSet(acknowledgementSet); doThrow(RuntimeException.class).when(codec).writeEvent(event1, outputStream); @@ -499,10 +528,10 @@ void output_will_skip_and_drop_failed_records() throws IOException { inOrder.verify(codec).writeEvent(event1, outputStream); inOrder.verify(codec).writeEvent(event2, outputStream); - verify(event1.getEventHandle()).release(false); - verify(event1.getEventHandle(), never()).release(true); - verify(event2.getEventHandle()).release(true); - verify(event2.getEventHandle(), never()).release(false); + verify(acknowledgementSet).release(eventHandle1, false); + verify(acknowledgementSet, never()).release(eventHandle1, true); + verify(acknowledgementSet).release(eventHandle2, true); + verify(acknowledgementSet, never()).release(eventHandle2, false); } @Test @@ -521,20 +550,26 @@ void output_will_release_only_new_handles_since_a_flush_when_S3_fails() throws I doNothing().when(codec).writeEvent(event, outputStream); final S3SinkService s3SinkService = createObjectUnderTest(); final List> records = generateEventRecords(1); + final List eventHandles = records.stream().map(Record::getData).map(Event::getEventHandle).map(this::castToDefaultHandle).collect(Collectors.toList()); + for (DefaultEventHandle eventHandle : eventHandles) { + eventHandle.setAcknowledgementSet(acknowledgementSet); + } s3SinkService.output(records); - final List> records2 = generateEventRecords(1); - s3SinkService.output(records2); - - final List eventHandles = records.stream().map(Record::getData).map(Event::getEventHandle).collect(Collectors.toList()); - for (EventHandle eventHandle : eventHandles) { - verify(eventHandle).release(false); + verify(acknowledgementSet).release(eventHandle, false); } - final List eventHandles2 = records2.stream().map(Record::getData).map(Event::getEventHandle).collect(Collectors.toList()); + final List> records2 = generateEventRecords(1); + final List eventHandles2 = records2.stream().map(Record::getData).map(Event::getEventHandle).map(this::castToDefaultHandle).collect(Collectors.toList()); + + for (DefaultEventHandle eventHandle : eventHandles2) { + eventHandle.setAcknowledgementSet(acknowledgementSet); + } + s3SinkService.output(records2); for (EventHandle eventHandle : eventHandles2) { - verify(eventHandle).release(false); + verify(acknowledgementSet).release(eventHandle, false); } + } private Collection> generateRandomStringEventRecord() { @@ -549,8 +584,6 @@ private List> generateEventRecords(final int numberOfRecords) { List> records = new ArrayList<>(); for (int i = 0; i < numberOfRecords; i++) { final JacksonEvent event = (JacksonEvent) JacksonEvent.fromMessage(UUID.randomUUID().toString()); - final EventHandle eventHandle = mock(EventHandle.class); - event.setEventHandle(eventHandle); records.add(new Record<>(event)); } return records; @@ -563,4 +596,4 @@ private byte[] generateByteArray() { } return bytes; } -} \ No newline at end of file +} diff --git a/data-prepper-plugins/sns-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/sns/SnsSinkServiceIT.java b/data-prepper-plugins/sns-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/sns/SnsSinkServiceIT.java index b3c235dab1..0ce6705ff0 100644 --- a/data-prepper-plugins/sns-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/sns/SnsSinkServiceIT.java +++ b/data-prepper-plugins/sns-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/sns/SnsSinkServiceIT.java @@ -17,7 +17,6 @@ import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.event.Event; -import org.opensearch.dataprepper.model.event.EventHandle; import org.opensearch.dataprepper.model.event.JacksonEvent; import org.opensearch.dataprepper.model.log.JacksonLog; import org.opensearch.dataprepper.model.plugin.PluginFactory; @@ -129,7 +128,6 @@ private Collection> setEventQueue(final int records) { private static Record createRecord() { final JacksonEvent event = JacksonLog.builder().withData("[{\"name\":\"test\"}]").build(); - event.setEventHandle(mock(EventHandle.class)); return new Record<>(event); } @@ -226,4 +224,4 @@ public void sns_sink_service_test_fail_to_push(final int recordCount) throws IOE final Map map = mapper.readValue(new String(Files.readAllBytes(Path.of(dlqFilePath))).replaceAll("(\\r|\\n)", ""), Map.class); assertThat(map.get("topic"),equalTo(topic)); } -} \ No newline at end of file +} diff --git a/data-prepper-plugins/sns-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/sns/SnsSinkService.java b/data-prepper-plugins/sns-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/sns/SnsSinkService.java index ed42d3ab4f..da29d71e17 100644 --- a/data-prepper-plugins/sns-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/sns/SnsSinkService.java +++ b/data-prepper-plugins/sns-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/sns/SnsSinkService.java @@ -121,9 +121,7 @@ void output(Collection> records) { for (Record record : records) { final Event event = record.getData(); processRecordsList.add(event); - if (event.getEventHandle() != null) { - bufferedEventHandles.add(event.getEventHandle()); - } + bufferedEventHandles.add(event.getEventHandle()); if (snsSinkConfig.getBatchSize() == processRecordsList.size()) { processRecords(); processRecordsList.clear(); From d2007bca5e568bd99079e5c5843bf84e393d78d3 Mon Sep 17 00:00:00 2001 From: Taylor Gray Date: Thu, 2 Nov 2023 11:14:41 -0500 Subject: [PATCH 22/53] Add acknowledgments for the ddb source (#3575) Signed-off-by: Taylor Gray --- .../enhanced/EnhancedSourceCoordinator.java | 3 +- .../DefaultAcknowledgementSet.java | 2 +- .../EnhancedLeaseBasedSourceCoordinator.java | 4 +- ...hancedLeaseBasedSourceCoordinatorTest.java | 2 +- .../source/dynamodb/ClientFactory.java | 21 +++++- .../source/dynamodb/DynamoDBService.java | 17 ++++- .../source/dynamodb/DynamoDBSource.java | 14 +++- .../source/dynamodb/DynamoDBSourceConfig.java | 30 +++++++- .../dynamodb/configuration/ExportConfig.java | 10 +++ .../converter/ExportRecordConverter.java | 5 +- .../dynamodb/converter/RecordConverter.java | 10 ++- .../converter/StreamRecordConverter.java | 5 +- .../partition/DataFilePartition.java | 2 +- .../dynamodb/export/DataFileCheckpointer.java | 6 +- .../dynamodb/export/DataFileLoader.java | 38 ++++++++-- .../export/DataFileLoaderFactory.java | 18 ++++- .../dynamodb/export/DataFileScheduler.java | 65 ++++++++++++----- .../dynamodb/export/ExportScheduler.java | 9 +-- .../source/dynamodb/stream/ShardConsumer.java | 33 +++++++-- .../dynamodb/stream/ShardConsumerFactory.java | 11 ++- .../source/dynamodb/stream/ShardManager.java | 3 +- .../dynamodb/stream/StreamCheckpointer.java | 8 ++- .../dynamodb/stream/StreamScheduler.java | 45 +++++++++--- .../source/dynamodb/DynamoDBServiceTest.java | 6 +- .../converter/ExportRecordConverterTest.java | 5 +- .../converter/StreamRecordConverterTest.java | 4 +- .../export/DataFileLoaderFactoryTest.java | 16 ++++- .../dynamodb/export/DataFileLoaderTest.java | 46 ++++++++++-- .../export/DataFileSchedulerTest.java | 72 +++++++++++++++++-- .../stream/ShardConsumerFactoryTest.java | 2 +- .../dynamodb/stream/ShardConsumerTest.java | 45 +++++++++++- .../dynamodb/stream/StreamSchedulerTest.java | 66 ++++++++++++++--- .../sink/opensearch/BulkRetryStrategy.java | 12 ++-- 33 files changed, 525 insertions(+), 110 deletions(-) diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/enhanced/EnhancedSourceCoordinator.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/enhanced/EnhancedSourceCoordinator.java index 9341a13e59..5423732b6f 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/enhanced/EnhancedSourceCoordinator.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/enhanced/EnhancedSourceCoordinator.java @@ -46,9 +46,10 @@ public interface EnhancedSourceCoordinator { * * @param partition The partition to be updated. * @param The progress state class + * @param ownershipTimeoutRenewal The amount of time to update ownership of the partition before another instance can acquire it. * @throws org.opensearch.dataprepper.model.source.coordinator.exceptions.PartitionUpdateException when the partition was not updated successfully */ - void saveProgressStateForPartition(EnhancedSourcePartition partition); + void saveProgressStateForPartition(EnhancedSourcePartition partition, Duration ownershipTimeoutRenewal); /** * This method is used to release the lease of a partition in the coordination store. diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java index 741f08939d..96e25fd9a8 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java @@ -139,7 +139,7 @@ public boolean release(final EventHandle eventHandle, final boolean result) { callbackFuture = executor.submit(() -> callback.accept(this.result)); return true; } else if (pendingAcknowledgments.size() == 0) { - LOG.warn("Acknowledgement set is not completed. Delaying callback until it is completed"); + LOG.debug("Acknowledgement set is not completed. Delaying callback until it is completed"); } } } finally { diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/sourcecoordination/enhanced/EnhancedLeaseBasedSourceCoordinator.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/sourcecoordination/enhanced/EnhancedLeaseBasedSourceCoordinator.java index 7ada401383..5d4fb3ae77 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/sourcecoordination/enhanced/EnhancedLeaseBasedSourceCoordinator.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/sourcecoordination/enhanced/EnhancedLeaseBasedSourceCoordinator.java @@ -127,7 +127,7 @@ public Optional acquireAvailablePartition(String partit @Override - public void saveProgressStateForPartition(EnhancedSourcePartition partition) { + public void saveProgressStateForPartition(EnhancedSourcePartition partition, final Duration ownershipTimeoutRenewal) { String partitionType = partition.getPartitionType() == null ? DEFAULT_GLOBAL_STATE_PARTITION_TYPE : partition.getPartitionType(); LOG.debug("Try to save progress for partition {} (Type {})", partition.getPartitionKey(), partitionType); @@ -140,7 +140,7 @@ public void saveProgressStateForPartition(EnhancedSourcePartition partiti final SourcePartitionStoreItem updateItem = partition.getSourcePartitionStoreItem(); // Also extend the timeout of the lease (ownership) if (updateItem.getPartitionOwnershipTimeout() != null) { - updateItem.setPartitionOwnershipTimeout(Instant.now().plus(DEFAULT_LEASE_TIMEOUT)); + updateItem.setPartitionOwnershipTimeout(Instant.now().plus(ownershipTimeoutRenewal == null ? DEFAULT_LEASE_TIMEOUT : ownershipTimeoutRenewal)); } updateItem.setPartitionProgressState(partition.convertPartitionProgressStatetoString(partition.getProgressState())); diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/sourcecoordination/enhanced/EnhancedLeaseBasedSourceCoordinatorTest.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/sourcecoordination/enhanced/EnhancedLeaseBasedSourceCoordinatorTest.java index 4e6e40c57a..6a1122ee56 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/sourcecoordination/enhanced/EnhancedLeaseBasedSourceCoordinatorTest.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/sourcecoordination/enhanced/EnhancedLeaseBasedSourceCoordinatorTest.java @@ -163,7 +163,7 @@ void test_saveProgressStateForPartition() { Optional sourcePartition = coordinator.acquireAvailablePartition(DEFAULT_PARTITION_TYPE); assertThat(sourcePartition.isPresent(), equalTo(true)); TestEnhancedSourcePartition partition = (TestEnhancedSourcePartition) sourcePartition.get(); - coordinator.saveProgressStateForPartition(partition); + coordinator.saveProgressStateForPartition(partition, null); verify(sourceCoordinationStore).tryAcquireAvailablePartition(anyString(), anyString(), any(Duration.class)); verify(sourceCoordinationStore).tryUpdateSourcePartitionItem(any(SourcePartitionStoreItem.class)); diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/ClientFactory.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/ClientFactory.java index 7d72049a6a..0704694d0c 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/ClientFactory.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/ClientFactory.java @@ -8,7 +8,10 @@ import org.opensearch.dataprepper.aws.api.AwsCredentialsOptions; import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; import org.opensearch.dataprepper.plugins.source.dynamodb.configuration.AwsAuthenticationConfig; +import org.opensearch.dataprepper.plugins.source.dynamodb.configuration.ExportConfig; import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration; +import software.amazon.awssdk.regions.Region; import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsClient; import software.amazon.awssdk.services.s3.S3Client; @@ -17,8 +20,11 @@ public class ClientFactory { private final AwsCredentialsProvider awsCredentialsProvider; private final AwsAuthenticationConfig awsAuthenticationConfig; + private final ExportConfig exportConfig; - public ClientFactory(AwsCredentialsSupplier awsCredentialsSupplier, AwsAuthenticationConfig awsAuthenticationConfig) { + public ClientFactory(final AwsCredentialsSupplier awsCredentialsSupplier, + final AwsAuthenticationConfig awsAuthenticationConfig, + final ExportConfig exportConfig) { awsCredentialsProvider = awsCredentialsSupplier.getProvider(AwsCredentialsOptions.builder() .withRegion(awsAuthenticationConfig.getAwsRegion()) .withStsRoleArn(awsAuthenticationConfig.getAwsStsRoleArn()) @@ -26,6 +32,7 @@ public ClientFactory(AwsCredentialsSupplier awsCredentialsSupplier, AwsAuthentic .withStsHeaderOverrides(awsAuthenticationConfig.getAwsStsHeaderOverrides()) .build()); this.awsAuthenticationConfig = awsAuthenticationConfig; + this.exportConfig = exportConfig; } @@ -47,8 +54,20 @@ public DynamoDbClient buildDynamoDBClient() { public S3Client buildS3Client() { return S3Client.builder() + .region(getS3ClientRegion()) .credentialsProvider(awsCredentialsProvider) + .overrideConfiguration(ClientOverrideConfiguration.builder() + .retryPolicy(retryPolicy -> retryPolicy.numRetries(5).build()) + .build()) .build(); } + private Region getS3ClientRegion() { + if (exportConfig != null && exportConfig.getAwsRegion() != null) { + return exportConfig.getAwsRegion(); + } + + return awsAuthenticationConfig.getAwsRegion(); + } + } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java index f09b5b8ded..1916a5344f 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java @@ -6,6 +6,7 @@ package org.opensearch.dataprepper.plugins.source.dynamodb; 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.event.Event; import org.opensearch.dataprepper.model.plugin.InvalidPluginConfigurationException; @@ -59,6 +60,8 @@ public class DynamoDBService { private final EnhancedSourceCoordinator coordinator; private final DynamoDbClient dynamoDbClient; + + private final DynamoDBSourceConfig dynamoDBSourceConfig; // private final DynamoDbStreamsClient dynamoDbStreamsClient; @@ -70,13 +73,21 @@ public class DynamoDBService { private final PluginMetrics pluginMetrics; + private final AcknowledgementSetManager acknowledgementSetManager; + static final Duration BUFFER_TIMEOUT = Duration.ofSeconds(60); static final int DEFAULT_BUFFER_BATCH_SIZE = 1_000; - public DynamoDBService(EnhancedSourceCoordinator coordinator, ClientFactory clientFactory, DynamoDBSourceConfig sourceConfig, PluginMetrics pluginMetrics) { + public DynamoDBService(final EnhancedSourceCoordinator coordinator, + final ClientFactory clientFactory, + final DynamoDBSourceConfig sourceConfig, + final PluginMetrics pluginMetrics, + final AcknowledgementSetManager acknowledgementSetManager) { this.coordinator = coordinator; this.pluginMetrics = pluginMetrics; + this.acknowledgementSetManager = acknowledgementSetManager; + this.dynamoDBSourceConfig = sourceConfig; // Initialize AWS clients dynamoDbClient = clientFactory.buildDynamoDBClient(); @@ -103,10 +114,10 @@ public void start(Buffer> buffer) { Runnable exportScheduler = new ExportScheduler(coordinator, dynamoDbClient, manifestFileReader, pluginMetrics); DataFileLoaderFactory loaderFactory = new DataFileLoaderFactory(coordinator, s3Client, pluginMetrics, buffer); - Runnable fileLoaderScheduler = new DataFileScheduler(coordinator, loaderFactory, pluginMetrics); + Runnable fileLoaderScheduler = new DataFileScheduler(coordinator, loaderFactory, pluginMetrics, acknowledgementSetManager, dynamoDBSourceConfig); ShardConsumerFactory consumerFactory = new ShardConsumerFactory(coordinator, dynamoDbStreamsClient, pluginMetrics, shardManager, buffer); - Runnable streamScheduler = new StreamScheduler(coordinator, consumerFactory, shardManager, pluginMetrics); + Runnable streamScheduler = new StreamScheduler(coordinator, consumerFactory, shardManager, pluginMetrics, acknowledgementSetManager, dynamoDBSourceConfig); // May consider start or shutdown the scheduler on demand // Currently, event after the exports are done, the related scheduler will not be shutdown diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSource.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSource.java index ffb1809ac6..7ecd565ec4 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSource.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSource.java @@ -7,6 +7,7 @@ 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; @@ -39,19 +40,26 @@ public class DynamoDBSource implements Source>, UsesEnhancedSource private final ClientFactory clientFactory; + private final AcknowledgementSetManager acknowledgementSetManager; + private EnhancedSourceCoordinator coordinator; private DynamoDBService dynamoDBService; @DataPrepperPluginConstructor - public DynamoDBSource(PluginMetrics pluginMetrics, final DynamoDBSourceConfig sourceConfig, final PluginFactory pluginFactory, final AwsCredentialsSupplier awsCredentialsSupplier) { + public DynamoDBSource(final PluginMetrics pluginMetrics, + final DynamoDBSourceConfig sourceConfig, + final PluginFactory pluginFactory, + final AwsCredentialsSupplier awsCredentialsSupplier, + final AcknowledgementSetManager acknowledgementSetManager) { LOG.info("Create DynamoDB Source"); this.pluginMetrics = pluginMetrics; this.sourceConfig = sourceConfig; this.pluginFactory = pluginFactory; + this.acknowledgementSetManager = acknowledgementSetManager; - clientFactory = new ClientFactory(awsCredentialsSupplier, sourceConfig.getAwsAuthenticationConfig()); + clientFactory = new ClientFactory(awsCredentialsSupplier, sourceConfig.getAwsAuthenticationConfig(), sourceConfig.getTableConfigs().get(0).getExportConfig()); } @Override @@ -61,7 +69,7 @@ public void start(Buffer> buffer) { coordinator.createPartition(new InitPartition()); // Create DynamoDB Service - dynamoDBService = new DynamoDBService(coordinator, clientFactory, sourceConfig, pluginMetrics); + dynamoDBService = new DynamoDBService(coordinator, clientFactory, sourceConfig, pluginMetrics, acknowledgementSetManager); dynamoDBService.init(); LOG.info("Start DynamoDB service"); diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSourceConfig.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSourceConfig.java index 6e4a8fe7ae..2a682ce280 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSourceConfig.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBSourceConfig.java @@ -7,10 +7,13 @@ import com.fasterxml.jackson.annotation.JsonProperty; import jakarta.validation.Valid; +import jakarta.validation.constraints.AssertTrue; import jakarta.validation.constraints.NotNull; import org.opensearch.dataprepper.plugins.source.dynamodb.configuration.AwsAuthenticationConfig; import org.opensearch.dataprepper.plugins.source.dynamodb.configuration.TableConfig; +import java.time.Duration; +import java.util.Collections; import java.util.List; /** @@ -19,14 +22,22 @@ public class DynamoDBSourceConfig { @JsonProperty("tables") - private List tableConfigs; - + private List tableConfigs = Collections.emptyList(); @JsonProperty("aws") @NotNull @Valid private AwsAuthenticationConfig awsAuthenticationConfig; + @JsonProperty("acknowledgments") + private boolean acknowledgments = false; + + @JsonProperty("shard_acknowledgment_timeout") + private Duration shardAcknowledgmentTimeout = Duration.ofMinutes(3); + + @JsonProperty("s3_data_file_acknowledgment_timeout") + private Duration dataFileAcknowledgmentTimeout = Duration.ofMinutes(5); + public DynamoDBSourceConfig() { } @@ -38,4 +49,19 @@ public List getTableConfigs() { public AwsAuthenticationConfig getAwsAuthenticationConfig() { return awsAuthenticationConfig; } + + public boolean isAcknowledgmentsEnabled() { + return acknowledgments; + } + + public Duration getShardAcknowledgmentTimeout() { + return shardAcknowledgmentTimeout; + } + + public Duration getDataFileAcknowledgmentTimeout() { return dataFileAcknowledgmentTimeout; } + + @AssertTrue(message = "Exactly one table must be configured for the DynamoDb source.") + boolean isExactlyOneTableConfigured() { + return tableConfigs.size() == 1; + } } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/configuration/ExportConfig.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/configuration/ExportConfig.java index cb3463a3b6..d34017599a 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/configuration/ExportConfig.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/configuration/ExportConfig.java @@ -7,15 +7,20 @@ import com.fasterxml.jackson.annotation.JsonProperty; import jakarta.validation.constraints.NotBlank; +import software.amazon.awssdk.regions.Region; public class ExportConfig { @JsonProperty("s3_bucket") @NotBlank(message = "Bucket Name is required for export") private String s3Bucket; + @JsonProperty("s3_prefix") private String s3Prefix; + @JsonProperty("s3_region") + private String s3Region; + public String getS3Bucket() { return s3Bucket; } @@ -23,4 +28,9 @@ public String getS3Bucket() { public String getS3Prefix() { return s3Prefix; } + + public Region getAwsRegion() { + return s3Region != null ? Region.of(s3Region) : null; + } + } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverter.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverter.java index 9f770dff69..c259ccbcaa 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverter.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverter.java @@ -10,6 +10,7 @@ import io.micrometer.core.instrument.Counter; import org.opensearch.dataprepper.buffer.common.BufferAccumulator; import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo; @@ -58,13 +59,13 @@ String getEventType() { return "EXPORT"; } - public void writeToBuffer(List lines) { + public void writeToBuffer(final AcknowledgementSet acknowledgementSet, List lines) { int eventCount = 0; for (String line : lines) { Map data = (Map) convertToMap(line).get(ITEM_KEY); try { - addToBuffer(data); + addToBuffer(acknowledgementSet, data); eventCount++; } catch (Exception e) { // will this cause too many logs? diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java index 4d5693db64..de79feadfe 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java @@ -6,6 +6,7 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.converter; import org.opensearch.dataprepper.buffer.common.BufferAccumulator; +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.event.EventMetadata; import org.opensearch.dataprepper.model.event.JacksonEvent; @@ -70,7 +71,7 @@ void flushBuffer() throws Exception { * @param eventName Event name * @throws Exception Exception if failed to write to buffer. */ - public void addToBuffer(Map data, Map keys, long eventCreationTimeMillis, String eventName) throws Exception { + public void addToBuffer(final AcknowledgementSet acknowledgementSet, Map data, Map keys, long eventCreationTimeMillis, String eventName) throws Exception { Event event = JacksonEvent.builder() .withEventType(getEventType()) .withData(data) @@ -90,13 +91,16 @@ public void addToBuffer(Map data, Map keys, long } else { eventMetadata.setAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE, partitionKey); } + if (acknowledgementSet != null) { + acknowledgementSet.add(event); + } bufferAccumulator.add(new Record<>(event)); } - public void addToBuffer(Map data) throws Exception { + public void addToBuffer(final AcknowledgementSet acknowledgementSet, Map data) throws Exception { // Export data doesn't have an event timestamp // Default to current timestamp when the event is added to buffer - addToBuffer(data, data, System.currentTimeMillis(), null); + addToBuffer(acknowledgementSet, data, data, System.currentTimeMillis(), null); } private String mapStreamEventNameToBulkAction(final String streamEventName) { diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java index 4b0bba0a4a..17c4e41e93 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java @@ -11,6 +11,7 @@ import io.micrometer.core.instrument.Counter; import org.opensearch.dataprepper.buffer.common.BufferAccumulator; import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo; import org.slf4j.Logger; @@ -53,7 +54,7 @@ String getEventType() { } - public void writeToBuffer(List records) { + public void writeToBuffer(final AcknowledgementSet acknowledgementSet, List records) { int eventCount = 0; for (Record record : records) { @@ -63,7 +64,7 @@ public void writeToBuffer(List records) { Map keys = convertKeys(record.dynamodb().keys()); try { - addToBuffer(data, keys, record.dynamodb().approximateCreationDateTime().toEpochMilli(), record.eventNameAsString()); + addToBuffer(acknowledgementSet, data, keys, record.dynamodb().approximateCreationDateTime().toEpochMilli(), record.eventNameAsString()); eventCount++; } catch (Exception e) { // will this cause too many logs? diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/coordination/partition/DataFilePartition.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/coordination/partition/DataFilePartition.java index 63222b0d22..14422c6e8c 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/coordination/partition/DataFilePartition.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/coordination/partition/DataFilePartition.java @@ -23,7 +23,7 @@ public class DataFilePartition extends EnhancedSourcePartition> bufferAccumulator = BufferAccumulator.create(builder.buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT); recordConverter = new ExportRecordConverter(bufferAccumulator, builder.tableInfo, builder.pluginMetrics); + this.acknowledgementSet = builder.acknowledgementSet; + this.dataFileAcknowledgmentTimeout = builder.dataFileAcknowledgmentTimeout; } public static Builder builder(final S3ObjectReader s3ObjectReader, final PluginMetrics pluginMetrics, final Buffer> buffer) { @@ -100,6 +105,10 @@ static class Builder { private String key; + private AcknowledgementSet acknowledgementSet; + + private Duration dataFileAcknowledgmentTimeout; + private int startLine; public Builder(final S3ObjectReader objectReader, final PluginMetrics pluginMetrics, final Buffer> buffer) { @@ -133,6 +142,16 @@ public Builder startLine(int startLine) { return this; } + public Builder acknowledgmentSet(AcknowledgementSet acknowledgementSet) { + this.acknowledgementSet = acknowledgementSet; + return this; + } + + public Builder acknowledgmentSetTimeout(Duration dataFileAcknowledgmentTimeout) { + this.dataFileAcknowledgmentTimeout = dataFileAcknowledgmentTimeout; + return this; + } + public DataFileLoader build() { return new DataFileLoader(this); } @@ -174,7 +193,7 @@ public void run() { if ((lineCount - startLine) % DEFAULT_BATCH_SIZE == 0) { // LOG.debug("Write to buffer for line " + (lineCount - DEFAULT_BATCH_SIZE) + " to " + lineCount); - recordConverter.writeToBuffer(lines); + recordConverter.writeToBuffer(acknowledgementSet, lines); lines.clear(); lastLineProcessed = lineCount; } @@ -189,17 +208,22 @@ public void run() { } if (!lines.isEmpty()) { // Do final checkpoint. - recordConverter.writeToBuffer(lines); + recordConverter.writeToBuffer(acknowledgementSet, lines); checkpointer.checkpoint(lineCount); } lines.clear(); - reader.close(); gzipInputStream.close(); inputStream.close(); - LOG.info("Complete loading s3://{}/{}", bucketName, key); - } catch (IOException e) { + + LOG.info("Completed loading s3://{}/{} to buffer", bucketName, key); + + if (acknowledgementSet != null) { + checkpointer.updateDatafileForAcknowledgmentWait(dataFileAcknowledgmentTimeout); + acknowledgementSet.complete(); + } + } catch (Exception e) { checkpointer.checkpoint(lineCount); String errorMessage = String.format("Loading of s3://%s/%s completed with Exception: %S", bucketName, key, e.getMessage()); diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactory.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactory.java index c8e5318f19..7fee416bbc 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactory.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactory.java @@ -6,6 +6,7 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.export; import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; import org.opensearch.dataprepper.model.buffer.Buffer; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.record.Record; @@ -14,6 +15,8 @@ import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo; import software.amazon.awssdk.services.s3.S3Client; +import java.time.Duration; + /** * Factory class for DataFileLoader thread. */ @@ -25,14 +28,20 @@ public class DataFileLoaderFactory { private final PluginMetrics pluginMetrics; private final Buffer> buffer; - public DataFileLoaderFactory(EnhancedSourceCoordinator coordinator, S3Client s3Client, PluginMetrics pluginMetrics, final Buffer> buffer) { + public DataFileLoaderFactory(final EnhancedSourceCoordinator coordinator, + final S3Client s3Client, + final PluginMetrics pluginMetrics, + final Buffer> buffer) { this.coordinator = coordinator; this.pluginMetrics = pluginMetrics; this.buffer = buffer; objectReader = new S3ObjectReader(s3Client); } - public Runnable createDataFileLoader(DataFilePartition dataFilePartition, TableInfo tableInfo) { + public Runnable createDataFileLoader(final DataFilePartition dataFilePartition, + final TableInfo tableInfo, + final AcknowledgementSet acknowledgementSet, + final Duration acknowledgmentTimeout) { DataFileCheckpointer checkpointer = new DataFileCheckpointer(coordinator, dataFilePartition); @@ -42,7 +51,10 @@ public Runnable createDataFileLoader(DataFilePartition dataFilePartition, TableI .key(dataFilePartition.getKey()) .tableInfo(tableInfo) .checkpointer(checkpointer) - .startLine(dataFilePartition.getProgressState().get().getLoaded()) + .acknowledgmentSet(acknowledgementSet) + .acknowledgmentSetTimeout(acknowledgmentTimeout) + // We can't checkpoint with acks enabled yet + .startLine(acknowledgementSet == null ? dataFilePartition.getProgressState().get().getLoaded() : 0) .build(); return loader; diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java index 79b99bf817..f02cf2f357 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileScheduler.java @@ -7,8 +7,11 @@ import io.micrometer.core.instrument.Counter; 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.source.coordinator.enhanced.EnhancedSourceCoordinator; import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourcePartition; +import org.opensearch.dataprepper.plugins.source.dynamodb.DynamoDBSourceConfig; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.DataFilePartition; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.GlobalState; import org.opensearch.dataprepper.plugins.source.dynamodb.model.LoadStatus; @@ -54,15 +57,26 @@ public class DataFileScheduler implements Runnable { private final PluginMetrics pluginMetrics; + private final AcknowledgementSetManager acknowledgementSetManager; + + private final DynamoDBSourceConfig dynamoDBSourceConfig; + private final Counter exportFileSuccessCounter; private final AtomicLong activeExportS3ObjectConsumersGauge; - public DataFileScheduler(EnhancedSourceCoordinator coordinator, DataFileLoaderFactory loaderFactory, PluginMetrics pluginMetrics) { + public DataFileScheduler(final EnhancedSourceCoordinator coordinator, + final DataFileLoaderFactory loaderFactory, + final PluginMetrics pluginMetrics, + final AcknowledgementSetManager acknowledgementSetManager, + final DynamoDBSourceConfig dynamoDBSourceConfig) { this.coordinator = coordinator; this.pluginMetrics = pluginMetrics; this.loaderFactory = loaderFactory; + this.acknowledgementSetManager = acknowledgementSetManager; + this.dynamoDBSourceConfig = dynamoDBSourceConfig; + executor = Executors.newFixedThreadPool(MAX_JOB_COUNT); this.exportFileSuccessCounter = pluginMetrics.counter(EXPORT_S3_OBJECTS_PROCESSED_COUNT); @@ -74,16 +88,36 @@ private void processDataFilePartition(DataFilePartition dataFilePartition) { String tableArn = getTableArn(exportArn); TableInfo tableInfo = getTableInfo(tableArn); - - Runnable loader = loaderFactory.createDataFileLoader(dataFilePartition, tableInfo); + + final boolean acknowledgmentsEnabled = dynamoDBSourceConfig.isAcknowledgmentsEnabled(); + + AcknowledgementSet acknowledgementSet = null; + if (acknowledgmentsEnabled) { + acknowledgementSet = acknowledgementSetManager.create((result) -> { + if (result == true) { + completeDataLoader(dataFilePartition).accept(null, null); + LOG.info("Received acknowledgment of completion from sink for data file {}", dataFilePartition.getKey()); + } else { + LOG.warn("Negative acknowledgment received for data file {}, retrying", dataFilePartition.getKey()); + coordinator.giveUpPartition(dataFilePartition); + } + }, dynamoDBSourceConfig.getDataFileAcknowledgmentTimeout()); + } + + Runnable loader = loaderFactory.createDataFileLoader(dataFilePartition, tableInfo, acknowledgementSet, dynamoDBSourceConfig.getDataFileAcknowledgmentTimeout()); CompletableFuture runLoader = CompletableFuture.runAsync(loader, executor); - runLoader.whenComplete(completeDataLoader(dataFilePartition)); + + if (!acknowledgmentsEnabled) { + runLoader.whenComplete(completeDataLoader(dataFilePartition)); + } else { + runLoader.whenComplete((v, ex) -> numOfWorkers.decrementAndGet()); + } numOfWorkers.incrementAndGet(); } @Override public void run() { - LOG.debug("Start running Data File Scheduler"); + LOG.debug("Starting Data File Scheduler to process S3 data files for export"); while (!Thread.currentThread().isInterrupted()) { try { @@ -145,7 +179,10 @@ private String getStreamArn(String exportArn) { private BiConsumer completeDataLoader(DataFilePartition dataFilePartition) { return (v, ex) -> { - numOfWorkers.decrementAndGet(); + + if (!dynamoDBSourceConfig.isAcknowledgmentsEnabled()) { + numOfWorkers.decrementAndGet(); + } if (ex == null) { exportFileSuccessCounter.increment(); // Update global state @@ -155,8 +192,7 @@ private BiConsumer completeDataLoader(DataFilePartition dataFilePartition) { } else { // The data loader must have already done one last checkpointing. - LOG.debug("Data Loader completed with exception"); - LOG.error("{}", ex); + LOG.error("Loading S3 data files completed with an exception: {}", ex); // Release the ownership coordinator.giveUpPartition(dataFilePartition); } @@ -190,27 +226,24 @@ private void updateState(String exportArn, int loaded) { GlobalState globalState = (GlobalState) globalPartition.get(); LoadStatus loadStatus = LoadStatus.fromMap(globalState.getProgressState().get()); - LOG.debug("Current status: total {} loaded {}", loadStatus.getTotalFiles(), loadStatus.getLoadedFiles()); - loadStatus.setLoadedFiles(loadStatus.getLoadedFiles() + 1); + LOG.info("Current status: total {} loaded {}", loadStatus.getTotalFiles(), loadStatus.getLoadedFiles()); + loadStatus.setLoadedRecords(loadStatus.getLoadedRecords() + loaded); globalState.setProgressState(loadStatus.toMap()); try { - coordinator.saveProgressStateForPartition(globalState); + coordinator.saveProgressStateForPartition(globalState, null); // if all load are completed. if (streamArn != null && loadStatus.getLoadedFiles() == loadStatus.getTotalFiles()) { - LOG.debug("All Exports are done, streaming can continue..."); + LOG.info("All Exports are done, streaming can continue..."); coordinator.createPartition(new GlobalState(streamArn, Optional.empty())); } break; } catch (Exception e) { - LOG.error("Failed to update the global status, looks like the status was out of dated, will retry.."); + LOG.error("Failed to update the global status, looks like the status was out of date, will retry.."); } - } - - } } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java index 9ea0d0dde9..9ace3aa6ae 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java @@ -139,7 +139,7 @@ public void run() { private BiConsumer completeExport(ExportPartition exportPartition) { return (status, ex) -> { if (ex != null) { - LOG.debug("Check export status for {} failed with error {}", exportPartition.getPartitionKey(), ex.getMessage()); + LOG.warn("Check export status for {} failed with error {}", exportPartition.getPartitionKey(), ex.getMessage()); // closeExportPartitionWithError(exportPartition); enhancedSourceCoordinator.giveUpPartition(exportPartition); } else { @@ -182,7 +182,7 @@ private BiConsumer completeExport(ExportPartition exportParti private void createDataFilePartitions(String exportArn, String bucketName, Map dataFileInfo) { - LOG.info("Totally {} data files generated for export {}", dataFileInfo.size(), exportArn); + LOG.info("Total of {} data files generated for export {}", dataFileInfo.size(), exportArn); AtomicInteger totalRecords = new AtomicInteger(); AtomicInteger totalFiles = new AtomicInteger(); dataFileInfo.forEach((key, size) -> { @@ -206,6 +206,7 @@ private void createDataFilePartitions(String exportArn, String bucketName, Map DEFAULT_CHECKPOINT_INTERVAL_MILLS) { - enhancedSourceCoordinator.saveProgressStateForPartition(exportPartition); + enhancedSourceCoordinator.saveProgressStateForPartition(exportPartition, null); lastCheckpointTime = System.currentTimeMillis(); } @@ -266,7 +267,7 @@ private String getOrCreateExportArn(ExportPartition exportPartition) { if (exportArn != null) { LOG.info("Export arn is " + exportArn); state.setExportArn(exportArn); - enhancedSourceCoordinator.saveProgressStateForPartition(exportPartition); + enhancedSourceCoordinator.saveProgressStateForPartition(exportPartition, null); } return exportArn; } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java index f80b145b8f..48e093f7b6 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java @@ -5,12 +5,12 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.stream; - import org.opensearch.dataprepper.buffer.common.BufferAccumulator; import org.opensearch.dataprepper.metrics.PluginMetrics; 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.acknowledgements.AcknowledgementSet; import org.opensearch.dataprepper.plugins.source.dynamodb.converter.StreamRecordConverter; import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo; import org.slf4j.Logger; @@ -79,15 +79,20 @@ public class ShardConsumer implements Runnable { private boolean waitForExport; + private final AcknowledgementSet acknowledgementSet; + + private final Duration shardAcknowledgmentTimeout; + private ShardConsumer(Builder builder) { this.dynamoDbStreamsClient = builder.dynamoDbStreamsClient; this.checkpointer = builder.checkpointer; this.shardIterator = builder.shardIterator; this.startTime = builder.startTime; this.waitForExport = builder.waitForExport; - final BufferAccumulator> bufferAccumulator = BufferAccumulator.create(builder.buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT); recordConverter = new StreamRecordConverter(bufferAccumulator, builder.tableInfo, builder.pluginMetrics); + this.acknowledgementSet = builder.acknowledgementSet; + this.shardAcknowledgmentTimeout = builder.dataFileAcknowledgmentTimeout; } public static Builder builder(final DynamoDbStreamsClient dynamoDbStreamsClient, final PluginMetrics pluginMetrics, final Buffer> buffer) { @@ -114,6 +119,9 @@ static class Builder { private boolean waitForExport; + private AcknowledgementSet acknowledgementSet; + private Duration dataFileAcknowledgmentTimeout; + public Builder(final DynamoDbStreamsClient dynamoDbStreamsClient, final PluginMetrics pluginMetrics, final Buffer> buffer) { this.dynamoDbStreamsClient = dynamoDbStreamsClient; this.pluginMetrics = pluginMetrics; @@ -145,6 +153,16 @@ public Builder waitForExport(boolean waitForExport) { return this; } + public Builder acknowledgmentSet(AcknowledgementSet acknowledgementSet) { + this.acknowledgementSet = acknowledgementSet; + return this; + } + + public Builder acknowledgmentSetTimeout(Duration dataFileAcknowledgmentTimeout) { + this.dataFileAcknowledgmentTimeout = dataFileAcknowledgmentTimeout; + return this; + } + public ShardConsumer build() { return new ShardConsumer(this); } @@ -154,7 +172,7 @@ public ShardConsumer build() { @Override public void run() { - LOG.info("Shard Consumer start to run..."); + LOG.debug("Shard Consumer start to run..."); long lastCheckpointTime = System.currentTimeMillis(); String sequenceNumber = ""; @@ -162,7 +180,7 @@ public void run() { while (!shouldStop) { if (shardIterator == null) { // End of Shard - LOG.debug("Reach end of shard"); + LOG.debug("Reached end of shard"); checkpointer.checkpoint(sequenceNumber); break; } @@ -211,7 +229,7 @@ public void run() { } else { records = response.records(); } - recordConverter.writeToBuffer(records); + recordConverter.writeToBuffer(acknowledgementSet, records); } try { // Idle between get records call. @@ -221,6 +239,11 @@ public void run() { } } + if (acknowledgementSet != null) { + checkpointer.updateShardForAcknowledgmentWait(shardAcknowledgmentTimeout); + acknowledgementSet.complete(); + } + // interrupted if (shouldStop) { // Do last checkpoint and then quit diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java index df605941e1..7cdceba3ff 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java @@ -6,6 +6,7 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.stream; import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; import org.opensearch.dataprepper.model.buffer.Buffer; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.record.Record; @@ -19,6 +20,7 @@ import org.slf4j.LoggerFactory; import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsClient; +import java.time.Duration; import java.time.Instant; import java.util.Optional; @@ -51,7 +53,9 @@ public ShardConsumerFactory(final EnhancedSourceCoordinator enhancedSourceCoordi } - public Runnable createConsumer(StreamPartition streamPartition) { + public Runnable createConsumer(final StreamPartition streamPartition, + final AcknowledgementSet acknowledgementSet, + final Duration shardAcknowledgmentTimeout) { LOG.info("Try to start a Shard Consumer for " + streamPartition.getShardId()); @@ -61,7 +65,8 @@ public Runnable createConsumer(StreamPartition streamPartition) { Instant startTime = null; boolean waitForExport = false; if (progressState.isPresent()) { - sequenceNumber = progressState.get().getSequenceNumber(); + // We can't checkpoint with acks yet + sequenceNumber = acknowledgementSet == null ? null : progressState.get().getSequenceNumber(); waitForExport = progressState.get().shouldWaitForExport(); if (progressState.get().getStartTime() != 0) { startTime = Instant.ofEpochMilli(progressState.get().getStartTime()); @@ -85,6 +90,8 @@ public Runnable createConsumer(StreamPartition streamPartition) { .shardIterator(shardIter) .startTime(startTime) .waitForExport(waitForExport) + .acknowledgmentSet(acknowledgementSet) + .acknowledgmentSetTimeout(shardAcknowledgmentTimeout) .build(); return shardConsumer; } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardManager.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardManager.java index 32ae8b573b..50c170972e 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardManager.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardManager.java @@ -129,9 +129,8 @@ public String getShardIterator(String streamArn, String shardId, String sequence .shardIteratorType(ShardIteratorType.AFTER_SEQUENCE_NUMBER) .sequenceNumber(sequenceNumber) .build(); - } else { - LOG.debug("Get Shard Iterator from beginning (TRIM_HORIZON)"); + LOG.info("Get Shard Iterator from beginning (TRIM_HORIZON) for shard {}", shardId); getShardIteratorRequest = GetShardIteratorRequest.builder() .shardId(shardId) .streamArn(streamArn) diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamCheckpointer.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamCheckpointer.java index ae807c2ca1..9e3113732a 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamCheckpointer.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamCheckpointer.java @@ -12,6 +12,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.time.Duration; import java.util.Optional; /** @@ -48,12 +49,10 @@ private void setSequenceNumber(String sequenceNumber) { * * @param sequenceNumber The last sequence number */ - public void checkpoint(String sequenceNumber) { LOG.debug("Checkpoint shard " + streamPartition.getShardId() + " with sequenceNumber " + sequenceNumber); setSequenceNumber(sequenceNumber); - coordinator.saveProgressStateForPartition(streamPartition); - + coordinator.saveProgressStateForPartition(streamPartition, null); } /** @@ -88,4 +87,7 @@ public boolean isExportDone() { return globalPartition.isPresent(); } + public void updateShardForAcknowledgmentWait(final Duration acknowledgmentSetTimeout) { + coordinator.saveProgressStateForPartition(streamPartition, acknowledgmentSetTimeout); + } } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamScheduler.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamScheduler.java index cd4d43baaf..ea799bfc24 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamScheduler.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamScheduler.java @@ -6,8 +6,11 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.stream; 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.source.coordinator.enhanced.EnhancedSourceCoordinator; import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourcePartition; +import org.opensearch.dataprepper.plugins.source.dynamodb.DynamoDBSourceConfig; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.StreamPartition; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.state.StreamProgressState; import org.slf4j.Logger; @@ -52,26 +55,53 @@ public class StreamScheduler implements Runnable { private final ShardManager shardManager; private final PluginMetrics pluginMetrics; private final AtomicLong activeChangeEventConsumers; + private final AcknowledgementSetManager acknowledgementSetManager; + private final DynamoDBSourceConfig dynamoDBSourceConfig; public StreamScheduler(final EnhancedSourceCoordinator coordinator, final ShardConsumerFactory consumerFactory, final ShardManager shardManager, - final PluginMetrics pluginMetrics) { + final PluginMetrics pluginMetrics, + final AcknowledgementSetManager acknowledgementSetManager, + final DynamoDBSourceConfig dynamoDBSourceConfig) { this.coordinator = coordinator; this.shardManager = shardManager; this.consumerFactory = consumerFactory; this.pluginMetrics = pluginMetrics; + this.acknowledgementSetManager = acknowledgementSetManager; + this.dynamoDBSourceConfig = dynamoDBSourceConfig; executor = Executors.newFixedThreadPool(MAX_JOB_COUNT); activeChangeEventConsumers = pluginMetrics.gauge(ACTIVE_CHANGE_EVENT_CONSUMERS, new AtomicLong()); } private void processStreamPartition(StreamPartition streamPartition) { - Runnable shardConsumer = consumerFactory.createConsumer(streamPartition); + final boolean acknowledgmentsEnabled = dynamoDBSourceConfig.isAcknowledgmentsEnabled(); + AcknowledgementSet acknowledgementSet = null; + + if (acknowledgmentsEnabled) { + acknowledgementSet = acknowledgementSetManager.create((result) -> { + if (result == true) { + LOG.info("Received acknowledgment of completion from sink for shard {}", streamPartition.getShardId()); + completeConsumer(streamPartition).accept(null, null); + } else { + LOG.warn("Negative acknowledgment received for shard {}, it will be retried", streamPartition.getShardId()); + coordinator.giveUpPartition(streamPartition); + } + }, dynamoDBSourceConfig.getShardAcknowledgmentTimeout()); + } + + Runnable shardConsumer = consumerFactory.createConsumer(streamPartition, acknowledgementSet, dynamoDBSourceConfig.getShardAcknowledgmentTimeout()); if (shardConsumer != null) { + CompletableFuture runConsumer = CompletableFuture.runAsync(shardConsumer, executor); - runConsumer.whenComplete(completeConsumer(streamPartition)); + + if (acknowledgmentsEnabled) { + runConsumer.whenComplete((v, ex) -> numOfWorkers.decrementAndGet()); + } else { + runConsumer.whenComplete(completeConsumer(streamPartition)); + } numOfWorkers.incrementAndGet(); } else { // If failed to create a new consumer. @@ -84,7 +114,6 @@ public void run() { LOG.debug("Stream Scheduler start to run..."); while (!Thread.currentThread().isInterrupted()) { try { - if (numOfWorkers.get() < MAX_JOB_COUNT) { final Optional sourcePartition = coordinator.acquireAvailablePartition(StreamPartition.PARTITION_TYPE); if (sourcePartition.isPresent()) { @@ -102,7 +131,7 @@ public void run() { break; } } catch (final Exception e) { - LOG.error("Received an exception while trying to get an S3 data file to process, backing off and retrying", e); + LOG.error("Received an exception while processing a shard for streams, backing off and retrying", e); try { Thread.sleep(DEFAULT_LEASE_INTERVAL_MILLIS); } catch (final InterruptedException ex) { @@ -122,7 +151,9 @@ public void run() { private BiConsumer completeConsumer(StreamPartition streamPartition) { return (v, ex) -> { - numOfWorkers.decrementAndGet(); + if (!dynamoDBSourceConfig.isAcknowledgmentsEnabled()) { + numOfWorkers.decrementAndGet(); + } if (ex == null) { LOG.info("Shard consumer for {} is completed", streamPartition.getShardId()); LOG.debug("Start creating new stream partitions for Child Shards"); @@ -147,8 +178,6 @@ private BiConsumer completeConsumer(StreamPartition streamPartition) { LOG.debug("Shard consumer completed with exception"); LOG.error(ex.toString()); coordinator.giveUpPartition(streamPartition); - - } }; } diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBServiceTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBServiceTest.java index 0032bcee92..9b436321fb 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBServiceTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBServiceTest.java @@ -11,6 +11,7 @@ import org.mockito.Mock; 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.event.Event; import org.opensearch.dataprepper.model.plugin.InvalidPluginConfigurationException; @@ -102,6 +103,9 @@ class DynamoDBServiceTest { @Mock private Buffer> buffer; + @Mock + private AcknowledgementSetManager acknowledgementSetManager; + private DynamoDBService dynamoDBService; private Collection keySchema; @@ -162,7 +166,7 @@ void setup() { } private DynamoDBService createObjectUnderTest() { - DynamoDBService objectUnderTest = new DynamoDBService(coordinator, clientFactory, sourceConfig, pluginMetrics); + DynamoDBService objectUnderTest = new DynamoDBService(coordinator, clientFactory, sourceConfig, pluginMetrics, acknowledgementSetManager); return objectUnderTest; } diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java index a31a642b8e..872e125ae1 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java @@ -31,6 +31,7 @@ import static org.hamcrest.Matchers.notNullValue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyDouble; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.BDDMockito.given; import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.times; @@ -107,7 +108,7 @@ void test_writeToBuffer() throws Exception { List data = generateData(numberOfRecords); ExportRecordConverter recordConverter = new ExportRecordConverter(bufferAccumulator, tableInfo, pluginMetrics); - recordConverter.writeToBuffer(data); + recordConverter.writeToBuffer(null, data); verify(bufferAccumulator, times(numberOfRecords)).add(any(Record.class)); verify(exportRecordSuccess).increment(anyDouble()); @@ -127,7 +128,7 @@ void test_writeSingleRecordToBuffer() throws Exception { doNothing().when(bufferAccumulator).add(recordArgumentCaptor.capture()); // doNothing().when(bufferAccumulator).flush(); - recordConverter.writeToBuffer(List.of(line)); + recordConverter.writeToBuffer(eq(null), List.of(line)); verify(bufferAccumulator).add(any(Record.class)); verify(bufferAccumulator).flush(); assertThat(recordArgumentCaptor.getValue().getData(), notNullValue()); diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java index cd02172da2..944ca4c99f 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java @@ -102,7 +102,7 @@ void test_writeToBuffer() throws Exception { StreamRecordConverter recordConverter = new StreamRecordConverter(bufferAccumulator, tableInfo, pluginMetrics); - recordConverter.writeToBuffer(records); + recordConverter.writeToBuffer(null, records); verify(bufferAccumulator, times(numberOfRecords)).add(any(Record.class)); verify(bufferAccumulator).flush(); verify(changeEventSuccessCounter).increment(anyDouble()); @@ -120,7 +120,7 @@ void test_writeSingleRecordToBuffer() throws Exception { StreamRecordConverter recordConverter = new StreamRecordConverter(bufferAccumulator, tableInfo, pluginMetrics); doNothing().when(bufferAccumulator).add(recordArgumentCaptor.capture()); - recordConverter.writeToBuffer(records); + recordConverter.writeToBuffer(null, records); verify(bufferAccumulator).add(any(Record.class)); verify(bufferAccumulator).flush(); diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactoryTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactoryTest.java index 061cb91a93..0c954740da 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactoryTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderFactoryTest.java @@ -11,6 +11,7 @@ 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.buffer.Buffer; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.record.Record; @@ -21,12 +22,14 @@ import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableMetadata; import software.amazon.awssdk.services.s3.S3Client; +import java.time.Duration; import java.util.Optional; import java.util.Random; import java.util.UUID; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.notNullValue; +import static org.mockito.Mockito.mock; @ExtendWith(MockitoExtension.class) class DataFileLoaderFactoryTest { @@ -84,10 +87,19 @@ void setup() { @Test void test_createDataFileLoader() { - DataFileLoaderFactory loaderFactory = new DataFileLoaderFactory(coordinator, s3Client, pluginMetrics, buffer); - Runnable loader = loaderFactory.createDataFileLoader(dataFilePartition, tableInfo); + Runnable loader = loaderFactory.createDataFileLoader(dataFilePartition, tableInfo, null, null); assertThat(loader, notNullValue()); + } + + @Test + void test_createDataFileLoader_with_acknowledgments() { + final AcknowledgementSet acknowledgementSet = mock(AcknowledgementSet.class); + final Duration acknowledgmentTimeout = Duration.ofSeconds(30); + DataFileLoaderFactory loaderFactory = new DataFileLoaderFactory(coordinator, s3Client, pluginMetrics, buffer); + + Runnable loader = loaderFactory.createDataFileLoader(dataFilePartition, tableInfo, acknowledgementSet, acknowledgmentTimeout); + assertThat(loader, notNullValue()); } } \ No newline at end of file diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java index 443e68f2fb..7c5c4bee56 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java @@ -14,6 +14,7 @@ import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.dataprepper.buffer.common.BufferAccumulator; import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; import org.opensearch.dataprepper.model.buffer.Buffer; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.record.Record; @@ -33,6 +34,7 @@ import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.charset.StandardCharsets; +import java.time.Duration; import java.util.Optional; import java.util.Random; import java.util.UUID; @@ -40,9 +42,11 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.BDDMockito.given; 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.times; import static org.mockito.Mockito.verify; @@ -68,7 +72,6 @@ class DataFileLoaderTest { @Mock private BufferAccumulator> bufferAccumulator; - @Mock private Counter testCounter; @@ -120,7 +123,7 @@ void setup() throws Exception { lenient().when(coordinator.createPartition(any(EnhancedSourcePartition.class))).thenReturn(true); lenient().doNothing().when(coordinator).completePartition(any(EnhancedSourcePartition.class)); - lenient().doNothing().when(coordinator).saveProgressStateForPartition(any(EnhancedSourcePartition.class)); + lenient().doNothing().when(coordinator).saveProgressStateForPartition(any(EnhancedSourcePartition.class), eq(null)); lenient().doNothing().when(coordinator).giveUpPartition(any(EnhancedSourcePartition.class)); doNothing().when(bufferAccumulator).add(any(Record.class)); @@ -166,14 +169,46 @@ private ResponseInputStream generateGzipInputStream(int numbe void test_run_loadFile_correctly() throws Exception { DataFileLoader loader; try ( - final MockedStatic bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class) - ) { + final MockedStatic bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class)) { + bufferAccumulatorMockedStatic.when(() -> BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT)).thenReturn(bufferAccumulator); + loader = DataFileLoader.builder(objectReader, pluginMetrics, buffer) + .bucketName(bucketName) + .key(manifestKey) + .checkpointer(checkpointer) + .tableInfo(tableInfo) + .build(); + } + + loader.run(); + + // Should call s3 getObject + verify(s3Client).getObject(any(GetObjectRequest.class)); + + // Should write to buffer + verify(bufferAccumulator, times(total)).add(any(Record.class)); + verify(bufferAccumulator).flush(); + + // Should do one last checkpoint when done. + verify(coordinator).saveProgressStateForPartition(any(DataFilePartition.class), eq(null)); + } + + @Test + void run_loadFile_with_acknowledgments_processes_correctly() throws Exception { + + final AcknowledgementSet acknowledgementSet = mock(AcknowledgementSet.class); + final Duration acknowledgmentTimeout = Duration.ofSeconds(30); + + DataFileLoader loader; + try ( + final MockedStatic bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class)) { bufferAccumulatorMockedStatic.when(() -> BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT)).thenReturn(bufferAccumulator); loader = DataFileLoader.builder(objectReader, pluginMetrics, buffer) .bucketName(bucketName) .key(manifestKey) .checkpointer(checkpointer) .tableInfo(tableInfo) + .acknowledgmentSet(acknowledgementSet) + .acknowledgmentSetTimeout(acknowledgmentTimeout) .build(); } @@ -187,8 +222,9 @@ void test_run_loadFile_correctly() throws Exception { verify(bufferAccumulator).flush(); // Should do one last checkpoint when done. - verify(coordinator).saveProgressStateForPartition(any(DataFilePartition.class)); + verify(coordinator).saveProgressStateForPartition(any(DataFilePartition.class), eq(null)); + verify(acknowledgementSet).complete(); } } \ No newline at end of file diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileSchedulerTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileSchedulerTest.java index 66378ee406..07dc9bfff9 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileSchedulerTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileSchedulerTest.java @@ -12,8 +12,11 @@ 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.acknowledgements.AcknowledgementSetManager; import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator; import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourcePartition; +import org.opensearch.dataprepper.plugins.source.dynamodb.DynamoDBSourceConfig; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.DataFilePartition; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.GlobalState; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.state.DataFileProgressState; @@ -21,6 +24,7 @@ import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo; import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableMetadata; +import java.time.Duration; import java.util.Optional; import java.util.UUID; import java.util.concurrent.ExecutorService; @@ -28,6 +32,7 @@ import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Consumer; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; @@ -35,7 +40,9 @@ import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.BDDMockito.given; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.opensearch.dataprepper.plugins.source.dynamodb.export.DataFileScheduler.ACTIVE_EXPORT_S3_OBJECT_CONSUMERS_GAUGE; @@ -47,6 +54,12 @@ class DataFileSchedulerTest { @Mock private EnhancedSourceCoordinator coordinator; + @Mock + private DynamoDBSourceConfig dynamoDBSourceConfig; + + @Mock + private AcknowledgementSetManager acknowledgementSetManager; + @Mock private PluginMetrics pluginMetrics; @@ -113,16 +126,18 @@ void setup() { lenient().when(coordinator.createPartition(any(EnhancedSourcePartition.class))).thenReturn(true); lenient().doNothing().when(coordinator).completePartition(any(EnhancedSourcePartition.class)); lenient().doNothing().when(coordinator).giveUpPartition(any(EnhancedSourcePartition.class)); - - lenient().when(loaderFactory.createDataFileLoader(any(DataFilePartition.class), any(TableInfo.class))).thenReturn(() -> System.out.println("Hello")); - } @Test public void test_run_DataFileLoader_correctly() throws InterruptedException { + given(loaderFactory.createDataFileLoader(any(DataFilePartition.class), any(TableInfo.class), eq(null), any(Duration.class))).willReturn(() -> System.out.println("Hello")); + + given(coordinator.acquireAvailablePartition(DataFilePartition.PARTITION_TYPE)).willReturn(Optional.of(dataFilePartition)).willReturn(Optional.empty()); + given(dynamoDBSourceConfig.isAcknowledgmentsEnabled()).willReturn(false); + given(dynamoDBSourceConfig.getDataFileAcknowledgmentTimeout()).willReturn(Duration.ofSeconds(10)); - scheduler = new DataFileScheduler(coordinator, loaderFactory, pluginMetrics); + scheduler = new DataFileScheduler(coordinator, loaderFactory, pluginMetrics, acknowledgementSetManager, dynamoDBSourceConfig); ExecutorService executorService = Executors.newSingleThreadExecutor(); final Future future = executorService.submit(() -> scheduler.run()); @@ -134,11 +149,11 @@ public void test_run_DataFileLoader_correctly() throws InterruptedException { // Should acquire data file partition verify(coordinator).acquireAvailablePartition(DataFilePartition.PARTITION_TYPE); // Should create a loader - verify(loaderFactory).createDataFileLoader(any(DataFilePartition.class), any(TableInfo.class)); + verify(loaderFactory).createDataFileLoader(any(DataFilePartition.class), any(TableInfo.class), eq(null), any(Duration.class)); // Need to call getPartition for 3 times (3 global states, 2 TableInfo) verify(coordinator, times(3)).getPartition(anyString()); // Should update global state with load status - verify(coordinator).saveProgressStateForPartition(any(GlobalState.class)); + verify(coordinator).saveProgressStateForPartition(any(GlobalState.class), eq(null)); // Should create a partition to inform streaming can start. verify(coordinator).createPartition(any(GlobalState.class)); // Should mask the partition as completed. @@ -148,14 +163,57 @@ public void test_run_DataFileLoader_correctly() throws InterruptedException { executorService.shutdownNow(); + } + @Test + void run_DataFileLoader_with_acknowledgments_enabled_processes_correctly() throws InterruptedException { + given(coordinator.acquireAvailablePartition(DataFilePartition.PARTITION_TYPE)).willReturn(Optional.of(dataFilePartition)).willReturn(Optional.empty()); + given(dynamoDBSourceConfig.isAcknowledgmentsEnabled()).willReturn(true); + + final Duration dataFileAcknowledgmentTimeout = Duration.ofSeconds(30); + given(dynamoDBSourceConfig.getDataFileAcknowledgmentTimeout()).willReturn(dataFileAcknowledgmentTimeout); + + final AcknowledgementSet acknowledgementSet = mock(AcknowledgementSet.class); + doAnswer(invocation -> { + Consumer consumer = invocation.getArgument(0); + consumer.accept(true); + return acknowledgementSet; + }).when(acknowledgementSetManager).create(any(Consumer.class), eq(dataFileAcknowledgmentTimeout)); + + given(loaderFactory.createDataFileLoader(any(DataFilePartition.class), any(TableInfo.class), eq(acknowledgementSet), eq(dataFileAcknowledgmentTimeout))).willReturn(() -> System.out.println("Hello")); + + scheduler = new DataFileScheduler(coordinator, loaderFactory, pluginMetrics, acknowledgementSetManager, dynamoDBSourceConfig); + + ExecutorService executorService = Executors.newSingleThreadExecutor(); + final Future future = executorService.submit(() -> scheduler.run()); + Thread.sleep(100); + executorService.shutdown(); + future.cancel(true); + assertThat(executorService.awaitTermination(1000, TimeUnit.MILLISECONDS), equalTo(true)); + + // Should acquire data file partition + verify(coordinator).acquireAvailablePartition(DataFilePartition.PARTITION_TYPE); + // Should create a loader + verify(loaderFactory).createDataFileLoader(any(DataFilePartition.class), any(TableInfo.class), eq(acknowledgementSet), eq(dataFileAcknowledgmentTimeout)); + // Need to call getPartition for 3 times (3 global states, 2 TableInfo) + verify(coordinator, times(3)).getPartition(anyString()); + // Should update global state with load status + verify(coordinator).saveProgressStateForPartition(any(GlobalState.class), eq(null)); + // Should create a partition to inform streaming can start. + verify(coordinator).createPartition(any(GlobalState.class)); + // Should mask the partition as completed. + verify(coordinator).completePartition(any(DataFilePartition.class)); + // Should update metrics. + verify(exportFileSuccess).increment(); + + executorService.shutdownNow(); } @Test void run_catches_exception_and_retries_when_exception_is_thrown_during_processing() throws InterruptedException { given(coordinator.acquireAvailablePartition(DataFilePartition.PARTITION_TYPE)).willThrow(RuntimeException.class); - scheduler = new DataFileScheduler(coordinator, loaderFactory, pluginMetrics); + scheduler = new DataFileScheduler(coordinator, loaderFactory, pluginMetrics, acknowledgementSetManager, dynamoDBSourceConfig); ExecutorService executorService = Executors.newSingleThreadExecutor(); final Future future = executorService.submit(() -> scheduler.run()); diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactoryTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactoryTest.java index c2f3bd1cf1..6e2a487d78 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactoryTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactoryTest.java @@ -90,7 +90,7 @@ public void test_create_shardConsumer_correctly() { ShardConsumerFactory consumerFactory = new ShardConsumerFactory(coordinator, dynamoDbStreamsClient, pluginMetrics, shardManager, buffer); - Runnable consumer = consumerFactory.createConsumer(streamPartition); + Runnable consumer = consumerFactory.createConsumer(streamPartition, null, null); assertThat(consumer, notNullValue()); } diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerTest.java index 417787becc..ec9df68189 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerTest.java @@ -14,6 +14,7 @@ import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.dataprepper.buffer.common.BufferAccumulator; import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; import org.opensearch.dataprepper.model.buffer.Buffer; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator; @@ -31,6 +32,7 @@ import software.amazon.awssdk.services.dynamodb.model.StreamRecord; import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsClient; +import java.time.Duration; import java.time.Instant; import java.util.ArrayList; import java.util.List; @@ -43,7 +45,9 @@ import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.BDDMockito.given; import static org.mockito.Mockito.doNothing; +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; @@ -96,7 +100,7 @@ class ShardConsumerTest { private final Random random = new Random(); - private final int total = random.nextInt(10); + private final int total = random.nextInt(10) + 1; @BeforeEach @@ -121,7 +125,7 @@ void setup() throws Exception { lenient().when(coordinator.createPartition(any(EnhancedSourcePartition.class))).thenReturn(true); lenient().doNothing().when(coordinator).completePartition(any(EnhancedSourcePartition.class)); - lenient().doNothing().when(coordinator).saveProgressStateForPartition(any(EnhancedSourcePartition.class)); + lenient().doNothing().when(coordinator).saveProgressStateForPartition(any(EnhancedSourcePartition.class), eq(null)); lenient().doNothing().when(coordinator).giveUpPartition(any(EnhancedSourcePartition.class)); doNothing().when(bufferAccumulator).add(any(org.opensearch.dataprepper.model.record.Record.class)); @@ -161,12 +165,47 @@ void test_run_shardConsumer_correctly() throws Exception { // Should call GetRecords verify(dynamoDbStreamsClient).getRecords(any(GetRecordsRequest.class)); + // Should write to buffer + verify(bufferAccumulator, times(total)).add(any(org.opensearch.dataprepper.model.record.Record.class)); + verify(bufferAccumulator).flush(); + // Should complete the consumer as reach to end of shard + verify(coordinator).saveProgressStateForPartition(any(StreamPartition.class), eq(null)); + } + + @Test + void test_run_shardConsumer_with_acknowledgments_correctly() throws Exception { + final AcknowledgementSet acknowledgementSet = mock(AcknowledgementSet.class); + final Duration acknowledgmentTimeout = Duration.ofSeconds(30); + + ShardConsumer shardConsumer; + try ( + final MockedStatic bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class) + ) { + bufferAccumulatorMockedStatic.when(() -> BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT)).thenReturn(bufferAccumulator); + shardConsumer = ShardConsumer.builder(dynamoDbStreamsClient, pluginMetrics, buffer) + .shardIterator(shardIterator) + .checkpointer(checkpointer) + .tableInfo(tableInfo) + .startTime(null) + .acknowledgmentSetTimeout(acknowledgmentTimeout) + .acknowledgmentSet(acknowledgementSet) + .waitForExport(false) + .build(); + } + + shardConsumer.run(); + + // Should call GetRecords + verify(dynamoDbStreamsClient).getRecords(any(GetRecordsRequest.class)); + // Should write to buffer verify(bufferAccumulator, times(total)).add(any(org.opensearch.dataprepper.model.record.Record.class)); verify(bufferAccumulator).flush(); // Should complete the consumer as reach to end of shard - verify(coordinator).saveProgressStateForPartition(any(StreamPartition.class)); + verify(coordinator).saveProgressStateForPartition(any(StreamPartition.class), eq(null)); + + verify(acknowledgementSet).complete(); } /** diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamSchedulerTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamSchedulerTest.java index a551a6052f..4c6dfdcd0d 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamSchedulerTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamSchedulerTest.java @@ -11,12 +11,16 @@ 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.acknowledgements.AcknowledgementSetManager; import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator; import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourcePartition; +import org.opensearch.dataprepper.plugins.source.dynamodb.DynamoDBSourceConfig; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.StreamPartition; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.state.StreamProgressState; import software.amazon.awssdk.services.dynamodb.streams.DynamoDbStreamsClient; +import java.time.Duration; import java.time.Instant; import java.util.List; import java.util.Optional; @@ -26,6 +30,7 @@ import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Consumer; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; @@ -33,7 +38,9 @@ import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.BDDMockito.given; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.opensearch.dataprepper.plugins.source.dynamodb.stream.StreamScheduler.ACTIVE_CHANGE_EVENT_CONSUMERS; @@ -47,6 +54,12 @@ class StreamSchedulerTest { @Mock private DynamoDbStreamsClient dynamoDbStreamsClient; + @Mock + private AcknowledgementSetManager acknowledgementSetManager; + + @Mock + private DynamoDBSourceConfig dynamoDBSourceConfig; + @Mock private ShardManager shardManager; @@ -90,10 +103,8 @@ void setup() { // Mock Coordinator methods lenient().when(coordinator.createPartition(any(EnhancedSourcePartition.class))).thenReturn(true); lenient().doNothing().when(coordinator).completePartition(any(EnhancedSourcePartition.class)); - lenient().doNothing().when(coordinator).saveProgressStateForPartition(any(EnhancedSourcePartition.class)); + lenient().doNothing().when(coordinator).saveProgressStateForPartition(any(EnhancedSourcePartition.class), eq(null)); lenient().doNothing().when(coordinator).giveUpPartition(any(EnhancedSourcePartition.class)); - - lenient().when(consumerFactory.createConsumer(any(StreamPartition.class))).thenReturn(() -> System.out.println("Hello")); lenient().when(shardManager.getChildShardIds(anyString(), anyString())).thenReturn(List.of(shardId)); when(pluginMetrics.gauge(eq(ACTIVE_CHANGE_EVENT_CONSUMERS), any(AtomicLong.class))).thenReturn(activeShardConsumers); @@ -103,9 +114,10 @@ void setup() { @Test public void test_normal_run() throws InterruptedException { - given(coordinator.acquireAvailablePartition(StreamPartition.PARTITION_TYPE)).willReturn(Optional.of(streamPartition)).willReturn(Optional.empty()); + when(consumerFactory.createConsumer(any(StreamPartition.class), eq(null), any(Duration.class))).thenReturn(() -> System.out.println("Hello")); + when(coordinator.acquireAvailablePartition(StreamPartition.PARTITION_TYPE)).thenReturn(Optional.of(streamPartition)).thenReturn(Optional.empty()); - scheduler = new StreamScheduler(coordinator, consumerFactory, shardManager, pluginMetrics); + scheduler = new StreamScheduler(coordinator, consumerFactory, shardManager, pluginMetrics, acknowledgementSetManager, dynamoDBSourceConfig); ExecutorService executorService = Executors.newSingleThreadExecutor(); @@ -117,7 +129,46 @@ public void test_normal_run() throws InterruptedException { // Should acquire the stream partition verify(coordinator).acquireAvailablePartition(StreamPartition.PARTITION_TYPE); // Should start a new consumer - verify(consumerFactory).createConsumer(any(StreamPartition.class)); + verify(consumerFactory).createConsumer(any(StreamPartition.class), eq(null), any(Duration.class)); + // Should create stream partition for child shards. + verify(coordinator).createPartition(any(StreamPartition.class)); + // Should mask the stream partition as completed. + verify(coordinator).completePartition(any(StreamPartition.class)); + + executorService.shutdownNow(); + } + + @Test + public void test_normal_run_with_acknowledgments() throws InterruptedException { + given(coordinator.acquireAvailablePartition(StreamPartition.PARTITION_TYPE)).willReturn(Optional.of(streamPartition)).willReturn(Optional.empty()); + given(dynamoDBSourceConfig.isAcknowledgmentsEnabled()).willReturn(true); + + final Duration shardAcknowledgmentTimeout = Duration.ofSeconds(30); + given(dynamoDBSourceConfig.getShardAcknowledgmentTimeout()).willReturn(shardAcknowledgmentTimeout); + + final AcknowledgementSet acknowledgementSet = mock(AcknowledgementSet.class); + doAnswer(invocation -> { + Consumer consumer = invocation.getArgument(0); + consumer.accept(true); + return acknowledgementSet; + }).when(acknowledgementSetManager).create(any(Consumer.class), eq(shardAcknowledgmentTimeout)); + + when(consumerFactory.createConsumer(any(StreamPartition.class), eq(acknowledgementSet), eq(shardAcknowledgmentTimeout))).thenReturn(() -> System.out.println("Hello")); + + scheduler = new StreamScheduler(coordinator, consumerFactory, shardManager, pluginMetrics, acknowledgementSetManager, dynamoDBSourceConfig); + + ExecutorService executorService = Executors.newSingleThreadExecutor(); + + final Future future = executorService.submit(() -> scheduler.run()); + Thread.sleep(3000); + executorService.shutdown(); + future.cancel(true); + assertThat(executorService.awaitTermination(1000, TimeUnit.MILLISECONDS), equalTo(true)); + + // Should acquire the stream partition + verify(coordinator).acquireAvailablePartition(StreamPartition.PARTITION_TYPE); + // Should start a new consumer + verify(consumerFactory).createConsumer(any(StreamPartition.class), any(AcknowledgementSet.class), any(Duration.class)); // Should create stream partition for child shards. verify(coordinator).createPartition(any(StreamPartition.class)); // Should mask the stream partition as completed. @@ -129,8 +180,7 @@ public void test_normal_run() throws InterruptedException { @Test void run_catches_exception_and_retries_when_exception_is_thrown_during_processing() throws InterruptedException { given(coordinator.acquireAvailablePartition(StreamPartition.PARTITION_TYPE)).willThrow(RuntimeException.class); - - scheduler = new StreamScheduler(coordinator, consumerFactory, shardManager, pluginMetrics); + scheduler = new StreamScheduler(coordinator, consumerFactory, shardManager, pluginMetrics, acknowledgementSetManager, dynamoDBSourceConfig); ExecutorService executorService = Executors.newSingleThreadExecutor(); final Future future = executorService.submit(() -> scheduler.run()); diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java index 451d153bfa..9f8f501fb6 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java @@ -5,21 +5,24 @@ package org.opensearch.dataprepper.plugins.sink.opensearch; -import com.linecorp.armeria.client.retry.Backoff; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.collect.ImmutableList; -import org.opensearch.dataprepper.metrics.PluginMetrics; +import com.linecorp.armeria.client.retry.Backoff; import io.micrometer.core.instrument.Counter; import org.opensearch.client.opensearch._types.OpenSearchException; import org.opensearch.client.opensearch.core.BulkRequest; import org.opensearch.client.opensearch.core.BulkResponse; import org.opensearch.client.opensearch.core.bulk.BulkResponseItem; +import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.plugins.sink.opensearch.bulk.AccumulatingBulkRequest; import org.opensearch.dataprepper.plugins.sink.opensearch.dlq.FailedBulkOperation; import org.opensearch.rest.RestStatus; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.io.IOException; +import java.time.Duration; import java.util.Arrays; import java.util.HashSet; import java.util.List; @@ -27,9 +30,6 @@ import java.util.Set; import java.util.function.BiConsumer; import java.util.function.Supplier; -import java.time.Duration; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; public final class BulkRetryStrategy { public static final String DOCUMENTS_SUCCESS = "documentsSuccess"; @@ -241,7 +241,7 @@ private BulkOperationRequestResponse handleRetriesAndFailures(final Accumulating if (e == null) { for (final BulkResponseItem bulkItemResponse : bulkResponse.items()) { if (Objects.nonNull(bulkItemResponse.error())) { - LOG.warn("operation = {}, error = {}", bulkItemResponse.operationType(), bulkItemResponse.error()); + LOG.warn("operation = {}, error = {}", bulkItemResponse.operationType(), bulkItemResponse.error().reason()); } } } From c53598f062fb3723fcab3c27fc4783d431d5d1dd Mon Sep 17 00:00:00 2001 From: David Venable Date: Fri, 3 Nov 2023 11:26:53 -0700 Subject: [PATCH 23/53] Update end-to-end tests to use the release Docker image or a custom image with a specific Java version. Resolves #3566 (#3576) Signed-off-by: David Venable --- ...per-log-analytics-basic-grok-e2e-tests.yml | 2 +- ...er-peer-forwarder-local-node-e2e-tests.yml | 2 +- ...repper-peer-forwarder-static-e2e-tests.yml | 2 +- ...ytics-raw-span-compatibility-e2e-tests.yml | 2 +- ...per-trace-analytics-raw-span-e2e-tests.yml | 2 +- ...tics-raw-span-peer-forwarder-e2e-tests.yml | 2 +- ...-trace-analytics-service-map-e2e-tests.yml | 2 +- buildSrc/build.gradle | 9 + .../gradle/end_to_end/DockerProviderTask.java | 24 ++ e2e-test/build.gradle | 65 ++-- e2e-test/log/build.gradle | 175 +++++------ e2e-test/peerforwarder/build.gradle | 214 ++++++------- .../resources/data_prepper_static.yml | 6 +- e2e-test/trace/build.gradle | 280 ++++++++++-------- .../resources/data_prepper_static.yml | 2 +- 15 files changed, 442 insertions(+), 347 deletions(-) create mode 100644 buildSrc/build.gradle create mode 100644 buildSrc/src/main/java/org/opensearch/dataprepper/gradle/end_to_end/DockerProviderTask.java diff --git a/.github/workflows/data-prepper-log-analytics-basic-grok-e2e-tests.yml b/.github/workflows/data-prepper-log-analytics-basic-grok-e2e-tests.yml index 5449ed7f6c..0ac943a231 100644 --- a/.github/workflows/data-prepper-log-analytics-basic-grok-e2e-tests.yml +++ b/.github/workflows/data-prepper-log-analytics-basic-grok-e2e-tests.yml @@ -13,7 +13,7 @@ jobs: build: strategy: matrix: - java: [11, 17, 21] + java: [11, 17, 21, docker] test: ['basicLogEndToEndTest', 'parallelGrokStringSubstituteTest'] fail-fast: false diff --git a/.github/workflows/data-prepper-peer-forwarder-local-node-e2e-tests.yml b/.github/workflows/data-prepper-peer-forwarder-local-node-e2e-tests.yml index de36390665..36f4aea9a6 100644 --- a/.github/workflows/data-prepper-peer-forwarder-local-node-e2e-tests.yml +++ b/.github/workflows/data-prepper-peer-forwarder-local-node-e2e-tests.yml @@ -13,7 +13,7 @@ jobs: build: strategy: matrix: - java: [11, 17, 21] + java: [11, 17, 21, docker] fail-fast: false runs-on: ubuntu-latest diff --git a/.github/workflows/data-prepper-peer-forwarder-static-e2e-tests.yml b/.github/workflows/data-prepper-peer-forwarder-static-e2e-tests.yml index 10efb39f01..0d81b0a615 100644 --- a/.github/workflows/data-prepper-peer-forwarder-static-e2e-tests.yml +++ b/.github/workflows/data-prepper-peer-forwarder-static-e2e-tests.yml @@ -13,7 +13,7 @@ jobs: build: strategy: matrix: - java: [11, 17, 21] + java: [11, 17, 21, docker] test: ['staticAggregateEndToEndTest', 'staticLogMetricsEndToEndTest'] fail-fast: false diff --git a/.github/workflows/data-prepper-trace-analytics-raw-span-compatibility-e2e-tests.yml b/.github/workflows/data-prepper-trace-analytics-raw-span-compatibility-e2e-tests.yml index e85d4b17ef..0dc90a9917 100644 --- a/.github/workflows/data-prepper-trace-analytics-raw-span-compatibility-e2e-tests.yml +++ b/.github/workflows/data-prepper-trace-analytics-raw-span-compatibility-e2e-tests.yml @@ -13,7 +13,7 @@ jobs: build: strategy: matrix: - java: [11, 17, 21] + java: [11, 17, 21, docker] fail-fast: false runs-on: ubuntu-latest diff --git a/.github/workflows/data-prepper-trace-analytics-raw-span-e2e-tests.yml b/.github/workflows/data-prepper-trace-analytics-raw-span-e2e-tests.yml index d7d705c93e..9538f3eb09 100644 --- a/.github/workflows/data-prepper-trace-analytics-raw-span-e2e-tests.yml +++ b/.github/workflows/data-prepper-trace-analytics-raw-span-e2e-tests.yml @@ -13,7 +13,7 @@ jobs: build: strategy: matrix: - java: [11, 17, 21] + java: [11, 17, 21, docker] otelVersion: ['0.9.0-alpha', '0.16.0-alpha'] fail-fast: false diff --git a/.github/workflows/data-prepper-trace-analytics-raw-span-peer-forwarder-e2e-tests.yml b/.github/workflows/data-prepper-trace-analytics-raw-span-peer-forwarder-e2e-tests.yml index 44291cdbf2..3cbe0d53a7 100644 --- a/.github/workflows/data-prepper-trace-analytics-raw-span-peer-forwarder-e2e-tests.yml +++ b/.github/workflows/data-prepper-trace-analytics-raw-span-peer-forwarder-e2e-tests.yml @@ -13,7 +13,7 @@ jobs: build: strategy: matrix: - java: [11, 17, 21] + java: [11, 17, 21, docker] fail-fast: false runs-on: ubuntu-latest diff --git a/.github/workflows/data-prepper-trace-analytics-service-map-e2e-tests.yml b/.github/workflows/data-prepper-trace-analytics-service-map-e2e-tests.yml index 0c5aa7eb16..6d6f7c1fc3 100644 --- a/.github/workflows/data-prepper-trace-analytics-service-map-e2e-tests.yml +++ b/.github/workflows/data-prepper-trace-analytics-service-map-e2e-tests.yml @@ -13,7 +13,7 @@ jobs: build: strategy: matrix: - java: [11, 17, 21] + java: [11, 17, 21, docker] fail-fast: false runs-on: ubuntu-latest diff --git a/buildSrc/build.gradle b/buildSrc/build.gradle new file mode 100644 index 0000000000..9d29afb944 --- /dev/null +++ b/buildSrc/build.gradle @@ -0,0 +1,9 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +plugins { + id 'java-gradle-plugin' + id 'java' +} diff --git a/buildSrc/src/main/java/org/opensearch/dataprepper/gradle/end_to_end/DockerProviderTask.java b/buildSrc/src/main/java/org/opensearch/dataprepper/gradle/end_to_end/DockerProviderTask.java new file mode 100644 index 0000000000..eb264c1fd7 --- /dev/null +++ b/buildSrc/src/main/java/org/opensearch/dataprepper/gradle/end_to_end/DockerProviderTask.java @@ -0,0 +1,24 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.gradle.end_to_end; + +import org.gradle.api.DefaultTask; +import org.gradle.api.provider.Property; +import org.gradle.api.tasks.Input; + +/** + * A task which can provide a Docker image to use for an end-to-end test. + */ +public abstract class DockerProviderTask extends DefaultTask { + /** + * The Docker image with both the name and tag in the standard string + * format - my-image:mytag + * + * @return The Docker image + */ + @Input + abstract Property getImageId(); +} diff --git a/e2e-test/build.gradle b/e2e-test/build.gradle index bc4331a09d..23578f3656 100644 --- a/e2e-test/build.gradle +++ b/e2e-test/build.gradle @@ -7,21 +7,24 @@ import com.bmuschko.gradle.docker.tasks.container.DockerCreateContainer import com.bmuschko.gradle.docker.tasks.container.DockerStartContainer import com.bmuschko.gradle.docker.tasks.container.DockerStopContainer +import com.bmuschko.gradle.docker.tasks.image.DockerBuildImage import com.bmuschko.gradle.docker.tasks.image.Dockerfile import com.bmuschko.gradle.docker.tasks.image.DockerPullImage import com.bmuschko.gradle.docker.tasks.network.DockerCreateNetwork import com.bmuschko.gradle.docker.tasks.network.DockerRemoveNetwork +import org.opensearch.dataprepper.gradle.end_to_end.DockerProviderTask plugins { id 'com.bmuschko.docker-remote-api' version '9.3.2' } + subprojects { apply plugin: 'com.bmuschko.docker-remote-api' ext { dataPrepperJarImageFilepath = 'bin/data-prepper/' - targetJavaVersion = project.hasProperty('endToEndJavaVersion') ? project.getProperty('endToEndJavaVersion') : '11' + targetJavaVersion = project.hasProperty('endToEndJavaVersion') ? project.getProperty('endToEndJavaVersion') : 'docker' targetOpenTelemetryVersion = project.hasProperty('openTelemetryVersion') ? project.getProperty('openTelemetryVersion') : "${libs.versions.opentelemetry.get()}" dataPrepperBaseImage = "eclipse-temurin:${targetJavaVersion}-jre" } @@ -46,29 +49,49 @@ subprojects { integrationTestRuntime.extendsFrom testRuntime } - task copyDataPrepperJar(type: Copy) { - dependsOn project(':data-prepper-main').jar - dependsOn project(':data-prepper-plugins').jar + tasks.register('copyDataPrepperArchive', Copy) { + dependsOn ':release:archives:linux:linuxx64DistTar' duplicatesStrategy = DuplicatesStrategy.EXCLUDE - from project(':data-prepper-main').jar.archivePath - from project(':data-prepper-main').configurations.runtimeClasspath - into("${project.buildDir}/docker/${dataPrepperJarImageFilepath}") + from project(':release:archives:linux').tasks.getByName('linuxx64DistTar').archivePath + into("${project.buildDir}/docker/") } - task createDataPrepperDockerFile(type: Dockerfile) { - dependsOn copyDataPrepperJar + tasks.register('createDataPrepperDockerFile', Dockerfile) { + dependsOn copyDataPrepperArchive + dependsOn ':release:archives:linux:linuxx64DistTar' destFile = project.file('build/docker/Dockerfile') + from(dataPrepperBaseImage) - workingDir('/app/data-prepper') - copyFile("${dataPrepperJarImageFilepath}", '/app/data-prepper/lib') - defaultCommand('java', '-Ddata-prepper.dir=/app/data-prepper', '-cp', '/app/data-prepper/lib/*', 'org.opensearch.dataprepper.DataPrepperExecute') + runCommand('mkdir -p /var/log/data-prepper') + addFile(project(':release:archives:linux').tasks.getByName('linuxx64DistTar').archiveFileName.get(), '/usr/share') + runCommand("mv /usr/share/${project(':release:archives:linux').tasks.getByName('linuxx64DistTar').archiveFileName.get().replace('.tar.gz', '')} /usr/share/data-prepper") + workingDir('/usr/share/data-prepper') + defaultCommand('bin/data-prepper') + } + + tasks.register('buildDataPrepperDockerImage', DockerBuildImage) { + dependsOn createDataPrepperDockerFile + dockerFile = file('build/docker/Dockerfile') + images.add('e2e-test-data-prepper') + } + + + tasks.register('dataPrepperDockerImage', DockerProviderTask) { + if(targetJavaVersion == 'docker') { + dependsOn ':release:docker:docker' + imageId = "${project.rootProject.name}:${project.version}" + } + else { + dependsOn 'createDataPrepperDockerFile' + imageId = buildDataPrepperDockerImage.getImageId() + } } - task createDataPrepperNetwork(type: DockerCreateNetwork) { + tasks.register('createDataPrepperNetwork', DockerCreateNetwork) { networkName = 'data_prepper_network' } - task removeDataPrepperNetwork(type: DockerRemoveNetwork) { + tasks.register('removeDataPrepperNetwork', DockerRemoveNetwork) { dependsOn createDataPrepperNetwork networkId = createDataPrepperNetwork.getNetworkId() } @@ -76,11 +99,11 @@ subprojects { /** * OpenSearch Docker tasks */ - task pullOpenSearchDockerImage(type: DockerPullImage) { + tasks.register('pullOpenSearchDockerImage', DockerPullImage) { image = "opensearchproject/opensearch:${libs.versions.opensearch.get()}" } - task createOpenSearchDockerContainer(type: DockerCreateContainer) { + tasks.register('createOpenSearchDockerContainer', DockerCreateContainer) { dependsOn createDataPrepperNetwork dependsOn pullOpenSearchDockerImage targetImageId pullOpenSearchDockerImage.image @@ -88,23 +111,23 @@ subprojects { hostConfig.portBindings = ['9200:9200', '9600:9600'] hostConfig.autoRemove = true hostConfig.network = createDataPrepperNetwork.getNetworkName() - envVars = ['discovery.type':'single-node'] + envVars = ['discovery.type': 'single-node'] } - task startOpenSearchDockerContainer(type: DockerStartContainer) { + tasks.register('startOpenSearchDockerContainer', DockerStartContainer) { dependsOn createOpenSearchDockerContainer targetContainerId createOpenSearchDockerContainer.getContainerId() doLast { - sleep(90*1000) + sleep(90 * 1000) } } - task stopOpenSearchDockerContainer(type: DockerStopContainer) { + tasks.register('stopOpenSearchDockerContainer', DockerStopContainer) { targetContainerId createOpenSearchDockerContainer.getContainerId() doLast { - sleep(5*1000) + sleep(5 * 1000) } } diff --git a/e2e-test/log/build.gradle b/e2e-test/log/build.gradle index 45624939fe..8198e41bee 100644 --- a/e2e-test/log/build.gradle +++ b/e2e-test/log/build.gradle @@ -8,120 +8,109 @@ import com.bmuschko.gradle.docker.tasks.container.DockerCreateContainer import com.bmuschko.gradle.docker.tasks.container.DockerRemoveContainer import com.bmuschko.gradle.docker.tasks.container.DockerStartContainer import com.bmuschko.gradle.docker.tasks.container.DockerStopContainer -import com.bmuschko.gradle.docker.tasks.image.DockerBuildImage /** - * End-to-end test docker network + * Represents the configurations needed for any end-to-end log test. */ - -def BASIC_GROK_PIPELINE_YAML = "basic-grok-e2e-pipeline.yml" -def PARALLEL_GROK_SUBSTITUTE_PIPELINE_YAML = "parallel-grok-substitute-e2e-pipeline.yml" -def DATA_PREPPER_CONFIG_YAML = "data_prepper.yml" - -/** - * DataPrepper Docker tasks - */ -task buildDataPrepperDockerImage(type: DockerBuildImage) { - dependsOn createDataPrepperDockerFile - dockerFile = file('build/docker/Dockerfile') - images.add('e2e-test-log-pipeline-image') +class LogTestConfiguration { + LogTestConfiguration( + String testName, + String description, + String testFilters, + String containerName, + String pipelineConfiguration, + String dataPrepperConfiguration) { + this.testName = testName + this.description = description + this.testFilters = testFilters + this.containerName = containerName + this.pipelineConfiguration = pipelineConfiguration + this.dataPrepperConfiguration = dataPrepperConfiguration + } + String testName + String description + String testFilters + String containerName + String pipelineConfiguration + String dataPrepperConfiguration } -def createDataPrepperDockerContainer(final String taskBaseName, final String dataPrepperName, final int sourcePort, - final int serverPort, final String pipelineConfigYAML, final String dataPrepperConfigYAML) { - return tasks.create("create${taskBaseName}", DockerCreateContainer) { - dependsOn buildDataPrepperDockerImage +List logTestConfigurations = [ + new LogTestConfiguration( + 'basicLogEndToEndTest', + 'Runs the basic grok end-to-end test.', + 'org.opensearch.dataprepper.integration.log.EndToEndBasicLogTest.testPipelineEndToEnd*', + 'data-prepper-basic-log', + 'basic-grok-e2e-pipeline.yml', + 'data_prepper.yml' + ), + new LogTestConfiguration( + 'parallelGrokStringSubstituteTest', + 'Runs the parallel grok and string substitute end-to-end test.', + 'org.opensearch.dataprepper.integration.log.ParallelGrokStringSubstituteLogTest.testPipelineEndToEnd*', + 'data-prepper-parallel-log', + 'parallel-grok-substitute-e2e-pipeline.yml', + 'data_prepper.yml' + ) +] + + +logTestConfigurations.each { testConfiguration -> + tasks.register("create${testConfiguration.testName}", DockerCreateContainer) { + dependsOn dataPrepperDockerImage dependsOn createDataPrepperNetwork - containerName = dataPrepperName - exposePorts("tcp", [2021, 4900]) - hostConfig.portBindings = [String.format('%d:2021', sourcePort), String.format('%d:4900', serverPort)] - hostConfig.binds = [(project.file("src/integrationTest/resources/${pipelineConfigYAML}").toString()):"/app/data-prepper/pipelines/pipelines.yaml", - (project.file("src/integrationTest/resources/${dataPrepperConfigYAML}").toString()):"/app/data-prepper/config/data-prepper-config.yaml"] + containerName = testConfiguration.containerName + exposePorts('tcp', [2021, 4900]) + hostConfig.portBindings = ['2021:2021', '4900:4900'] + hostConfig.binds = [ + (project.file("src/integrationTest/resources/${testConfiguration.pipelineConfiguration}").toString()) : '/usr/share/data-prepper/pipelines/log-pipeline.yaml', + (project.file("src/integrationTest/resources/${testConfiguration.dataPrepperConfiguration}").toString()): '/usr/share/data-prepper/config/data-prepper-config.yaml' + ] hostConfig.network = createDataPrepperNetwork.getNetworkName() - cmd = ['java', '-Ddata-prepper.dir=/app/data-prepper', '-cp', '/app/data-prepper/lib/*', 'org.opensearch.dataprepper.DataPrepperExecute'] - targetImageId buildDataPrepperDockerImage.getImageId() + targetImageId dataPrepperDockerImage.imageId } -} -def startDataPrepperDockerContainer(final DockerCreateContainer createDataPrepperDockerContainerTask) { - return tasks.create("start${createDataPrepperDockerContainerTask.getName()}", DockerStartContainer) { - dependsOn createDataPrepperDockerContainerTask - targetContainerId createDataPrepperDockerContainerTask.getContainerId() + tasks.register("start${testConfiguration.testName}", DockerStartContainer) { + dependsOn "create${testConfiguration.testName}" + dependsOn 'startOpenSearchDockerContainer' + mustRunAfter 'startOpenSearchDockerContainer' + targetContainerId tasks.getByName("create${testConfiguration.testName}").getContainerId() } -} -def stopDataPrepperDockerContainer(final DockerStartContainer startDataPrepperDockerContainerTask) { - return tasks.create("stop${startDataPrepperDockerContainerTask.getName()}", DockerStopContainer) { - targetContainerId startDataPrepperDockerContainerTask.getContainerId() + tasks.register("stop${testConfiguration.testName}", DockerStopContainer) { + dependsOn "${testConfiguration.testName}" + targetContainerId tasks.getByName("create${testConfiguration.testName}").getContainerId() } -} -def removeDataPrepperDockerContainer(final DockerStopContainer stopDataPrepperDockerContainerTask) { - return tasks.create("remove${stopDataPrepperDockerContainerTask.getName()}", DockerRemoveContainer) { - targetContainerId stopDataPrepperDockerContainerTask.getContainerId() + tasks.register("remove${testConfiguration.testName}", DockerRemoveContainer) { + dependsOn "stop${testConfiguration.testName}" + targetContainerId tasks.getByName("stop${testConfiguration.testName}").getContainerId() } -} -/** - * End to end test. Spins up OpenSearch and DataPrepper docker containers, then runs the integ test - * Stops the docker containers when finished - */ -task basicLogEndToEndTest(type: Test) { - dependsOn build - dependsOn startOpenSearchDockerContainer - def createDataPrepperTask = createDataPrepperDockerContainer( - "basicLogDataPrepper", "dataprepper", 2021, 4900, "${BASIC_GROK_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_YAML}") - def startDataPrepperTask = startDataPrepperDockerContainer(createDataPrepperTask as DockerCreateContainer) - dependsOn startDataPrepperTask - startDataPrepperTask.mustRunAfter 'startOpenSearchDockerContainer' - // wait for data-preppers to be ready - doFirst { - sleep(15*1000) - } - - description = 'Runs the basic grok end-to-end test.' - group = 'verification' - testClassesDirs = sourceSets.integrationTest.output.classesDirs - classpath = sourceSets.integrationTest.runtimeClasspath + tasks.register(testConfiguration.testName, Test) { + dependsOn build + dependsOn startOpenSearchDockerContainer + dependsOn "start${testConfiguration.testName}" - filter { - includeTestsMatching "org.opensearch.dataprepper.integration.log.EndToEndBasicLogTest.testPipelineEndToEnd*" - } - - finalizedBy stopOpenSearchDockerContainer - def stopDataPrepperTask = stopDataPrepperDockerContainer(startDataPrepperTask as DockerStartContainer) - finalizedBy stopDataPrepperTask - finalizedBy removeDataPrepperDockerContainer(stopDataPrepperTask as DockerStopContainer) - finalizedBy removeDataPrepperNetwork -} + // Wait for Data Prepper image to be ready + doFirst { + sleep(15 * 1000) + } -task parallelGrokStringSubstituteTest(type: Test) { - dependsOn build - dependsOn startOpenSearchDockerContainer - def createDataPrepperTask = createDataPrepperDockerContainer( - "ParallelGrokSubstLogDataPrepper", "dataprepper-pgsts-test", 2021, 4900, "${PARALLEL_GROK_SUBSTITUTE_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_YAML}") - def startDataPrepperTask = startDataPrepperDockerContainer(createDataPrepperTask as DockerCreateContainer) - dependsOn startDataPrepperTask - startDataPrepperTask.mustRunAfter 'startOpenSearchDockerContainer' - // wait for data-preppers to be ready - doFirst { - sleep(15*1000) - } + description = testConfiguration.description + group = 'verification' + testClassesDirs = sourceSets.integrationTest.output.classesDirs + classpath = sourceSets.integrationTest.runtimeClasspath - description = 'Runs the parallel grok and string substitute end-to-end test.' - group = 'verification' - testClassesDirs = sourceSets.integrationTest.output.classesDirs - classpath = sourceSets.integrationTest.runtimeClasspath + filter { + includeTestsMatching testConfiguration.testFilters + } - filter { - includeTestsMatching "org.opensearch.dataprepper.integration.log.ParallelGrokStringSubstituteLogTest.testPipelineEndToEnd*" + finalizedBy stopOpenSearchDockerContainer + finalizedBy "remove${testConfiguration.testName}" + finalizedBy removeDataPrepperNetwork } - finalizedBy stopOpenSearchDockerContainer - def stopDataPrepperTask = stopDataPrepperDockerContainer(startDataPrepperTask as DockerStartContainer) - finalizedBy stopDataPrepperTask - finalizedBy removeDataPrepperDockerContainer(stopDataPrepperTask as DockerStopContainer) - finalizedBy removeDataPrepperNetwork } dependencies { diff --git a/e2e-test/peerforwarder/build.gradle b/e2e-test/peerforwarder/build.gradle index fd017b03fc..2754f4ab34 100644 --- a/e2e-test/peerforwarder/build.gradle +++ b/e2e-test/peerforwarder/build.gradle @@ -8,139 +8,155 @@ import com.bmuschko.gradle.docker.tasks.container.DockerCreateContainer import com.bmuschko.gradle.docker.tasks.container.DockerRemoveContainer import com.bmuschko.gradle.docker.tasks.container.DockerStartContainer import com.bmuschko.gradle.docker.tasks.container.DockerStopContainer -import com.bmuschko.gradle.docker.tasks.image.DockerBuildImage /** * End-to-end test docker network */ -def AGGREGATE_PIPELINE_YAML = "aggregate-e2e-pipeline.yml" -def LOG_METRICS_PIPELINE_YAML = "log-metrics-pipeline.yml" -def DATA_PREPPER_CONFIG_LOCAL_NODE = "data_prepper_local_node.yml" -def DATA_PREPPER_CONFIG_STATIC = "data_prepper_static.yml" +def AGGREGATE_PIPELINE_YAML = 'aggregate-e2e-pipeline.yml' +def DATA_PREPPER_CONFIG_LOCAL_NODE = 'data_prepper_local_node.yml' +def DATA_PREPPER_CONFIG_STATIC = 'data_prepper_static.yml' + /** - * DataPrepper Docker tasks + * Represents the configurations needed for peer-forwarder end-to-end tests. */ - -task buildDataPrepperDockerImage(type: DockerBuildImage) { - dependsOn createDataPrepperDockerFile - dockerFile = file('build/docker/Dockerfile') - images.add('integ-test-pipeline-image') -} - -def createDataPrepperDockerContainer(final String taskBaseName, final String dataPrepperName, final int sourcePort, - final String pipelineConfigYAML, final String dataPrepperConfigYAML) { - return tasks.create("create${taskBaseName}", DockerCreateContainer) { - dependsOn buildDataPrepperDockerImage - dependsOn createDataPrepperNetwork - containerName = dataPrepperName - exposePorts('tcp', [2021]) - hostConfig.portBindings = [String.format('%d:2021', sourcePort)] - hostConfig.network = createDataPrepperNetwork.getNetworkName() - hostConfig.binds = [ - (project.file("src/integrationTest/resources/${pipelineConfigYAML}").toString()) : "/app/data-prepper/pipelines/pipelines.yaml", - "/tmp" : "/tmp", - (project.file("src/integrationTest/resources/${dataPrepperConfigYAML}").toString()): "/app/data-prepper/config/data-prepper-config.yaml", - (project.file('src/integrationTest/resources/default_certificate.pem').toString()) : '/app/data-prepper/config/default_certificate.pem', - (project.file('src/integrationTest/resources/default_private_key.pem').toString()): '/app/data-prepper/config/default_private_key.pem' - ] - cmd = ['java', '-Ddata-prepper.dir=/app/data-prepper', '-cp', '/app/data-prepper/lib/*', 'org.opensearch.dataprepper.DataPrepperExecute'] - targetImageId buildDataPrepperDockerImage.getImageId() +class PeerForwarderTestConfiguration { + PeerForwarderTestConfiguration( + String testName, + String description, + String testFilters, + String containerName, + String pipelineConfiguration, + String dataPrepperConfiguration) { + this.testName = testName + this.description = description + this.testFilters = testFilters + this.containerName = containerName + this.pipelineConfiguration = pipelineConfiguration + this.dataPrepperConfiguration = dataPrepperConfiguration } + String testName + String description + String testFilters + String containerName + String pipelineConfiguration + String dataPrepperConfiguration } -def startDataPrepperDockerContainer(final DockerCreateContainer createDataPrepperDockerContainerTask) { - return tasks.create("start${createDataPrepperDockerContainerTask.getName()}", DockerStartContainer) { - dependsOn createDataPrepperDockerContainerTask - targetContainerId createDataPrepperDockerContainerTask.getContainerId() - } -} +List peerForwarderTestConfigurations = [ + new PeerForwarderTestConfiguration( + 'localAggregateEndToEndTest', + 'Runs the local log aggregation end-to-end test.', + 'org.opensearch.dataprepper.integration.peerforwarder.EndToEndPeerForwarderTest.testAggregatePipelineWithSingleNodeEndToEnd*', + 'data-prepper', + AGGREGATE_PIPELINE_YAML, + DATA_PREPPER_CONFIG_LOCAL_NODE + ), + new PeerForwarderTestConfiguration( + 'staticAggregateEndToEndTest', + 'Runs the local log aggregation end-to-end test.', + 'org.opensearch.dataprepper.integration.peerforwarder.EndToEndPeerForwarderTest.testAggregatePipelineWithMultipleNodesEndToEnd*', + 'node.data-prepper.example.com', + AGGREGATE_PIPELINE_YAML, + DATA_PREPPER_CONFIG_STATIC + ), + new PeerForwarderTestConfiguration( + 'staticLogMetricsEndToEndTest', + 'Runs the local log aggregation end-to-end test.', + 'org.opensearch.dataprepper.integration.peerforwarder.EndToEndLogMetricsTest.testLogMetricsPipelineWithMultipleNodesEndToEnd*', + 'node.data-prepper.example.com', + 'log-metrics-pipeline.yml', + DATA_PREPPER_CONFIG_STATIC + ) +] -def stopDataPrepperDockerContainer(final DockerStartContainer startDataPrepperDockerContainerTask) { - return tasks.create("stop${startDataPrepperDockerContainerTask.getName()}", DockerStopContainer) { - targetContainerId startDataPrepperDockerContainerTask.getContainerId() - } +/** + * Creates a container name from a base name. If the name looks like a DNS entry + * it adds the index after the first part. Otherwise, it appends the index to the end. + * + * @param name The base container name + * @param index The index of the container to create + * @return The container name to use. + */ +static def createContainerName(String name, int index) { + def nameParts = name.split('\\.', 2) + + return "${nameParts[0]}${index}${nameParts.length==2 ? '.' + nameParts[1] : ''}" } -def removeDataPrepperDockerContainer(final DockerStopContainer stopDataPrepperDockerContainerTask) { - return tasks.create("remove${stopDataPrepperDockerContainerTask.getName()}", DockerRemoveContainer) { - targetContainerId stopDataPrepperDockerContainerTask.getContainerId() +peerForwarderTestConfigurations.each { testConfiguration -> + tasks.register("start${testConfiguration.testName}All") + tasks.register("remove${testConfiguration.testName}All") + + (0..1).each { containerIndex -> + tasks.register("create${testConfiguration.testName}${containerIndex}", DockerCreateContainer) { + dependsOn dataPrepperDockerImage + dependsOn createDataPrepperNetwork + containerName = createContainerName(testConfiguration.containerName, containerIndex) + exposePorts('tcp', [2021]) + hostConfig.portBindings = ["${2021+containerIndex}:2021"] + hostConfig.binds = [ + (project.file("src/integrationTest/resources/${testConfiguration.pipelineConfiguration}").toString()) : '/usr/share/data-prepper/pipelines/test-pipeline.yaml', + (project.file("src/integrationTest/resources/${testConfiguration.dataPrepperConfiguration}").toString()): '/usr/share/data-prepper/config/data-prepper-config.yaml', + (project.file('src/integrationTest/resources/default_certificate.pem').toString()) : '/usr/share/data-prepper/config/default_certificate.pem', + (project.file('src/integrationTest/resources/default_private_key.pem').toString()): '/usr/share/data-prepper/config/default_private_key.pem' + ] + hostConfig.network = createDataPrepperNetwork.getNetworkName() + targetImageId dataPrepperDockerImage.imageId + } + + tasks.register("start${testConfiguration.testName}${containerIndex}", DockerStartContainer) { + dependsOn "create${testConfiguration.testName}${containerIndex}" + dependsOn 'startOpenSearchDockerContainer' + mustRunAfter 'startOpenSearchDockerContainer' + targetContainerId tasks.getByName("create${testConfiguration.testName}${containerIndex}").getContainerId() + } + + tasks.named("start${testConfiguration.testName}All").configure { + dependsOn "start${testConfiguration.testName}${containerIndex}" + } + + tasks.register("stop${testConfiguration.testName}${containerIndex}", DockerStopContainer) { + dependsOn "${testConfiguration.testName}" + targetContainerId tasks.getByName("create${testConfiguration.testName}${containerIndex}").getContainerId() + } + + tasks.register("remove${testConfiguration.testName}${containerIndex}", DockerRemoveContainer) { + dependsOn "stop${testConfiguration.testName}${containerIndex}" + targetContainerId tasks.getByName("stop${testConfiguration.testName}${containerIndex}").getContainerId() + } + + tasks.named("remove${testConfiguration.testName}All").configure { + dependsOn "remove${testConfiguration.testName}${containerIndex}" + } } -} -/** - * End to end test. Spins up OpenSearch and DataPrepper docker containers, then runs the integ test - * Stops the docker containers when finished - */ -def createEndToEndTest(final String testName, final String includeTestsMatchPattern, - final DockerCreateContainer createDataPrepper1Task, final DockerCreateContainer createDataPrepper2Task) { - return tasks.create(testName, Test) { + tasks.register(testConfiguration.testName, Test) { dependsOn build dependsOn startOpenSearchDockerContainer - def startDataPrepper1Task = startDataPrepperDockerContainer(createDataPrepper1Task as DockerCreateContainer) - def startDataPrepper2Task = startDataPrepperDockerContainer(createDataPrepper2Task as DockerCreateContainer) - dependsOn startDataPrepper1Task - dependsOn startDataPrepper2Task - startDataPrepper1Task.mustRunAfter 'startOpenSearchDockerContainer' - startDataPrepper2Task.mustRunAfter 'startOpenSearchDockerContainer' - // wait for data-preppers to be ready + dependsOn "start${testConfiguration.testName}All" + + // Wait for Data Prepper image to be ready doFirst { sleep(15 * 1000) } - description = 'Runs the raw span integration tests.' + description = testConfiguration.description group = 'verification' testClassesDirs = sourceSets.integrationTest.output.classesDirs classpath = sourceSets.integrationTest.runtimeClasspath filter { - includeTestsMatching includeTestsMatchPattern + includeTestsMatching testConfiguration.testFilters } finalizedBy stopOpenSearchDockerContainer - def stopDataPrepper1Task = stopDataPrepperDockerContainer(startDataPrepper1Task as DockerStartContainer) - def stopDataPrepper2Task = stopDataPrepperDockerContainer(startDataPrepper2Task as DockerStartContainer) - finalizedBy stopDataPrepper1Task - finalizedBy stopDataPrepper2Task - finalizedBy removeDataPrepperDockerContainer(stopDataPrepper1Task as DockerStopContainer) - finalizedBy removeDataPrepperDockerContainer(stopDataPrepper2Task as DockerStopContainer) + finalizedBy "remove${testConfiguration.testName}All" finalizedBy removeDataPrepperNetwork } } -// Discovery mode: LOCAL_NODE -def includeLocalAggregateTestsMatchPattern = "org.opensearch.dataprepper.integration.peerforwarder.EndToEndPeerForwarderTest.testAggregatePipelineWithSingleNodeEndToEnd*" - -def createLocalAggregateDataPrepper1Task = createDataPrepperDockerContainer( - "localAggregateDataPrepper1", "dataprepper1", 2021, "${AGGREGATE_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_LOCAL_NODE}") -def createLocalAggregateDataPrepper2Task = createDataPrepperDockerContainer( - "localAggregateDataPrepper2", "dataprepper2", 2022, "${AGGREGATE_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_LOCAL_NODE}") - -def localAggregateEndToEndTest = createEndToEndTest("localAggregateEndToEndTest", includeLocalAggregateTestsMatchPattern, - createLocalAggregateDataPrepper1Task, createLocalAggregateDataPrepper2Task) - -// Discovery mode: STATIC with SSL & mTLS -def includeStaticAggregateTestsMatchPattern = "org.opensearch.dataprepper.integration.peerforwarder.EndToEndPeerForwarderTest.testAggregatePipelineWithMultipleNodesEndToEnd*" - -def createAggregateDataPrepper1Task = createDataPrepperDockerContainer( - "staticAggregateDataPrepper1", "node1.data-prepper.example.com", 2021, "${AGGREGATE_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_STATIC}") -def createAggregateDataPrepper2Task = createDataPrepperDockerContainer( - "staticAggregateDataPrepper2", "node2.data-prepper.example.com", 2022, "${AGGREGATE_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_STATIC}") - -def staticAggregateEndToEndTest = createEndToEndTest("staticAggregateEndToEndTest", includeStaticAggregateTestsMatchPattern, - createAggregateDataPrepper1Task, createAggregateDataPrepper2Task) - -// Discovery mode: STATIC - Log pipeline metrics e2e test -def includeLocalLogMetricsTestsMatchPattern = "org.opensearch.dataprepper.integration.peerforwarder.EndToEndLogMetricsTest.testLogMetricsPipelineWithMultipleNodesEndToEnd*" - -def createStaticLogMetricsDataPrepper1Task = createDataPrepperDockerContainer( - "staticLogMetricsDataPrepper1", "node1.data-prepper.example.com", 2021, "${LOG_METRICS_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_STATIC}") -def createStaticLogMetricsDataPrepper2Task = createDataPrepperDockerContainer( - "staticLogMetricsDataPrepper2", "node2.data-prepper.example.com", 2022, "${LOG_METRICS_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_STATIC}") - -def staticLogMetricsEndToEndTest = createEndToEndTest("staticLogMetricsEndToEndTest", includeLocalLogMetricsTestsMatchPattern, - createStaticLogMetricsDataPrepper1Task, createStaticLogMetricsDataPrepper2Task) dependencies { integrationTestImplementation project(':data-prepper-api') diff --git a/e2e-test/peerforwarder/src/integrationTest/resources/data_prepper_static.yml b/e2e-test/peerforwarder/src/integrationTest/resources/data_prepper_static.yml index 910ef43191..e16b43b45c 100644 --- a/e2e-test/peerforwarder/src/integrationTest/resources/data_prepper_static.yml +++ b/e2e-test/peerforwarder/src/integrationTest/resources/data_prepper_static.yml @@ -2,9 +2,9 @@ ssl: false peer_forwarder: port: 4994 ssl: true - ssl_certificate_file: "/app/data-prepper/config/default_certificate.pem" - ssl_key_file: "/app/data-prepper/config/default_private_key.pem" + ssl_certificate_file: "/usr/share/data-prepper/config/default_certificate.pem" + ssl_key_file: "/usr/share/data-prepper/config/default_private_key.pem" discovery_mode: static - static_endpoints: ["node1.data-prepper.example.com", "node2.data-prepper.example.com"] + static_endpoints: ["node0.data-prepper.example.com", "node1.data-prepper.example.com"] authentication: mutual_tls: diff --git a/e2e-test/trace/build.gradle b/e2e-test/trace/build.gradle index c0ff41e62f..b83d977973 100644 --- a/e2e-test/trace/build.gradle +++ b/e2e-test/trace/build.gradle @@ -8,163 +8,197 @@ import com.bmuschko.gradle.docker.tasks.container.DockerCreateContainer import com.bmuschko.gradle.docker.tasks.container.DockerRemoveContainer import com.bmuschko.gradle.docker.tasks.container.DockerStartContainer import com.bmuschko.gradle.docker.tasks.container.DockerStopContainer -import com.bmuschko.gradle.docker.tasks.image.DockerBuildImage import com.bmuschko.gradle.docker.tasks.image.DockerPullImage +import org.opensearch.dataprepper.gradle.end_to_end.DockerProviderTask /** - * End-to-end test docker network + * Represents a configuration for a trace test container. */ - -def RAW_SPAN_PIPELINE_YAML = "raw-span-e2e-pipeline.yml" -def SERVICE_MAP_PIPELINE_YAML = "service-map-e2e-pipeline.yml" -def DATA_PREPPER_CONFIG_YAML = "data_prepper.yml" -def DATA_PREPPER_CONFIG_STATIC_YAML = "data_prepper_static.yml" -def RAW_SPAN_PIPELINE_FROM_BUILD_YAML = "raw-span-e2e-pipeline-from-build.yml" -def RAW_SPAN_PIPELINE_LATEST_RELEASE_YAML = "raw-span-e2e-pipeline-latest-release.yml" +class TraceTestContainerConfiguration { + TraceTestContainerConfiguration( + TaskProvider imageTask, + String pipelineConfiguration, + String dataPrepperConfiguration) { + this.imageTask = imageTask + this.pipelineConfiguration = pipelineConfiguration + this.dataPrepperConfiguration = dataPrepperConfiguration + } + TaskProvider imageTask + String pipelineConfiguration + String dataPrepperConfiguration +} /** - * DataPrepper Docker tasks + * Represents the configurations needed for any end-to-end trace test. */ -task buildDataPrepperDockerImage(type: DockerBuildImage) { - dependsOn createDataPrepperDockerFile - dockerFile = file('build/docker/Dockerfile') - images.add('integ-test-pipeline-image') -} - -def createDataPrepperDockerContainer(final String taskBaseName, final String dataPrepperName, final int grpcPort, - final String pipelineConfigYAML, final String dataPrepperConfigYAML) { - return tasks.create("create${taskBaseName}", DockerCreateContainer) { - dependsOn buildDataPrepperDockerImage - dependsOn createDataPrepperNetwork - containerName = dataPrepperName - exposePorts("tcp", [21890]) - hostConfig.portBindings = [String.format('%d:21890', grpcPort)] - hostConfig.network = createDataPrepperNetwork.getNetworkName() - hostConfig.binds = [(project.file("src/integrationTest/resources/${pipelineConfigYAML}").toString()):"/app/data-prepper/pipelines/pipelines.yaml", - (project.file("src/integrationTest/resources/${dataPrepperConfigYAML}").toString()):"/app/data-prepper/config/data-prepper-config.yaml"] - cmd = ['java', '-Ddata-prepper.dir=/app/data-prepper', '-cp', '/app/data-prepper/lib/*', 'org.opensearch.dataprepper.DataPrepperExecute'] - targetImageId buildDataPrepperDockerImage.getImageId() +class TraceTestConfiguration { + TraceTestConfiguration( + String testName, + String description, + String testFilters, + List containerConfigurations, + String containerNamePrefix) { + this.testName = testName + this.description = description + this.testFilters = testFilters + this.containerConfigurations = containerConfigurations + this.containerNamePrefix = containerNamePrefix } + String testName + String description + String testFilters + List containerConfigurations + String containerNamePrefix } -def startDataPrepperDockerContainer(final DockerCreateContainer createDataPrepperDockerContainerTask) { - return tasks.create("start${createDataPrepperDockerContainerTask.getName()}", DockerStartContainer) { - dependsOn createDataPrepperDockerContainerTask - targetContainerId createDataPrepperDockerContainerTask.getContainerId() - } -} +def RAW_SPAN_PIPELINE_YAML = 'raw-span-e2e-pipeline.yml' +def DATA_PREPPER_CONFIG_YAML = 'data_prepper.yml' +def RAW_SPAN_TESTS_PATTERN = 'org.opensearch.dataprepper.integration.trace.EndToEndRawSpanTest.testPipelineEndToEnd*' +def DATA_PREPPER_CONFIG_STATIC_YAML = 'data_prepper_static.yml' +def RELEASED_DATA_PREPPER_DOCKER_IMAGE = 'opensearchproject/data-prepper:latest' -def stopDataPrepperDockerContainer(final DockerStartContainer startDataPrepperDockerContainerTask) { - return tasks.create("stop${startDataPrepperDockerContainerTask.getName()}", DockerStopContainer) { - targetContainerId startDataPrepperDockerContainerTask.getContainerId() - } -} -def removeDataPrepperDockerContainer(final DockerStopContainer stopDataPrepperDockerContainerTask) { - return tasks.create("remove${stopDataPrepperDockerContainerTask.getName()}", DockerRemoveContainer) { - targetContainerId stopDataPrepperDockerContainerTask.getContainerId() - } +tasks.register('pullDataPrepperDockerImage', DockerPullImage) { + image = RELEASED_DATA_PREPPER_DOCKER_IMAGE } -task pullDataPrepperDockerImage(type: DockerPullImage) { - image = 'opensearchproject/data-prepper:latest' +tasks.register('latestDataPrepperDockerImage', DockerProviderTask) { + dependsOn 'pullDataPrepperDockerImage' + imageId = RELEASED_DATA_PREPPER_DOCKER_IMAGE } -def createDataPrepperDockerContainerFromPullImage(final String taskBaseName, final String dataPrepperName, final int grpcPort, - final String pipelineConfigYAML, final String dataPrepperConfigYAML) { - return tasks.create("create${taskBaseName}", DockerCreateContainer) { - dependsOn createDataPrepperNetwork - dependsOn pullDataPrepperDockerImage - containerName = dataPrepperName - hostConfig.portBindings = [String.format('%d:21890', grpcPort)] - exposePorts('tcp', [21890]) - hostConfig.network = createDataPrepperNetwork.getNetworkName() - hostConfig.binds = [(project.file("src/integrationTest/resources/${pipelineConfigYAML}").toString()):"/usr/share/data-prepper/pipelines.yaml", - (project.file("src/integrationTest/resources/${dataPrepperConfigYAML}").toString()):"/usr/share/data-prepper/data-prepper-config.yaml"] - targetImageId pullDataPrepperDockerImage.image + +List traceTestConfigurations = [ + new TraceTestConfiguration( + 'rawSpanEndToEndTest', + 'Runs the raw span integration tests.', + RAW_SPAN_TESTS_PATTERN, + [new TraceTestContainerConfiguration( + tasks.named('dataPrepperDockerImage'), + RAW_SPAN_PIPELINE_YAML, + DATA_PREPPER_CONFIG_YAML), + new TraceTestContainerConfiguration( + tasks.named('dataPrepperDockerImage'), + RAW_SPAN_PIPELINE_YAML, + DATA_PREPPER_CONFIG_YAML)], + 'data-prepper-raw', + ), + new TraceTestConfiguration( + 'rawSpanPeerForwarderEndToEndTest', + 'Runs the raw span with peer-forwarder integration tests.', + RAW_SPAN_TESTS_PATTERN, + [new TraceTestContainerConfiguration( + tasks.named('dataPrepperDockerImage'), + RAW_SPAN_PIPELINE_YAML, + DATA_PREPPER_CONFIG_STATIC_YAML), + new TraceTestContainerConfiguration( + tasks.named('dataPrepperDockerImage'), + RAW_SPAN_PIPELINE_YAML, + DATA_PREPPER_CONFIG_STATIC_YAML)], + 'data-prepper' + ), + new TraceTestConfiguration( + 'rawSpanLatestReleaseCompatibilityEndToEndTest', + 'Runs the raw span integration tests with the latest released Data Prepper as a peer.', + RAW_SPAN_TESTS_PATTERN, + [new TraceTestContainerConfiguration( + tasks.named('dataPrepperDockerImage'), + 'raw-span-e2e-pipeline-from-build.yml', + DATA_PREPPER_CONFIG_STATIC_YAML), + new TraceTestContainerConfiguration( + tasks.named('latestDataPrepperDockerImage'), + 'raw-span-e2e-pipeline-latest-release.yml', + DATA_PREPPER_CONFIG_STATIC_YAML) + ], + 'data-prepper' + ), + new TraceTestConfiguration( + 'serviceMapPeerForwarderEndToEndTest', + 'Runs the service map with peer-forwarder integration tests.', + 'org.opensearch.dataprepper.integration.trace.EndToEndServiceMapTest.testPipelineEndToEnd*', + [new TraceTestContainerConfiguration( + tasks.named('dataPrepperDockerImage'), + 'service-map-e2e-pipeline.yml', + DATA_PREPPER_CONFIG_STATIC_YAML), + new TraceTestContainerConfiguration( + tasks.named('dataPrepperDockerImage'), + 'service-map-e2e-pipeline.yml', + DATA_PREPPER_CONFIG_STATIC_YAML)], + 'data-prepper' + ) +] + + +traceTestConfigurations.each { testConfiguration -> + tasks.register("start${testConfiguration.testName}All") + tasks.register("remove${testConfiguration.testName}All") + + (0.. + tasks.register("create${testConfiguration.testName}${containerIndex}", DockerCreateContainer) { + dependsOn testConfiguration.containerConfigurations.get(containerIndex).imageTask + dependsOn createDataPrepperNetwork + containerName = "${testConfiguration.containerNamePrefix}-${containerIndex}" + exposePorts('tcp', [21890]) + hostConfig.portBindings = ["${21890+containerIndex}:21890"] + hostConfig.binds = [ + (project.file("src/integrationTest/resources/${testConfiguration.containerConfigurations.get(containerIndex).pipelineConfiguration}").toString()) : '/usr/share/data-prepper/pipelines/trace-pipeline.yaml', + (project.file("src/integrationTest/resources/${testConfiguration.containerConfigurations.get(containerIndex).dataPrepperConfiguration}").toString()): '/usr/share/data-prepper/config/data-prepper-config.yaml' + ] + hostConfig.network = createDataPrepperNetwork.getNetworkName() + targetImageId testConfiguration.containerConfigurations.get(containerIndex).imageTask.get().imageId + } + + tasks.register("start${testConfiguration.testName}${containerIndex}", DockerStartContainer) { + dependsOn "create${testConfiguration.testName}${containerIndex}" + dependsOn 'startOpenSearchDockerContainer' + mustRunAfter 'startOpenSearchDockerContainer' + targetContainerId tasks.getByName("create${testConfiguration.testName}${containerIndex}").getContainerId() + } + + tasks.named("start${testConfiguration.testName}All").configure { + dependsOn "start${testConfiguration.testName}${containerIndex}" + } + + tasks.register("stop${testConfiguration.testName}${containerIndex}", DockerStopContainer) { + dependsOn "${testConfiguration.testName}" + targetContainerId tasks.getByName("create${testConfiguration.testName}${containerIndex}").getContainerId() + } + + tasks.register("remove${testConfiguration.testName}${containerIndex}", DockerRemoveContainer) { + dependsOn "stop${testConfiguration.testName}${containerIndex}" + targetContainerId tasks.getByName("stop${testConfiguration.testName}${containerIndex}").getContainerId() + } + + tasks.named("remove${testConfiguration.testName}All").configure { + dependsOn "remove${testConfiguration.testName}${containerIndex}" + } } -} -/** - * End to end test. Spins up OpenSearch and DataPrepper docker containers, then runs the integ test - * Stops the docker containers when finished - */ -def createEndToEndTest(final String testName, final String includeTestsMatchPattern, - final DockerCreateContainer createDataPrepper1Task, final DockerCreateContainer createDataPrepper2Task) { - return tasks.create(testName, Test) { + + tasks.register(testConfiguration.testName, Test) { dependsOn build dependsOn startOpenSearchDockerContainer - def startDataPrepper1Task = startDataPrepperDockerContainer(createDataPrepper1Task as DockerCreateContainer) - def startDataPrepper2Task = startDataPrepperDockerContainer(createDataPrepper2Task as DockerCreateContainer) - dependsOn startDataPrepper1Task - dependsOn startDataPrepper2Task - startDataPrepper1Task.mustRunAfter 'startOpenSearchDockerContainer' - startDataPrepper2Task.mustRunAfter 'startOpenSearchDockerContainer' - // wait for data-preppers to be ready + dependsOn "start${testConfiguration.testName}All" + + // Wait for Data Prepper image to be ready doFirst { - sleep(15*1000) + sleep(15 * 1000) } - description = 'Runs the raw span integration tests.' + description = testConfiguration.description group = 'verification' testClassesDirs = sourceSets.integrationTest.output.classesDirs classpath = sourceSets.integrationTest.runtimeClasspath filter { - includeTestsMatching includeTestsMatchPattern + includeTestsMatching testConfiguration.testFilters } - def stopDataPrepper1Task = stopDataPrepperDockerContainer(startDataPrepper1Task as DockerStartContainer) - def stopDataPrepper2Task = stopDataPrepperDockerContainer(startDataPrepper2Task as DockerStartContainer) - finalizedBy stopDataPrepper1Task - finalizedBy stopDataPrepper2Task - finalizedBy removeDataPrepperDockerContainer(stopDataPrepper1Task as DockerStopContainer) - finalizedBy removeDataPrepperDockerContainer(stopDataPrepper2Task as DockerStopContainer) finalizedBy stopOpenSearchDockerContainer + finalizedBy "remove${testConfiguration.testName}All" finalizedBy removeDataPrepperNetwork } } -// raw span e2e test -def includeRawSpanTestsMatchPattern = "org.opensearch.dataprepper.integration.trace.EndToEndRawSpanTest.testPipelineEndToEnd*" - -def createRawSpanDataPrepper1Task = createDataPrepperDockerContainer( - "rawSpanDataPrepper1", "dataprepper1", 21890, "${RAW_SPAN_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_YAML}") -def createRawSpanDataPrepper2Task = createDataPrepperDockerContainer( - "rawSpanDataPrepper2", "dataprepper2", 21891, "${RAW_SPAN_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_YAML}") - -def rawSpanEndToEndTest = createEndToEndTest("rawSpanEndToEndTest", includeRawSpanTestsMatchPattern, - createRawSpanDataPrepper1Task, createRawSpanDataPrepper2Task) - -// raw span with peer forwarding e2e test -def createRawSpanPeerForwarderDataPrepper1Task = createDataPrepperDockerContainer( - "rawSpanPeerForwarderDataPrepper1", "dataprepper1", 21890, "${RAW_SPAN_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_STATIC_YAML}") -def createRawSpanPeerForwarderDataPrepper2Task = createDataPrepperDockerContainer( - "rawSpanPeerForwarderDataPrepper2", "dataprepper2", 21891, "${RAW_SPAN_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_STATIC_YAML}") - -def rawSpanPeerForwarderEndToEndTest = createEndToEndTest("rawSpanPeerForwarderEndToEndTest", includeRawSpanTestsMatchPattern, - createRawSpanPeerForwarderDataPrepper1Task, createRawSpanPeerForwarderDataPrepper2Task) - -// raw span compatibility e2e test -def rawSpanDataPrepperEventFromBuild = createDataPrepperDockerContainer( - "rawSpanDataPrepperEventFromBuild", "dataprepper1", 21890, "${RAW_SPAN_PIPELINE_FROM_BUILD_YAML}", "${DATA_PREPPER_CONFIG_STATIC_YAML}") -def rawSpanDataPrepperLatestFromPull = createDataPrepperDockerContainerFromPullImage( - "rawSpanDataPrepperLatestFromPull", "dataprepper2", 21891, "${RAW_SPAN_PIPELINE_LATEST_RELEASE_YAML}", "${DATA_PREPPER_CONFIG_STATIC_YAML}") - -def rawSpanLatestReleaseCompatibilityEndToEndTest = createEndToEndTest("rawSpanLatestReleaseCompatibilityEndToEndTest", - includeRawSpanTestsMatchPattern, - rawSpanDataPrepperEventFromBuild, rawSpanDataPrepperLatestFromPull) - -// service map e2e -def includeServiceMapTestsMatchPattern = "org.opensearch.dataprepper.integration.trace.EndToEndServiceMapTest.testPipelineEndToEnd*" - -// service map with peer forwarding e2e test -def createServiceMapPeerForwarderDataPrepper1Task = createDataPrepperDockerContainer( - "serviceMapPeerForwarderDataPrepper1", "dataprepper1", 21890, "${SERVICE_MAP_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_STATIC_YAML}") -def createServiceMapPeerForwarderDataPrepper2Task = createDataPrepperDockerContainer( - "serviceMapPeerForwarderDataPrepper2", "dataprepper2", 21891, "${SERVICE_MAP_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_STATIC_YAML}") - -def serviceMapPeerForwarderEndToEndTest = createEndToEndTest("serviceMapPeerForwarderEndToEndTest", includeServiceMapTestsMatchPattern, - createServiceMapPeerForwarderDataPrepper1Task, createServiceMapPeerForwarderDataPrepper2Task) dependencies { integrationTestImplementation project(':data-prepper-api') diff --git a/e2e-test/trace/src/integrationTest/resources/data_prepper_static.yml b/e2e-test/trace/src/integrationTest/resources/data_prepper_static.yml index 9add9cf0b2..b40680acc1 100644 --- a/e2e-test/trace/src/integrationTest/resources/data_prepper_static.yml +++ b/e2e-test/trace/src/integrationTest/resources/data_prepper_static.yml @@ -3,4 +3,4 @@ peer_forwarder: port: 4994 ssl: false discovery_mode: static - static_endpoints: ["dataprepper1", "dataprepper2"] + static_endpoints: ["data-prepper-0", "data-prepper-1"] From 6681e75d8b8cfa3985e4a11f9fa9d6238562e462 Mon Sep 17 00:00:00 2001 From: kkondaka <41027584+kkondaka@users.noreply.github.com> Date: Sat, 4 Nov 2023 08:04:22 -0700 Subject: [PATCH 24/53] Add ProgressCheck callbacks to end-to-end acknowledgements (#3565) Add ProgressCheck callbacks to end-to-end acknowledgements Signed-off-by: Krishna Kondaka Co-authored-by: Krishna Kondaka --- .../acknowledgements/AcknowledgementSet.java | 15 ++ .../model/acknowledgements/ProgressCheck.java | 17 ++ .../AcknowledgementAppConfig.java | 6 +- .../DefaultAcknowledgementSet.java | 47 ++++- .../DefaultAcknowledgementSetManager.java | 12 +- .../DefaultProgressCheck.java | 21 +++ ...DefaultAcknowledgementSetManagerTests.java | 135 ++++++++++++-- .../DefaultAcknowledgementSetTests.java | 51 +++++- .../kafka/source/KafkaSourceJsonTypeIT.java | 4 +- .../consumer/KafkaCustomConsumerTest.java | 6 +- .../sink/opensearch/OpenSearchSink.java | 3 +- .../source/s3/S3ScanObjectWorkerIT.java | 10 +- .../plugins/source/s3/SqsWorkerIT.java | 164 +++++++++++++++++- .../plugins/source/s3/SqsWorker.java | 58 ++++++- .../source/s3/configuration/SqsOptions.java | 18 ++ .../plugins/source/s3/SqsWorkerTest.java | 39 ++++- 16 files changed, 547 insertions(+), 59 deletions(-) create mode 100644 data-prepper-api/src/main/java/org/opensearch/dataprepper/model/acknowledgements/ProgressCheck.java create mode 100644 data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultProgressCheck.java diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/acknowledgements/AcknowledgementSet.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/acknowledgements/AcknowledgementSet.java index c95c2e5f88..efd36e123d 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/acknowledgements/AcknowledgementSet.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/acknowledgements/AcknowledgementSet.java @@ -8,6 +8,9 @@ import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.event.EventHandle; +import java.time.Duration; +import java.util.function.Consumer; + /** * AcknowledgmentSet keeps track of set of events that * belong to the batch of events that a source creates. @@ -58,4 +61,16 @@ public interface AcknowledgementSet { * initial events are going through the pipeline line. */ public void complete(); + + /** + * adds progress check callback to the acknowledgement set. When added + * the callback is called every progressCheckInterval time with the + * indication of current progress as a ratio of pending number of + * acknowledgements over total acknowledgements + * + * @param progressCheckCallback progress check callback to be called + * @param progressCheckInterval frequency of invocation of progress check callback + * @since 2.6 + */ + public void addProgressCheck(final Consumer progressCheckCallback, final Duration progressCheckInterval); } diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/acknowledgements/ProgressCheck.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/acknowledgements/ProgressCheck.java new file mode 100644 index 0000000000..07a2f18c03 --- /dev/null +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/acknowledgements/ProgressCheck.java @@ -0,0 +1,17 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.model.acknowledgements; + +public interface ProgressCheck { + /** + * Returns the pending ratio + * + * @return returns the ratio of pending to the total acknowledgements + * @since 2.6 + */ + Double getRatio(); +} + diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementAppConfig.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementAppConfig.java index 21032873a4..2d32cb116c 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementAppConfig.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/AcknowledgementAppConfig.java @@ -8,7 +8,7 @@ import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; -import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.ThreadFactory; @@ -23,7 +23,7 @@ CallbackTheadFactory callbackTheadFactory() { } @Bean(name = "acknowledgementCallbackExecutor") - ExecutorService acknowledgementCallbackExecutor(final CallbackTheadFactory callbackTheadFactory) { - return Executors.newFixedThreadPool(MAX_THREADS, callbackTheadFactory); + ScheduledExecutorService acknowledgementCallbackExecutor(final CallbackTheadFactory callbackTheadFactory) { + return Executors.newScheduledThreadPool(MAX_THREADS, callbackTheadFactory); } } diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java index 96e25fd9a8..1c3ef6032d 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSet.java @@ -6,6 +6,7 @@ package org.opensearch.dataprepper.acknowledgements; import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; +import org.opensearch.dataprepper.model.acknowledgements.ProgressCheck; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.event.EventHandle; import org.opensearch.dataprepper.model.event.InternalEventHandle; @@ -18,37 +19,61 @@ import java.time.Instant; import java.util.HashMap; import java.util.Map; -import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.Future; +import java.util.concurrent.ScheduledFuture; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.locks.ReentrantLock; +import java.util.concurrent.TimeUnit; import java.util.function.Consumer; public class DefaultAcknowledgementSet implements AcknowledgementSet { private static final Logger LOG = LoggerFactory.getLogger(DefaultAcknowledgementSet.class); private final Consumer callback; + private Consumer progressCheckCallback; private final Instant expiryTime; - private final ExecutorService executor; + private final ScheduledExecutorService scheduledExecutor; // This lock protects all the non-final members private final ReentrantLock lock; private boolean result; private final Map pendingAcknowledgments; private Future callbackFuture; private final DefaultAcknowledgementSetMetrics metrics; + private ScheduledFuture progressCheckFuture; private boolean completed; + private AtomicInteger totalEventsAdded; - public DefaultAcknowledgementSet(final ExecutorService executor, final Consumer callback, final Duration expiryTime, final DefaultAcknowledgementSetMetrics metrics) { + public DefaultAcknowledgementSet(final ScheduledExecutorService scheduledExecutor, + final Consumer callback, + final Duration expiryTime, + final DefaultAcknowledgementSetMetrics metrics) { this.callback = callback; this.result = true; - this.executor = executor; + this.totalEventsAdded = new AtomicInteger(0); + this.scheduledExecutor = scheduledExecutor; this.expiryTime = Instant.now().plusMillis(expiryTime.toMillis()); this.callbackFuture = null; this.metrics = metrics; this.completed = false; + this.progressCheckCallback = null; pendingAcknowledgments = new HashMap<>(); lock = new ReentrantLock(true); } + public void addProgressCheck(final Consumer progressCheckCallback, final Duration progressCheckInterval) { + this.progressCheckCallback = progressCheckCallback; + this.progressCheckFuture = scheduledExecutor.scheduleAtFixedRate(this::checkProgress, 0L, progressCheckInterval.toMillis(), TimeUnit.MILLISECONDS); + } + + public void checkProgress() { + lock.lock(); + int numberOfEventsPending = pendingAcknowledgments.size(); + lock.unlock(); + if (progressCheckCallback != null) { + progressCheckCallback.accept(new DefaultProgressCheck((double)numberOfEventsPending/totalEventsAdded.get())); + } + } + @Override public void add(Event event) { lock.lock(); @@ -59,6 +84,7 @@ public void add(Event event) { InternalEventHandle internalEventHandle = (InternalEventHandle)(DefaultEventHandle)eventHandle; internalEventHandle.setAcknowledgementSet(this); pendingAcknowledgments.put(eventHandle, new AtomicInteger(1)); + totalEventsAdded.incrementAndGet(); } } } finally { @@ -88,6 +114,9 @@ public boolean isDone() { return true; } if (Instant.now().isAfter(expiryTime)) { + if (progressCheckFuture != null) { + progressCheckFuture.cancel(false); + } if (callbackFuture != null) { callbackFuture.cancel(true); callbackFuture = null; @@ -112,7 +141,10 @@ public void complete() { try { completed = true; if (pendingAcknowledgments.size() == 0) { - callbackFuture = executor.submit(() -> callback.accept(this.result)); + if (progressCheckFuture != null) { + progressCheckFuture.cancel(false); + } + callbackFuture = scheduledExecutor.submit(() -> callback.accept(this.result)); } } finally { lock.unlock(); @@ -136,7 +168,10 @@ public boolean release(final EventHandle eventHandle, final boolean result) { if (pendingAcknowledgments.get(eventHandle).decrementAndGet() == 0) { pendingAcknowledgments.remove(eventHandle); if (completed && pendingAcknowledgments.size() == 0) { - callbackFuture = executor.submit(() -> callback.accept(this.result)); + if (progressCheckFuture != null) { + progressCheckFuture.cancel(false); + } + callbackFuture = scheduledExecutor.submit(() -> callback.accept(this.result)); return true; } else if (pendingAcknowledgments.size() == 0) { LOG.debug("Acknowledgement set is not completed. Delaying callback until it is completed"); diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManager.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManager.java index 104945960e..3f2e3761bd 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManager.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManager.java @@ -15,27 +15,27 @@ import javax.inject.Named; import java.time.Duration; import java.util.Objects; -import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; import java.util.function.Consumer; @Named public class DefaultAcknowledgementSetManager implements AcknowledgementSetManager { private static final int DEFAULT_WAIT_TIME_MS = 15 * 1000; private final AcknowledgementSetMonitor acknowledgementSetMonitor; - private final ExecutorService executor; + private final ScheduledExecutorService scheduledExecutor; private final AcknowledgementSetMonitorThread acknowledgementSetMonitorThread; private PluginMetrics pluginMetrics; private DefaultAcknowledgementSetMetrics metrics; @Inject public DefaultAcknowledgementSetManager( - @Named("acknowledgementCallbackExecutor") final ExecutorService callbackExecutor) { + @Named("acknowledgementCallbackExecutor") final ScheduledExecutorService callbackExecutor) { this(callbackExecutor, Duration.ofMillis(DEFAULT_WAIT_TIME_MS)); } - public DefaultAcknowledgementSetManager(final ExecutorService callbackExecutor, final Duration waitTime) { + public DefaultAcknowledgementSetManager(final ScheduledExecutorService callbackExecutor, final Duration waitTime) { this.acknowledgementSetMonitor = new AcknowledgementSetMonitor(); - this.executor = Objects.requireNonNull(callbackExecutor); + this.scheduledExecutor = Objects.requireNonNull(callbackExecutor); acknowledgementSetMonitorThread = new AcknowledgementSetMonitorThread(acknowledgementSetMonitor, waitTime); acknowledgementSetMonitorThread.start(); pluginMetrics = PluginMetrics.fromNames("acknowledgementSetManager", "acknowledgements"); @@ -43,7 +43,7 @@ public DefaultAcknowledgementSetManager(final ExecutorService callbackExecutor, } public AcknowledgementSet create(final Consumer callback, final Duration timeout) { - AcknowledgementSet acknowledgementSet = new DefaultAcknowledgementSet(executor, callback, timeout, metrics); + AcknowledgementSet acknowledgementSet = new DefaultAcknowledgementSet(scheduledExecutor, callback, timeout, metrics); acknowledgementSetMonitor.add(acknowledgementSet); metrics.increment(DefaultAcknowledgementSetMetrics.CREATED_METRIC_NAME); return acknowledgementSet; diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultProgressCheck.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultProgressCheck.java new file mode 100644 index 0000000000..87b7a8226d --- /dev/null +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/acknowledgements/DefaultProgressCheck.java @@ -0,0 +1,21 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.acknowledgements; + +import org.opensearch.dataprepper.model.acknowledgements.ProgressCheck; + +public class DefaultProgressCheck implements ProgressCheck { + double ratio; + + public DefaultProgressCheck(double ratio) { + this.ratio = ratio; + } + + @Override + public Double getRatio() { + return ratio; + } +} diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManagerTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManagerTests.java index c9fd556214..1b87d6c849 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManagerTests.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetManagerTests.java @@ -7,6 +7,7 @@ import org.opensearch.dataprepper.model.event.JacksonEvent; import org.opensearch.dataprepper.model.event.DefaultEventHandle; +import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; import org.junit.jupiter.api.BeforeEach; @@ -22,14 +23,14 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; import java.time.Duration; -import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.Executors; @ExtendWith(MockitoExtension.class) class DefaultAcknowledgementSetManagerTests { private static final Duration TEST_TIMEOUT = Duration.ofMillis(400); - DefaultAcknowledgementSetManager acknowledgementSetManager; - private ExecutorService callbackExecutor; + private DefaultAcknowledgementSetManager acknowledgementSetManager; + private ScheduledExecutorService callbackExecutor; @Mock JacksonEvent event1; @@ -38,17 +39,24 @@ class DefaultAcknowledgementSetManagerTests { @Mock JacksonEvent event3; - DefaultEventHandle eventHandle1; - DefaultEventHandle eventHandle2; - DefaultEventHandle eventHandle3; - Boolean result; + private PluginMetrics pluginMetrics; + private DefaultEventHandle eventHandle1; + private DefaultEventHandle eventHandle2; + private DefaultEventHandle eventHandle3; + private DefaultEventHandle eventHandle4; + private DefaultEventHandle eventHandle5; + private DefaultEventHandle eventHandle6; + private Boolean result; + private double currentRatio; @BeforeEach void setup() { - callbackExecutor = Executors.newFixedThreadPool(2); + currentRatio = 0; + callbackExecutor = Executors.newScheduledThreadPool(2); event1 = mock(JacksonEvent.class); eventHandle1 = mock(DefaultEventHandle.class); lenient().when(event1.getEventHandle()).thenReturn(eventHandle1); + pluginMetrics = mock(PluginMetrics.class); event2 = mock(JacksonEvent.class); eventHandle2 = mock(DefaultEventHandle.class); @@ -76,8 +84,6 @@ void testBasic() { assertThat(acknowledgementSetManager.getAcknowledgementSetMonitor().getSize(), equalTo(0)); assertThat(result, equalTo(true)); }); - assertThat(acknowledgementSetManager.getAcknowledgementSetMonitor().getSize(), equalTo(0)); - assertThat(result, equalTo(true)); } @Test @@ -85,7 +91,10 @@ void testExpirations() throws InterruptedException { acknowledgementSetManager.releaseEventReference(eventHandle2, true); Thread.sleep(TEST_TIMEOUT.multipliedBy(5).toMillis()); assertThat(acknowledgementSetManager.getAcknowledgementSetMonitor().getSize(), equalTo(0)); - assertThat(result, equalTo(null)); + await().atMost(TEST_TIMEOUT.multipliedBy(5)) + .untilAsserted(() -> { + assertThat(result, equalTo(null)); + }); } @Test @@ -106,7 +115,107 @@ void testMultipleAcknowledgementSets() { assertThat(acknowledgementSetManager.getAcknowledgementSetMonitor().getSize(), equalTo(0)); assertThat(result, equalTo(true)); }); - assertThat(acknowledgementSetManager.getAcknowledgementSetMonitor().getSize(), equalTo(0)); - assertThat(result, equalTo(true)); } + + @Test + void testWithProgressCheckCallbacks() { + eventHandle3 = mock(DefaultEventHandle.class); + lenient().when(event3.getEventHandle()).thenReturn(eventHandle3); + + eventHandle4 = mock(DefaultEventHandle.class); + JacksonEvent event4 = mock(JacksonEvent.class); + lenient().when(event4.getEventHandle()).thenReturn(eventHandle4); + + eventHandle5 = mock(DefaultEventHandle.class); + JacksonEvent event5 = mock(JacksonEvent.class); + lenient().when(event5.getEventHandle()).thenReturn(eventHandle5); + + eventHandle6 = mock(DefaultEventHandle.class); + JacksonEvent event6 = mock(JacksonEvent.class); + lenient().when(event6.getEventHandle()).thenReturn(eventHandle6); + + AcknowledgementSet acknowledgementSet2 = acknowledgementSetManager.create((flag) -> { result = flag; }, Duration.ofMillis(10000)); + acknowledgementSet2.addProgressCheck((progressCheck) -> {currentRatio = progressCheck.getRatio();}, Duration.ofSeconds(1)); + acknowledgementSet2.add(event3); + acknowledgementSet2.add(event4); + acknowledgementSet2.add(event5); + acknowledgementSet2.add(event6); + lenient().when(eventHandle3.getAcknowledgementSet()).thenReturn(acknowledgementSet2); + lenient().when(eventHandle4.getAcknowledgementSet()).thenReturn(acknowledgementSet2); + lenient().when(eventHandle5.getAcknowledgementSet()).thenReturn(acknowledgementSet2); + lenient().when(eventHandle6.getAcknowledgementSet()).thenReturn(acknowledgementSet2); + acknowledgementSet2.complete(); + acknowledgementSetManager.releaseEventReference(eventHandle3, true); + await().atMost(TEST_TIMEOUT.multipliedBy(5)) + .untilAsserted(() -> { + assertThat(currentRatio, equalTo(0.75)); + }); + acknowledgementSetManager.releaseEventReference(eventHandle4, true); + await().atMost(TEST_TIMEOUT.multipliedBy(5)) + .untilAsserted(() -> { + assertThat(currentRatio, equalTo(0.5)); + }); + acknowledgementSetManager.releaseEventReference(eventHandle5, true); + await().atMost(TEST_TIMEOUT.multipliedBy(5)) + .untilAsserted(() -> { + assertThat(currentRatio, equalTo(0.25)); + }); + acknowledgementSetManager.releaseEventReference(eventHandle6, true); + await().atMost(TEST_TIMEOUT.multipliedBy(5)) + .untilAsserted(() -> { + assertThat(result, equalTo(true)); + }); + + } + + @Test + void testWithProgressCheckCallbacks_AcksExpire() { + eventHandle3 = mock(DefaultEventHandle.class); + lenient().when(event3.getEventHandle()).thenReturn(eventHandle3); + + eventHandle4 = mock(DefaultEventHandle.class); + JacksonEvent event4 = mock(JacksonEvent.class); + lenient().when(event4.getEventHandle()).thenReturn(eventHandle4); + + eventHandle5 = mock(DefaultEventHandle.class); + JacksonEvent event5 = mock(JacksonEvent.class); + lenient().when(event5.getEventHandle()).thenReturn(eventHandle5); + + eventHandle6 = mock(DefaultEventHandle.class); + JacksonEvent event6 = mock(JacksonEvent.class); + lenient().when(event6.getEventHandle()).thenReturn(eventHandle6); + + AcknowledgementSet acknowledgementSet2 = acknowledgementSetManager.create((flag) -> { result = flag; }, Duration.ofSeconds(10)); + acknowledgementSet2.addProgressCheck((progressCheck) -> {currentRatio = progressCheck.getRatio();}, Duration.ofSeconds(1)); + acknowledgementSet2.add(event3); + acknowledgementSet2.add(event4); + acknowledgementSet2.add(event5); + acknowledgementSet2.add(event6); + lenient().when(eventHandle3.getAcknowledgementSet()).thenReturn(acknowledgementSet2); + lenient().when(eventHandle4.getAcknowledgementSet()).thenReturn(acknowledgementSet2); + lenient().when(eventHandle5.getAcknowledgementSet()).thenReturn(acknowledgementSet2); + lenient().when(eventHandle6.getAcknowledgementSet()).thenReturn(acknowledgementSet2); + acknowledgementSet2.complete(); + acknowledgementSetManager.releaseEventReference(eventHandle3, true); + await().atMost(TEST_TIMEOUT.multipliedBy(5)) + .untilAsserted(() -> { + assertThat(currentRatio, equalTo(0.75)); + }); + acknowledgementSetManager.releaseEventReference(eventHandle4, true); + await().atMost(TEST_TIMEOUT.multipliedBy(5)) + .untilAsserted(() -> { + assertThat(currentRatio, equalTo(0.5)); + }); + acknowledgementSetManager.releaseEventReference(eventHandle5, true); + await().atMost(TEST_TIMEOUT.multipliedBy(5)) + .untilAsserted(() -> { + assertThat(currentRatio, equalTo(0.25)); + }); + await().atMost(TEST_TIMEOUT.multipliedBy(5)) + .untilAsserted(() -> { + assertThat(result, equalTo(null)); + }); + + } + } diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetTests.java index 8a4aa1485a..5deba46be8 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetTests.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/acknowledgements/DefaultAcknowledgementSetTests.java @@ -6,6 +6,7 @@ package org.opensearch.dataprepper.acknowledgements; import org.awaitility.Awaitility; +import static org.awaitility.Awaitility.await; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -16,7 +17,7 @@ import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; import java.time.Duration; -import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Consumer; @@ -41,7 +42,9 @@ class DefaultAcknowledgementSetTests { @Mock private DefaultEventHandle handle2; - private ExecutorService executor; + private double currentRatio; + + private ScheduledExecutorService executor; private Boolean acknowledgementSetResult; private final Duration TEST_TIMEOUT = Duration.ofMillis(5000); private AtomicBoolean callbackInterrupted; @@ -76,7 +79,7 @@ private DefaultAcknowledgementSet createObjectUnderTestWithCallback(Consumer callbackInterrupted.get()); assertThat(callbackInterrupted.get(), equalTo(true)); } + + @Test + void testDefaultAcknowledgementSetWithProgressCheck() throws Exception { + defaultAcknowledgementSet = createObjectUnderTestWithCallback( + (flag) -> { + acknowledgementSetResult = flag; + } + ); + defaultAcknowledgementSet.addProgressCheck( + (progressCheck) -> { + currentRatio = progressCheck.getRatio(); + }, + Duration.ofSeconds(1) + ); + defaultAcknowledgementSet.add(event); + defaultAcknowledgementSet.add(event2); + defaultAcknowledgementSet.complete(); + lenient().doAnswer(a -> { + AcknowledgementSet acknowledgementSet = a.getArgument(0); + lenient().when(handle.getAcknowledgementSet()).thenReturn(acknowledgementSet); + return null; + }).when(handle).setAcknowledgementSet(any(AcknowledgementSet.class)); + assertThat(handle, not(equalTo(null))); + assertThat(handle.getAcknowledgementSet(), equalTo(defaultAcknowledgementSet)); + await().atMost(Duration.ofSeconds(5)) + .untilAsserted(() -> { + assertThat(currentRatio, equalTo(1.0)); + }); + assertThat(defaultAcknowledgementSet.release(handle, true), equalTo(false)); + await().atMost(Duration.ofSeconds(5)) + .untilAsserted(() -> { + assertThat(currentRatio, equalTo(0.5)); + }); + assertThat(defaultAcknowledgementSet.release(handle2, true), equalTo(true)); + Awaitility.waitAtMost(Duration.ofSeconds(10)) + .pollDelay(Duration.ofMillis(500)) + .until(() -> defaultAcknowledgementSet.isDone()); + await().atMost(Duration.ofSeconds(5)) + .untilAsserted(() -> { + assertThat(acknowledgementSetResult, equalTo(true)); + }); + } } diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceJsonTypeIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceJsonTypeIT.java index 17ab0d85bb..6cd7a5215f 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceJsonTypeIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSourceJsonTypeIT.java @@ -42,7 +42,7 @@ import java.util.List; import java.util.Map; import java.util.Properties; -import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicBoolean; @@ -107,7 +107,7 @@ public void setup() throws Throwable { buffer = mock(Buffer.class); encryptionConfig = mock(EncryptionConfig.class); receivedRecords = new ArrayList<>(); - ExecutorService executor = Executors.newFixedThreadPool(2); + ScheduledExecutorService executor = Executors.newScheduledThreadPool(2); acknowledgementSetManager = new DefaultAcknowledgementSetManager(executor); pipelineDescription = mock(PipelineDescription.class); when(sourceConfig.getAcknowledgementsEnabled()).thenReturn(false); 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 e32ce36836..a5913042de 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 @@ -45,7 +45,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicBoolean; @@ -72,7 +72,7 @@ public class KafkaCustomConsumerTest { @Mock private KafkaConsumerConfig sourceConfig; - private ExecutorService callbackExecutor; + private ScheduledExecutorService callbackExecutor; private AcknowledgementSetManager acknowledgementSetManager; @Mock @@ -134,7 +134,7 @@ public void setUp() { }).when(negCounter).increment(); doAnswer((i)-> {return posCount;}).when(posCounter).count(); doAnswer((i)-> {return negCount;}).when(negCounter).count(); - callbackExecutor = Executors.newFixedThreadPool(2); + callbackExecutor = Executors.newScheduledThreadPool(2); acknowledgementSetManager = new DefaultAcknowledgementSetManager(callbackExecutor, Duration.ofMillis(2000)); sourceConfig = mock(KafkaConsumerConfig.class); 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 8820de4afe..8853674d90 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 @@ -236,8 +236,7 @@ private void doInitializeInternal() throws IOException { .setParameter("filter_path", "errors,took,items.*.error,items.*.status,items.*._index,items.*._id") .build()); bulkApiWrapper = BulkApiWrapperFactory.getWrapper(openSearchSinkConfig.getIndexConfiguration(), filteringOpenSearchClient); - bulkRetryStrategy = new BulkRetryStrategy( - bulkRequest -> bulkApiWrapper.bulk(bulkRequest.getRequest()), + bulkRetryStrategy = new BulkRetryStrategy(bulkRequest -> bulkApiWrapper.bulk(bulkRequest.getRequest()), this::logFailureForBulkRequests, pluginMetrics, maxRetries, diff --git a/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanObjectWorkerIT.java b/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanObjectWorkerIT.java index 49f5c687b6..48c5862155 100644 --- a/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanObjectWorkerIT.java +++ b/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanObjectWorkerIT.java @@ -56,7 +56,7 @@ import java.util.List; import java.util.Objects; import java.util.Optional; -import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.ThreadLocalRandom; import java.util.function.Consumer; @@ -209,7 +209,7 @@ private ScanObjectWorker createObjectUnderTest(final RecordsGenerator recordsGen lenient().when(s3ScanSchedulingOptions.getInterval()).thenReturn(Duration.ofHours(1)); lenient().when(s3ScanSchedulingOptions.getCount()).thenReturn(1); - ExecutorService executor = Executors.newFixedThreadPool(2); + ScheduledExecutorService executor = Executors.newScheduledThreadPool(2); acknowledgementSetManager = new DefaultAcknowledgementSetManager(executor); return new ScanObjectWorker(s3Client,List.of(scanOptions),createObjectUnderTest(s3ObjectRequest) @@ -257,7 +257,7 @@ void parseS3Object_parquet_correctly_with_bucket_scan_and_loads_data_into_Buffer startTimeAndRangeScanOptions, Boolean.TRUE); - final ExecutorService executorService = Executors.newSingleThreadExecutor(); + final ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor(); executorService.submit(objectUnderTest::run); await().atMost(Duration.ofSeconds(30)).until(() -> waitForAllRecordsToBeProcessed(numberOfRecords)); @@ -299,7 +299,7 @@ void parseS3Object_correctly_with_bucket_scan_and_loads_data_into_Buffer( final int expectedWrites = numberOfRecords / numberOfRecordsToAccumulate + (numberOfRecords % numberOfRecordsToAccumulate != 0 ? 1 : 0); - final ExecutorService executorService = Executors.newSingleThreadExecutor(); + final ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor(); executorService.submit(scanObjectWorker::run); await().atMost(Duration.ofSeconds(30)).until(() -> waitForAllRecordsToBeProcessed(numberOfRecords)); @@ -346,7 +346,7 @@ void parseS3Object_correctly_with_bucket_scan_and_loads_data_into_Buffer_and_del final int expectedWrites = numberOfRecords / numberOfRecordsToAccumulate; - final ExecutorService executorService = Executors.newSingleThreadExecutor(); + final ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor(); executorService.submit(scanObjectWorker::run); diff --git a/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerIT.java b/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerIT.java index 474929f71e..dc09cd989c 100644 --- a/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerIT.java +++ b/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerIT.java @@ -35,11 +35,14 @@ import java.io.IOException; import java.time.Duration; import java.time.Instant; +import java.util.List; +import java.util.ArrayList; import java.util.UUID; -import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicBoolean; +import static org.awaitility.Awaitility.await; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.greaterThanOrEqualTo; @@ -60,6 +63,8 @@ class SqsWorkerIT { private SqsClient sqsClient; @Mock private S3Service s3Service; + @Mock + private SqsOptions sqsOptions; private S3SourceConfig s3SourceConfig; private PluginMetrics pluginMetrics; private S3ObjectGenerator s3ObjectGenerator; @@ -69,8 +74,11 @@ class SqsWorkerIT { private Double receivedCount = 0.0; private Double deletedCount = 0.0; private Double ackCallbackCount = 0.0; + private Double visibilityTimeoutChangedCount = 0.0; private Event event; private AtomicBoolean ready = new AtomicBoolean(false); + private int numEventsAdded; + private List events; @BeforeEach void setUp() { @@ -80,6 +88,7 @@ void setUp() { .build(); bucket = System.getProperty("tests.s3source.bucket"); s3ObjectGenerator = new S3ObjectGenerator(s3Client, bucket); + events = new ArrayList<>(); sqsClient = SqsClient.builder() .region(Region.of(System.getProperty("tests.s3source.region"))) @@ -100,14 +109,14 @@ void setUp() { lenient().when(pluginMetrics.summary(anyString())).thenReturn(distributionSummary); when(pluginMetrics.timer(anyString())).thenReturn(sqsMessageDelayTimer); - final SqsOptions sqsOptions = mock(SqsOptions.class); + sqsOptions = mock(SqsOptions.class); when(sqsOptions.getSqsUrl()).thenReturn(System.getProperty("tests.s3source.queue.url")); when(sqsOptions.getVisibilityTimeout()).thenReturn(Duration.ofSeconds(60)); when(sqsOptions.getMaximumMessages()).thenReturn(10); when(sqsOptions.getWaitTime()).thenReturn(Duration.ofSeconds(10)); when(s3SourceConfig.getSqsOptions()).thenReturn(sqsOptions); lenient().when(s3SourceConfig.getOnErrorOption()).thenReturn(OnErrorOption.DELETE_MESSAGES); - when(s3SourceConfig.getNotificationSource()).thenReturn(NotificationSourceOption.S3); + lenient().when(s3SourceConfig.getNotificationSource()).thenReturn(NotificationSourceOption.S3); } private SqsWorker createObjectUnderTest() { @@ -183,7 +192,7 @@ void processSqsMessages_should_return_at_least_one_message_with_acks_with_callba ackSet.add(event); return null; }).when(s3Service).addS3Object(any(S3ObjectReference.class), any(AcknowledgementSet.class)); - ExecutorService executor = Executors.newFixedThreadPool(2); + ScheduledExecutorService executor = Executors.newScheduledThreadPool(2); acknowledgementSetManager = new DefaultAcknowledgementSetManager(executor); final SqsWorker objectUnderTest = createObjectUnderTest(); Thread sinkThread = new Thread(() -> { @@ -251,12 +260,12 @@ void processSqsMessages_should_return_at_least_one_message_with_acks_with_callba this.notify(); } try { - Thread.sleep(4000); + Thread.sleep(2000); } catch (Exception e){} return null; }).when(s3Service).addS3Object(any(S3ObjectReference.class), any(AcknowledgementSet.class)); - ExecutorService executor = Executors.newFixedThreadPool(2); + ScheduledExecutorService executor = Executors.newScheduledThreadPool(2); acknowledgementSetManager = new DefaultAcknowledgementSetManager(executor); final SqsWorker objectUnderTest = createObjectUnderTest(); Thread sinkThread = new Thread(() -> { @@ -281,6 +290,149 @@ void processSqsMessages_should_return_at_least_one_message_with_acks_with_callba assertThat(ackCallbackCount, equalTo((double)1.0)); } + @ParameterizedTest + @ValueSource(ints = {1}) + void processSqsMessages_with_acks_and_progress_check_callbacks(final int numberOfObjectsToWrite) throws IOException, InterruptedException { + writeToS3(numberOfObjectsToWrite); + + when(s3SourceConfig.getAcknowledgements()).thenReturn(true); + final Counter receivedCounter = mock(Counter.class); + final Counter deletedCounter = mock(Counter.class); + final Counter ackCallbackCounter = mock(Counter.class); + final Counter visibilityTimeoutChangedCounter = mock(Counter.class); + when(pluginMetrics.counter(SqsWorker.SQS_MESSAGES_RECEIVED_METRIC_NAME)).thenReturn(receivedCounter); + when(pluginMetrics.counter(SqsWorker.SQS_MESSAGES_DELETED_METRIC_NAME)).thenReturn(deletedCounter); + when(pluginMetrics.counter(SqsWorker.ACKNOWLEDGEMENT_SET_CALLACK_METRIC_NAME)).thenReturn(ackCallbackCounter); + when(pluginMetrics.counter(SqsWorker.SQS_VISIBILITY_TIMEOUT_CHANGED_COUNT_METRIC_NAME)).thenReturn(visibilityTimeoutChangedCounter); + lenient().doAnswer((val) -> { + receivedCount += (double)val.getArgument(0); + return null; + }).when(receivedCounter).increment(any(Double.class)); + + lenient().doAnswer((val) -> { + if (val.getArgument(0) != null) { + deletedCount += (double)val.getArgument(0); + } + return null; + }).when(deletedCounter).increment(any(Double.class)); + ackCallbackCount = 0.0; + lenient().doAnswer((val) -> { + ackCallbackCount += 1; + return null; + }).when(ackCallbackCounter).increment(); + lenient().doAnswer((val) -> { + visibilityTimeoutChangedCount += 1; + return null; + }).when(visibilityTimeoutChangedCounter).increment(); + numEventsAdded = 0; + + doAnswer((val) -> { + AcknowledgementSet ackSet = val.getArgument(1); + S3ObjectReference s3ObjectReference = val.getArgument(0); + assertThat(s3ObjectReference.getBucketName(), equalTo(bucket)); + assertThat(s3ObjectReference.getKey(), startsWith("s3 source/sqs/")); + event = (Event)JacksonEvent.fromMessage(val.getArgument(0).toString()); + + ackSet.add(event); + synchronized(events) { + events.add(event); + } + try { + Thread.sleep(2000); + } catch (Exception e) {} + return null; + }).when(s3Service).addS3Object(any(S3ObjectReference.class), any(AcknowledgementSet.class)); + ScheduledExecutorService executor = Executors.newScheduledThreadPool(2); + when(sqsOptions.getVisibilityTimeout()).thenReturn(Duration.ofSeconds(6)); + when(sqsOptions.getMaxVisibilityTimeoutExtension()).thenReturn(Duration.ofSeconds(60)); + when(sqsOptions.getVisibilityDuplicateProtection()).thenReturn(true); + acknowledgementSetManager = new DefaultAcknowledgementSetManager(executor); + final SqsWorker objectUnderTest = createObjectUnderTest(); + final int sqsMessagesProcessed = objectUnderTest.processSqsMessages(); + synchronized(events) { + for (Event e: events) { + if (e.getEventHandle() != null) { + e.getEventHandle().release(true); + } + } + } + await().atMost(Duration.ofSeconds(20)) + .untilAsserted(() -> { + assertThat(visibilityTimeoutChangedCount, greaterThanOrEqualTo((double)numberOfObjectsToWrite)); + assertThat(deletedCount, equalTo((double)numberOfObjectsToWrite)); + assertThat(ackCallbackCount, equalTo((double)numberOfObjectsToWrite)); + }); + } + + @ParameterizedTest + @ValueSource(ints = {1}) + void processSqsMessages_with_acks_and_progress_check_callbacks_expires(final int numberOfObjectsToWrite) throws IOException, InterruptedException { + writeToS3(numberOfObjectsToWrite); + + when(s3SourceConfig.getAcknowledgements()).thenReturn(true); + final Counter receivedCounter = mock(Counter.class); + final Counter deletedCounter = mock(Counter.class); + final Counter ackCallbackCounter = mock(Counter.class); + final Counter visibilityTimeoutChangedCounter = mock(Counter.class); + when(pluginMetrics.counter(SqsWorker.SQS_MESSAGES_RECEIVED_METRIC_NAME)).thenReturn(receivedCounter); + when(pluginMetrics.counter(SqsWorker.SQS_MESSAGES_DELETED_METRIC_NAME)).thenReturn(deletedCounter); + when(pluginMetrics.counter(SqsWorker.ACKNOWLEDGEMENT_SET_CALLACK_METRIC_NAME)).thenReturn(ackCallbackCounter); + when(pluginMetrics.counter(SqsWorker.SQS_VISIBILITY_TIMEOUT_CHANGED_COUNT_METRIC_NAME)).thenReturn(visibilityTimeoutChangedCounter); + lenient().doAnswer((val) -> { + receivedCount += (double)val.getArgument(0); + return null; + }).when(receivedCounter).increment(any(Double.class)); + + lenient().doAnswer((val) -> { + if (val.getArgument(0) != null) { + deletedCount += (double)val.getArgument(0); + } + return null; + }).when(deletedCounter).increment(any(Double.class)); + lenient().when(deletedCounter.count()).thenReturn(deletedCount); + ackCallbackCount = 0.0; + lenient().doAnswer((val) -> { + ackCallbackCount += 1; + return null; + }).when(ackCallbackCounter).increment(); + lenient().doAnswer((val) -> { + visibilityTimeoutChangedCount += 1; + return null; + }).when(visibilityTimeoutChangedCounter).increment(); + numEventsAdded = 0; + + doAnswer((val) -> { + AcknowledgementSet ackSet = val.getArgument(1); + S3ObjectReference s3ObjectReference = val.getArgument(0); + assertThat(s3ObjectReference.getBucketName(), equalTo(bucket)); + assertThat(s3ObjectReference.getKey(), startsWith("s3 source/sqs/")); + event = (Event)JacksonEvent.fromMessage(val.getArgument(0).toString()); + + ackSet.add(event); + synchronized(events) { + events.add(event); + } + try { + Thread.sleep(2000); + } catch (Exception e) {} + return null; + }).when(s3Service).addS3Object(any(S3ObjectReference.class), any(AcknowledgementSet.class)); + ScheduledExecutorService executor = Executors.newScheduledThreadPool(2); + when(sqsOptions.getVisibilityTimeout()).thenReturn(Duration.ofSeconds(6)); + when(sqsOptions.getMaxVisibilityTimeoutExtension()).thenReturn(Duration.ofSeconds(60)); + when(sqsOptions.getVisibilityDuplicateProtection()).thenReturn(true); + acknowledgementSetManager = new DefaultAcknowledgementSetManager(executor); + final SqsWorker objectUnderTest = createObjectUnderTest(); + final int sqsMessagesProcessed = objectUnderTest.processSqsMessages(); + await().atMost(Duration.ofSeconds(10)) + .untilAsserted(() -> { + assertThat(visibilityTimeoutChangedCount, greaterThanOrEqualTo((double)numberOfObjectsToWrite)); + assertThat(deletedCount, equalTo(0.0)); + assertThat(ackCallbackCount, equalTo(0.0)); + }); + + } + /** The EventBridge test is disabled by default * To run this test run only this one test with S3 bucket configured to use EventBridge to send notifications to SQS */ diff --git a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorker.java b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorker.java index 06a38d2393..a0f7a0bb16 100644 --- a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorker.java +++ b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorker.java @@ -28,6 +28,7 @@ import software.amazon.awssdk.services.sqs.SqsClient; import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequest; import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequestEntry; +import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityRequest; import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchResponse; import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchResultEntry; import software.amazon.awssdk.services.sqs.model.Message; @@ -40,7 +41,9 @@ 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.Optional; import java.util.stream.Collectors; @@ -51,6 +54,7 @@ public class SqsWorker implements Runnable { static final String SQS_MESSAGES_FAILED_METRIC_NAME = "sqsMessagesFailed"; static final String SQS_MESSAGES_DELETE_FAILED_METRIC_NAME = "sqsMessagesDeleteFailed"; static final String SQS_MESSAGE_DELAY_METRIC_NAME = "sqsMessageDelay"; + static final String SQS_VISIBILITY_TIMEOUT_CHANGED_COUNT_METRIC_NAME = "sqsVisibilityTimeoutChangedCount"; static final String ACKNOWLEDGEMENT_SET_CALLACK_METRIC_NAME = "acknowledgementSetCallbackCounter"; private final S3SourceConfig s3SourceConfig; @@ -64,6 +68,7 @@ public class SqsWorker implements Runnable { private final Counter sqsMessagesFailedCounter; private final Counter sqsMessagesDeleteFailedCounter; private final Counter acknowledgementSetCallbackCounter; + private final Counter sqsVisibilityTimeoutChangedCount; private final Timer sqsMessageDelayTimer; private final Backoff standardBackoff; private int failedAttemptCount; @@ -72,6 +77,7 @@ public class SqsWorker implements Runnable { private final ObjectMapper objectMapper = new ObjectMapper(); private volatile boolean isStopped = false; + private Map parsedMessageVisibilityTimesMap; public SqsWorker(final AcknowledgementSetManager acknowledgementSetManager, final SqsClient sqsClient, @@ -89,6 +95,7 @@ public SqsWorker(final AcknowledgementSetManager acknowledgementSetManager, objectCreatedFilter = new S3ObjectCreatedFilter(); evenBridgeObjectCreatedFilter = new EventBridgeObjectCreatedFilter(); failedAttemptCount = 0; + parsedMessageVisibilityTimesMap = new HashMap<>(); sqsMessagesReceivedCounter = pluginMetrics.counter(SQS_MESSAGES_RECEIVED_METRIC_NAME); sqsMessagesDeletedCounter = pluginMetrics.counter(SQS_MESSAGES_DELETED_METRIC_NAME); @@ -96,6 +103,7 @@ public SqsWorker(final AcknowledgementSetManager acknowledgementSetManager, sqsMessagesDeleteFailedCounter = pluginMetrics.counter(SQS_MESSAGES_DELETE_FAILED_METRIC_NAME); sqsMessageDelayTimer = pluginMetrics.timer(SQS_MESSAGE_DELAY_METRIC_NAME); acknowledgementSetCallbackCounter = pluginMetrics.counter(ACKNOWLEDGEMENT_SET_CALLACK_METRIC_NAME); + sqsVisibilityTimeoutChangedCount = pluginMetrics.counter(SQS_VISIBILITY_TIMEOUT_CHANGED_COUNT_METRIC_NAME); } @Override @@ -226,16 +234,48 @@ && isEventBridgeEventTypeCreated(parsedMessage)) { for (ParsedMessage parsedMessage : parsedMessagesToRead) { List waitingForAcknowledgements = new ArrayList<>(); AcknowledgementSet acknowledgementSet = null; + final int visibilityTimeout = (int)sqsOptions.getVisibilityTimeout().getSeconds(); + final int maxVisibilityTimeout = (int)sqsOptions.getMaxVisibilityTimeoutExtension().getSeconds(); + final int progressCheckInterval = visibilityTimeout/2 - 1; if (endToEndAcknowledgementsEnabled) { - // Acknowledgement Set timeout is slightly smaller than the visibility timeout; - int timeout = (int) sqsOptions.getVisibilityTimeout().getSeconds() - 2; - acknowledgementSet = acknowledgementSetManager.create((result) -> { - acknowledgementSetCallbackCounter.increment(); - // Delete only if this is positive acknowledgement - if (result == true) { - deleteSqsMessages(waitingForAcknowledgements); - } - }, Duration.ofSeconds(timeout)); + int expiryTimeout = visibilityTimeout - 2; + final boolean visibilityDuplicateProtectionEnabled = sqsOptions.getVisibilityDuplicateProtection(); + if (visibilityDuplicateProtectionEnabled) { + expiryTimeout = maxVisibilityTimeout; + } + acknowledgementSet = acknowledgementSetManager.create( + (result) -> { + acknowledgementSetCallbackCounter.increment(); + // Delete only if this is positive acknowledgement + if (visibilityDuplicateProtectionEnabled) { + parsedMessageVisibilityTimesMap.remove(parsedMessage); + } + if (result == true) { + deleteSqsMessages(waitingForAcknowledgements); + } + }, + Duration.ofSeconds(expiryTimeout)); + if (visibilityDuplicateProtectionEnabled) { + acknowledgementSet.addProgressCheck( + (ratio) -> { + final int newVisibilityTimeoutSeconds = visibilityTimeout; + int newValue = parsedMessageVisibilityTimesMap.getOrDefault(parsedMessage, visibilityTimeout) + progressCheckInterval; + if (newValue >= maxVisibilityTimeout) { + return; + } + parsedMessageVisibilityTimesMap.put(parsedMessage, newValue); + final ChangeMessageVisibilityRequest changeMessageVisibilityRequest = ChangeMessageVisibilityRequest.builder() + .visibilityTimeout(newVisibilityTimeoutSeconds) + .queueUrl(sqsOptions.getSqsUrl()) + .receiptHandle(parsedMessage.getMessage().receiptHandle()) + .build(); + + LOG.info("Setting visibility timeout for message {} to {}", parsedMessage.getMessage().messageId(), newVisibilityTimeoutSeconds); + sqsClient.changeMessageVisibility(changeMessageVisibilityRequest); + sqsVisibilityTimeoutChangedCount.increment(); + }, + Duration.ofSeconds(progressCheckInterval)); + } } final S3ObjectReference s3ObjectReference = populateS3Reference(parsedMessage.getBucketName(), parsedMessage.getObjectKey()); final Optional deleteMessageBatchRequestEntry = processS3Object(parsedMessage, s3ObjectReference, acknowledgementSet); diff --git a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/configuration/SqsOptions.java b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/configuration/SqsOptions.java index ee2cbd0395..c4acd3abfd 100644 --- a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/configuration/SqsOptions.java +++ b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/configuration/SqsOptions.java @@ -16,7 +16,9 @@ public class SqsOptions { private static final int DEFAULT_MAXIMUM_MESSAGES = 10; + private static final Boolean DEFAULT_VISIBILITY_DUPLICATE_PROTECTION = false; private static final Duration DEFAULT_VISIBILITY_TIMEOUT_SECONDS = Duration.ofSeconds(30); + private static final Duration DEFAULT_MAX_VISIBILITY_TIMEOUT_EXTENSION = Duration.ofSeconds(1800); // 30 minutes private static final Duration DEFAULT_WAIT_TIME_SECONDS = Duration.ofSeconds(20); private static final Duration DEFAULT_POLL_DELAY_SECONDS = Duration.ofSeconds(0); @@ -34,6 +36,14 @@ public class SqsOptions { @DurationMax(seconds = 43200) private Duration visibilityTimeout = DEFAULT_VISIBILITY_TIMEOUT_SECONDS; + @JsonProperty("visibility_duplication_protection") + private Boolean visibilityDuplicateProtection = DEFAULT_VISIBILITY_DUPLICATE_PROTECTION; + + @JsonProperty("max_visibility_timeout_extension") + @DurationMin(seconds = 30) + @DurationMax(seconds = 3600) + private Duration maxVisibilityTimeoutExtension = DEFAULT_MAX_VISIBILITY_TIMEOUT_EXTENSION; + @JsonProperty("wait_time") @DurationMin(seconds = 0) @DurationMax(seconds = 20) @@ -55,6 +65,14 @@ public Duration getVisibilityTimeout() { return visibilityTimeout; } + public Duration getMaxVisibilityTimeoutExtension() { + return maxVisibilityTimeoutExtension; + } + + public Boolean getVisibilityDuplicateProtection() { + return visibilityDuplicateProtection; + } + public Duration getWaitTime() { return waitTime; } diff --git a/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerTest.java b/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerTest.java index 61ac59308d..9fc800eac0 100644 --- a/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerTest.java +++ b/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerTest.java @@ -91,6 +91,7 @@ class SqsWorkerTest { private Timer sqsMessageDelayTimer; private AcknowledgementSetManager acknowledgementSetManager; private AcknowledgementSet acknowledgementSet; + private SqsOptions sqsOptions; @BeforeEach void setUp() { @@ -105,7 +106,7 @@ void setUp() { AwsAuthenticationOptions awsAuthenticationOptions = mock(AwsAuthenticationOptions.class); when(awsAuthenticationOptions.getAwsRegion()).thenReturn(Region.US_EAST_1); - SqsOptions sqsOptions = mock(SqsOptions.class); + sqsOptions = mock(SqsOptions.class); when(sqsOptions.getSqsUrl()).thenReturn("https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue"); when(s3SourceConfig.getAwsAuthenticationOptions()).thenReturn(awsAuthenticationOptions); @@ -217,6 +218,42 @@ void processSqsMessages_should_return_number_of_messages_processed_with_acknowle assertThat(actualDelay, greaterThanOrEqualTo(Duration.ofHours(1).minus(Duration.ofSeconds(5)))); } + @ParameterizedTest + @ValueSource(strings = {"ObjectCreated:Put", "ObjectCreated:Post", "ObjectCreated:Copy", "ObjectCreated:CompleteMultipartUpload"}) + void processSqsMessages_should_return_number_of_messages_processed_with_acknowledgements_and_progress_check(final String eventName) throws IOException { + when(sqsOptions.getVisibilityDuplicateProtection()).thenReturn(true); + when(sqsOptions.getVisibilityTimeout()).thenReturn(Duration.ofSeconds(6)); + when(acknowledgementSetManager.create(any(), any(Duration.class))).thenReturn(acknowledgementSet); + when(s3SourceConfig.getAcknowledgements()).thenReturn(true); + sqsWorker = new SqsWorker(acknowledgementSetManager, sqsClient, s3Service, s3SourceConfig, pluginMetrics, backoff); + Instant startTime = Instant.now().minus(1, ChronoUnit.HOURS); + final Message message = mock(Message.class); + when(message.body()).thenReturn(createEventNotification(eventName, startTime)); + final String testReceiptHandle = UUID.randomUUID().toString(); + when(message.messageId()).thenReturn(testReceiptHandle); + when(message.receiptHandle()).thenReturn(testReceiptHandle); + + final ReceiveMessageResponse receiveMessageResponse = mock(ReceiveMessageResponse.class); + when(sqsClient.receiveMessage(any(ReceiveMessageRequest.class))).thenReturn(receiveMessageResponse); + when(receiveMessageResponse.messages()).thenReturn(Collections.singletonList(message)); + + final int messagesProcessed = sqsWorker.processSqsMessages(); + final ArgumentCaptor deleteMessageBatchRequestArgumentCaptor = ArgumentCaptor.forClass(DeleteMessageBatchRequest.class); + + final ArgumentCaptor durationArgumentCaptor = ArgumentCaptor.forClass(Duration.class); + verify(sqsMessageDelayTimer).record(durationArgumentCaptor.capture()); + Duration actualDelay = durationArgumentCaptor.getValue(); + + assertThat(messagesProcessed, equalTo(1)); + verify(s3Service).addS3Object(any(S3ObjectReference.class), any()); + verify(acknowledgementSetManager).create(any(), any(Duration.class)); + verify(acknowledgementSet).addProgressCheck(any(), any(Duration.class)); + verify(sqsMessagesReceivedCounter).increment(1); + verifyNoInteractions(sqsMessagesDeletedCounter); + assertThat(actualDelay, lessThanOrEqualTo(Duration.ofHours(1).plus(Duration.ofSeconds(5)))); + assertThat(actualDelay, greaterThanOrEqualTo(Duration.ofHours(1).minus(Duration.ofSeconds(5)))); + } + @ParameterizedTest @ValueSource(strings = {"", "{\"foo\": \"bar\""}) void processSqsMessages_should_not_interact_with_S3Service_if_input_is_not_valid_JSON(String inputString) { From 377671cf368e2eaaa1cd58ab894cc985aa411100 Mon Sep 17 00:00:00 2001 From: Roman Kvasnytskyi Date: Sat, 4 Nov 2023 21:24:12 +0100 Subject: [PATCH 25/53] Replace Guava Cache with Caffeine for OpenSearch integration (#3586) Signed-off-by: Roman Kvasnytskyi --- data-prepper-plugins/opensearch/build.gradle | 2 +- .../opensearch/index/DynamicIndexManager.java | 14 +++++------ .../otel-trace-raw-processor/build.gradle | 2 +- .../oteltrace/OTelTraceRawProcessor.java | 25 ++++++++----------- 4 files changed, 19 insertions(+), 24 deletions(-) diff --git a/data-prepper-plugins/opensearch/build.gradle b/data-prepper-plugins/opensearch/build.gradle index e74771543b..17c632a359 100644 --- a/data-prepper-plugins/opensearch/build.gradle +++ b/data-prepper-plugins/opensearch/build.gradle @@ -15,7 +15,6 @@ dependencies { implementation libs.opensearch.java implementation 'com.fasterxml.jackson.core:jackson-databind' implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml' - implementation libs.guava.core implementation 'javax.ws.rs:javax.ws.rs-api:2.1.1' implementation 'software.amazon.awssdk:auth' implementation 'software.amazon.awssdk:http-client-spi' @@ -30,6 +29,7 @@ dependencies { implementation 'software.amazon.awssdk:s3' implementation 'software.amazon.awssdk:opensearchserverless' implementation libs.commons.lang3 + implementation 'com.github.ben-manes.caffeine:caffeine:3.1.8' implementation 'software.amazon.awssdk:apache-client' testImplementation testLibs.junit.vintage testImplementation libs.commons.io diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/DynamicIndexManager.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/DynamicIndexManager.java index a2e613a13b..4bcc58c924 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/DynamicIndexManager.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/DynamicIndexManager.java @@ -5,21 +5,20 @@ package org.opensearch.dataprepper.plugins.sink.opensearch.index; +import com.github.benmanes.caffeine.cache.Cache; +import com.github.benmanes.caffeine.cache.Caffeine; import org.opensearch.client.RestHighLevelClient; import org.opensearch.client.opensearch.OpenSearchClient; import org.opensearch.dataprepper.plugins.sink.opensearch.OpenSearchSinkConfiguration; - -import com.google.common.cache.Cache; -import com.google.common.cache.CacheBuilder; import java.io.IOException; import java.util.concurrent.TimeUnit; import static com.google.common.base.Preconditions.checkNotNull; public class DynamicIndexManager extends AbstractIndexManager { - private Cache indexManagerCache; - final int CACHE_EXPIRE_AFTER_ACCESS_TIME_MINUTES = 30; - final int APPROXIMATE_INDEX_MANAGER_SIZE = 32; + private final Cache indexManagerCache; + private static final int CACHE_EXPIRE_AFTER_ACCESS_TIME_MINUTES = 30; + private static final int APPROXIMATE_INDEX_MANAGER_SIZE = 32; private final long cacheSizeInKB = 1024; protected RestHighLevelClient restHighLevelClient; protected OpenSearchClient openSearchClient; @@ -47,9 +46,8 @@ public DynamicIndexManager(final IndexType indexType, this.restHighLevelClient = restHighLevelClient; this.openSearchSinkConfiguration = openSearchSinkConfiguration; this.clusterSettingsParser = clusterSettingsParser; - CacheBuilder cacheBuilder = CacheBuilder.newBuilder() + Caffeine cacheBuilder = Caffeine.newBuilder() .recordStats() - .concurrencyLevel(1) .maximumWeight(cacheSizeInKB) .expireAfterAccess(CACHE_EXPIRE_AFTER_ACCESS_TIME_MINUTES, TimeUnit.MINUTES) .weigher((k, v) -> APPROXIMATE_INDEX_MANAGER_SIZE); diff --git a/data-prepper-plugins/otel-trace-raw-processor/build.gradle b/data-prepper-plugins/otel-trace-raw-processor/build.gradle index 72e9f05f85..05ef3b5f4b 100644 --- a/data-prepper-plugins/otel-trace-raw-processor/build.gradle +++ b/data-prepper-plugins/otel-trace-raw-processor/build.gradle @@ -18,7 +18,7 @@ dependencies { implementation libs.armeria.grpc implementation 'com.fasterxml.jackson.core:jackson-databind' implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml' - implementation libs.guava.core + implementation 'com.github.ben-manes.caffeine:caffeine:3.1.8' testImplementation testLibs.junit.vintage testImplementation 'org.assertj:assertj-core:3.24.2' testImplementation testLibs.mockito.inline diff --git a/data-prepper-plugins/otel-trace-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/oteltrace/OTelTraceRawProcessor.java b/data-prepper-plugins/otel-trace-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/oteltrace/OTelTraceRawProcessor.java index 7d89ab2d05..a156590107 100644 --- a/data-prepper-plugins/otel-trace-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/oteltrace/OTelTraceRawProcessor.java +++ b/data-prepper-plugins/otel-trace-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/oteltrace/OTelTraceRawProcessor.java @@ -5,6 +5,8 @@ package org.opensearch.dataprepper.plugins.processor.oteltrace; +import com.github.benmanes.caffeine.cache.Cache; +import com.github.benmanes.caffeine.cache.Caffeine; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; @@ -14,8 +16,6 @@ import org.opensearch.dataprepper.model.processor.Processor; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.trace.Span; -import com.google.common.cache.Cache; -import com.google.common.cache.CacheBuilder; import io.micrometer.core.instrument.util.StringUtils; import org.opensearch.dataprepper.plugins.processor.oteltrace.model.SpanSet; import org.opensearch.dataprepper.plugins.processor.oteltrace.model.TraceGroup; @@ -62,14 +62,13 @@ public OTelTraceRawProcessor(final OtelTraceRawProcessorConfig otelTraceRawProce final PluginMetrics pluginMetrics) { super(pluginMetrics); traceFlushInterval = SEC_TO_MILLIS * otelTraceRawProcessorConfig.getTraceFlushIntervalSeconds(); - final int numProcessWorkers = pipelineDescription.getNumberOfProcessWorkers(); - traceIdTraceGroupCache = CacheBuilder.newBuilder() - .concurrencyLevel(numProcessWorkers) - .maximumSize(otelTraceRawProcessorConfig.getTraceGroupCacheMaxSize()) - .expireAfterWrite(otelTraceRawProcessorConfig.getTraceGroupCacheTimeToLive().toMillis(), TimeUnit.MILLISECONDS) - .build(); - - pluginMetrics.gauge(TRACE_GROUP_CACHE_COUNT_METRIC_NAME, traceIdTraceGroupCache, cache -> (double) cache.size()); + traceIdTraceGroupCache = Caffeine.newBuilder() + .maximumSize(otelTraceRawProcessorConfig.getTraceGroupCacheMaxSize()) + .expireAfterWrite(otelTraceRawProcessorConfig.getTraceGroupCacheTimeToLive().toMillis(), TimeUnit.MILLISECONDS) + .build(); + + + pluginMetrics.gauge(TRACE_GROUP_CACHE_COUNT_METRIC_NAME, traceIdTraceGroupCache, cache -> (double) cache.estimatedSize()); pluginMetrics.gauge(SPAN_SET_COUNT_METRIC_NAME, traceIdSpanSetMap, cache -> (double) cache.size()); LOG.info("Configured Trace Raw Processor with a trace flush interval of {} ms.", traceFlushInterval); @@ -109,9 +108,7 @@ private void processSpan(final Span span, final Collection spanSet) { spanSet.addAll(rootSpanAndChildren); } else { final Optional populatedChildSpanOptional = processChildSpan(span); - if (populatedChildSpanOptional.isPresent()) { - spanSet.add(populatedChildSpanOptional.get()); - } + populatedChildSpanOptional.ifPresent(spanSet::add); } } @@ -212,7 +209,7 @@ private List getTracesToFlushByGarbageCollection() { entryIterator.remove(); } } - if (recordsToFlush.size() > 0) { + if (!recordsToFlush.isEmpty()) { LOG.info("Flushing {} records", recordsToFlush.size()); } } finally { From fa691699b8bc91653ce1796a451c19c4f5c9b520 Mon Sep 17 00:00:00 2001 From: Taylor Gray Date: Sun, 5 Nov 2023 10:31:01 -0600 Subject: [PATCH 26/53] =?UTF-8?q?Change=20s3=20scan=20and=20opensearch=20t?= =?UTF-8?q?o=20only=20save=20state=20every=205=20minutes,=20fix=E2=80=A6?= =?UTF-8?q?=20(#3581)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Change s3 scan and opensearch to only save state every 5 minutes, fix bug where any action was valid in OpenSearch sink Signed-off-by: Taylor Gray --- .../dataprepper/model/event/JacksonEvent.java | 1 + .../dynamodb/export/DataFileLoaderTest.java | 69 ++++++++----------- .../worker/NoSearchContextWorker.java | 10 ++- .../source/opensearch/worker/PitWorker.java | 9 ++- .../opensearch/worker/ScrollWorker.java | 10 ++- .../opensearch/worker/WorkerCommonUtils.java | 3 +- .../worker/NoSearchContextWorkerTest.java | 4 +- .../opensearch/worker/PitWorkerTest.java | 6 +- .../opensearch/worker/ScrollWorkerTest.java | 4 +- .../sink/opensearch/BulkRetryStrategy.java | 2 +- .../sink/opensearch/OpenSearchSink.java | 1 - .../opensearch/index/IndexConfiguration.java | 8 ++- .../plugins/source/s3/S3ObjectWorker.java | 13 ++-- .../plugins/source/s3/S3ObjectWorkerTest.java | 6 +- 14 files changed, 81 insertions(+), 65 deletions(-) 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 846c4c3011..1f36e01e25 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 @@ -322,6 +322,7 @@ public static boolean isValidFormatExpressions(final String format, final Expres if (Objects.isNull(expressionEvaluator)) { return false; } + int fromIndex = 0; int position = 0; while ((position = format.indexOf("${", fromIndex)) != -1) { diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java index 7c5c4bee56..a92b80645b 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java @@ -5,11 +5,11 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.export; -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.MockedConstruction; import org.mockito.MockedStatic; import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.dataprepper.buffer.common.BufferAccumulator; @@ -18,8 +18,7 @@ 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.enhanced.EnhancedSourceCoordinator; -import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourcePartition; +import org.opensearch.dataprepper.plugins.source.dynamodb.converter.ExportRecordConverter; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.DataFilePartition; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.state.DataFileProgressState; import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo; @@ -41,14 +40,12 @@ import java.util.zip.GZIPOutputStream; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.anyList; import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.BDDMockito.given; -import static org.mockito.Mockito.doNothing; -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.times; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.opensearch.dataprepper.plugins.source.dynamodb.export.DataFileLoader.BUFFER_TIMEOUT; @@ -57,9 +54,6 @@ @ExtendWith(MockitoExtension.class) class DataFileLoaderTest { - @Mock - private EnhancedSourceCoordinator coordinator; - @Mock private S3Client s3Client; @@ -73,12 +67,13 @@ class DataFileLoaderTest { private BufferAccumulator> bufferAccumulator; @Mock - private Counter testCounter; + private ExportRecordConverter exportRecordConverter; + @Mock + private DataFileCheckpointer checkpointer; - private S3ObjectReader objectReader; - private DataFileCheckpointer checkpointer; + private S3ObjectReader objectReader; private DataFilePartition dataFilePartition; @@ -93,7 +88,6 @@ class DataFileLoaderTest { private final String manifestKey = UUID.randomUUID().toString(); private final String bucketName = UUID.randomUUID().toString(); - private final String prefix = UUID.randomUUID().toString(); private final String exportArn = tableArn + "/export/01693291918297-bfeccbea"; @@ -102,7 +96,7 @@ class DataFileLoaderTest { private final int total = random.nextInt(10); @BeforeEach - void setup() throws Exception { + void setup() { DataFileProgressState state = new DataFileProgressState(); state.setLoaded(0); @@ -120,18 +114,6 @@ void setup() throws Exception { when(s3Client.getObject(any(GetObjectRequest.class))).thenReturn(generateGzipInputStream(total)); objectReader = new S3ObjectReader(s3Client); - - lenient().when(coordinator.createPartition(any(EnhancedSourcePartition.class))).thenReturn(true); - lenient().doNothing().when(coordinator).completePartition(any(EnhancedSourcePartition.class)); - lenient().doNothing().when(coordinator).saveProgressStateForPartition(any(EnhancedSourcePartition.class), eq(null)); - lenient().doNothing().when(coordinator).giveUpPartition(any(EnhancedSourcePartition.class)); - - doNothing().when(bufferAccumulator).add(any(Record.class)); - doNothing().when(bufferAccumulator).flush(); - checkpointer = new DataFileCheckpointer(coordinator, dataFilePartition); - - given(pluginMetrics.counter(anyString())).willReturn(testCounter); - } private ResponseInputStream generateGzipInputStream(int numberOfRecords) { @@ -166,10 +148,13 @@ private ResponseInputStream generateGzipInputStream(int numbe } @Test - void test_run_loadFile_correctly() throws Exception { + void test_run_loadFile_correctly() { DataFileLoader loader; try ( - final MockedStatic bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class)) { + final MockedStatic bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class); + final MockedConstruction recordConverterMockedConstruction = mockConstruction(ExportRecordConverter.class, (mock, context) -> { + exportRecordConverter = mock; + })) { bufferAccumulatorMockedStatic.when(() -> BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT)).thenReturn(bufferAccumulator); loader = DataFileLoader.builder(objectReader, pluginMetrics, buffer) .bucketName(bucketName) @@ -184,23 +169,24 @@ void test_run_loadFile_correctly() throws Exception { // Should call s3 getObject verify(s3Client).getObject(any(GetObjectRequest.class)); - // Should write to buffer - verify(bufferAccumulator, times(total)).add(any(Record.class)); - verify(bufferAccumulator).flush(); + verify(exportRecordConverter).writeToBuffer(eq(null), anyList()); - // Should do one last checkpoint when done. - verify(coordinator).saveProgressStateForPartition(any(DataFilePartition.class), eq(null)); + verify(checkpointer).checkpoint(total); + verify(checkpointer, never()).updateDatafileForAcknowledgmentWait(any(Duration.class)); } @Test - void run_loadFile_with_acknowledgments_processes_correctly() throws Exception { + void run_loadFile_with_acknowledgments_processes_correctly() { final AcknowledgementSet acknowledgementSet = mock(AcknowledgementSet.class); final Duration acknowledgmentTimeout = Duration.ofSeconds(30); DataFileLoader loader; try ( - final MockedStatic bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class)) { + final MockedStatic bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class); + final MockedConstruction recordConverterMockedConstruction = mockConstruction(ExportRecordConverter.class, (mock, context) -> { + exportRecordConverter = mock; + })) { bufferAccumulatorMockedStatic.when(() -> BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT)).thenReturn(bufferAccumulator); loader = DataFileLoader.builder(objectReader, pluginMetrics, buffer) .bucketName(bucketName) @@ -217,12 +203,11 @@ void run_loadFile_with_acknowledgments_processes_correctly() throws Exception { // Should call s3 getObject verify(s3Client).getObject(any(GetObjectRequest.class)); - // Should write to buffer - verify(bufferAccumulator, times(total)).add(any(Record.class)); - verify(bufferAccumulator).flush(); + verify(exportRecordConverter).writeToBuffer(eq(acknowledgementSet), anyList()); + + verify(checkpointer).checkpoint(total); + verify(checkpointer).updateDatafileForAcknowledgmentWait(acknowledgmentTimeout); - // Should do one last checkpoint when done. - verify(coordinator).saveProgressStateForPartition(any(DataFilePartition.class), eq(null)); verify(acknowledgementSet).complete(); } diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorker.java b/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorker.java index 121be0a0b5..b46f50ab3e 100644 --- a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorker.java +++ b/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorker.java @@ -31,6 +31,7 @@ import java.util.Optional; import static org.opensearch.dataprepper.plugins.source.opensearch.worker.WorkerCommonUtils.BACKOFF_ON_EXCEPTION; +import static org.opensearch.dataprepper.plugins.source.opensearch.worker.WorkerCommonUtils.DEFAULT_CHECKPOINT_INTERVAL_MILLS; import static org.opensearch.dataprepper.plugins.source.opensearch.worker.WorkerCommonUtils.calculateExponentialBackoffAndJitter; import static org.opensearch.dataprepper.plugins.source.opensearch.worker.WorkerCommonUtils.completeIndexPartition; import static org.opensearch.dataprepper.plugins.source.opensearch.worker.WorkerCommonUtils.createAcknowledgmentSet; @@ -128,6 +129,8 @@ public void run() { private void processIndex(final SourcePartition openSearchIndexPartition, final AcknowledgementSet acknowledgementSet) { final String indexName = openSearchIndexPartition.getPartitionKey(); + long lastCheckpointTime = System.currentTimeMillis(); + LOG.info("Started processing for index: '{}'", indexName); Optional openSearchIndexProgressStateOptional = openSearchIndexPartition.getPartitionState(); @@ -165,7 +168,12 @@ private void processIndex(final SourcePartition op }); openSearchIndexProgressState.setSearchAfter(searchWithSearchAfterResults.getNextSearchAfter()); - sourceCoordinator.saveProgressStateForPartition(indexName, openSearchIndexProgressState); + + if (System.currentTimeMillis() - lastCheckpointTime > DEFAULT_CHECKPOINT_INTERVAL_MILLS) { + LOG.debug("Renew ownership of index {}", indexName); + sourceCoordinator.saveProgressStateForPartition(indexName, openSearchIndexProgressState); + lastCheckpointTime = System.currentTimeMillis(); + } } while (searchWithSearchAfterResults.getDocuments().size() == searchConfiguration.getBatchSize()); try { diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorker.java b/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorker.java index ee49aee262..9050b6fd87 100644 --- a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorker.java +++ b/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorker.java @@ -35,6 +35,7 @@ import java.util.Optional; import static org.opensearch.dataprepper.plugins.source.opensearch.worker.WorkerCommonUtils.BACKOFF_ON_EXCEPTION; +import static org.opensearch.dataprepper.plugins.source.opensearch.worker.WorkerCommonUtils.DEFAULT_CHECKPOINT_INTERVAL_MILLS; import static org.opensearch.dataprepper.plugins.source.opensearch.worker.WorkerCommonUtils.calculateExponentialBackoffAndJitter; import static org.opensearch.dataprepper.plugins.source.opensearch.worker.WorkerCommonUtils.completeIndexPartition; import static org.opensearch.dataprepper.plugins.source.opensearch.worker.WorkerCommonUtils.createAcknowledgmentSet; @@ -150,6 +151,7 @@ public void run() { private void processIndex(final SourcePartition openSearchIndexPartition, final AcknowledgementSet acknowledgementSet) { final String indexName = openSearchIndexPartition.getPartitionKey(); + long lastCheckpointTime = System.currentTimeMillis(); LOG.info("Starting processing for index: '{}'", indexName); Optional openSearchIndexProgressStateOptional = openSearchIndexPartition.getPartitionState(); @@ -203,7 +205,12 @@ private void processIndex(final SourcePartition op openSearchIndexProgressState.setSearchAfter(searchWithSearchAfterResults.getNextSearchAfter()); openSearchIndexProgressState.setKeepAlive(Duration.ofMillis(openSearchIndexProgressState.getKeepAlive()).plus(EXTEND_KEEP_ALIVE_DURATION).toMillis()); - sourceCoordinator.saveProgressStateForPartition(indexName, openSearchIndexProgressState); + + if (System.currentTimeMillis() - lastCheckpointTime > DEFAULT_CHECKPOINT_INTERVAL_MILLS) { + LOG.debug("Renew ownership of index {}", indexName); + sourceCoordinator.saveProgressStateForPartition(indexName, openSearchIndexProgressState); + lastCheckpointTime = System.currentTimeMillis(); + } } while (searchWithSearchAfterResults.getDocuments().size() == searchConfiguration.getBatchSize()); try { diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorker.java b/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorker.java index ce34521205..be12d40dc6 100644 --- a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorker.java +++ b/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorker.java @@ -34,6 +34,7 @@ import java.util.Optional; import static org.opensearch.dataprepper.plugins.source.opensearch.worker.WorkerCommonUtils.BACKOFF_ON_EXCEPTION; +import static org.opensearch.dataprepper.plugins.source.opensearch.worker.WorkerCommonUtils.DEFAULT_CHECKPOINT_INTERVAL_MILLS; import static org.opensearch.dataprepper.plugins.source.opensearch.worker.WorkerCommonUtils.calculateExponentialBackoffAndJitter; import static org.opensearch.dataprepper.plugins.source.opensearch.worker.WorkerCommonUtils.completeIndexPartition; import static org.opensearch.dataprepper.plugins.source.opensearch.worker.WorkerCommonUtils.createAcknowledgmentSet; @@ -145,6 +146,8 @@ public void run() { private void processIndex(final SourcePartition openSearchIndexPartition, final AcknowledgementSet acknowledgementSet) { final String indexName = openSearchIndexPartition.getPartitionKey(); + long lastCheckpointTime = System.currentTimeMillis(); + LOG.info("Started processing for index: '{}'", indexName); final Integer batchSize = openSearchSourceConfiguration.getSearchConfiguration().getBatchSize(); @@ -168,7 +171,12 @@ private void processIndex(final SourcePartition op .build()); writeDocumentsToBuffer(searchScrollResponse.getDocuments(), acknowledgementSet); - sourceCoordinator.saveProgressStateForPartition(indexName, null); + + if (System.currentTimeMillis() - lastCheckpointTime > DEFAULT_CHECKPOINT_INTERVAL_MILLS) { + LOG.debug("Renew ownership of index {}", indexName); + sourceCoordinator.saveProgressStateForPartition(indexName, null); + lastCheckpointTime = System.currentTimeMillis(); + } } catch (final Exception e) { deleteScroll(createScrollResponse.getScrollId()); throw e; diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtils.java b/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtils.java index 2bde6c0370..7c5a132017 100644 --- a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtils.java +++ b/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtils.java @@ -27,7 +27,8 @@ public class WorkerCommonUtils { static final Duration BACKOFF_ON_EXCEPTION = Duration.ofSeconds(60); - static final Duration ACKNOWLEDGEMENT_SET_TIMEOUT = Duration.ofHours(2); + static final long DEFAULT_CHECKPOINT_INTERVAL_MILLS = 5 * 60_000; + static final Duration ACKNOWLEDGEMENT_SET_TIMEOUT = Duration.ofHours(1); static final Duration STARTING_BACKOFF = Duration.ofMillis(500); static final Duration MAX_BACKOFF = Duration.ofSeconds(60); static final int BACKOFF_RATE = 2; diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorkerTest.java b/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorkerTest.java index 2397aa87b0..e57d40f266 100644 --- a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorkerTest.java +++ b/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorkerTest.java @@ -206,7 +206,7 @@ void run_with_getNextPartition_with_non_empty_partition_processes_and_closes_tha assertThat(executorService.awaitTermination(100, TimeUnit.MILLISECONDS), equalTo(true)); verify(searchAccessor, times(2)).searchWithoutSearchContext(any(NoSearchContextSearchRequest.class)); - verify(sourceCoordinator, times(2)).saveProgressStateForPartition(eq(partitionKey), any(OpenSearchIndexProgressState.class)); + verify(sourceCoordinator, times(0)).saveProgressStateForPartition(eq(partitionKey), any(OpenSearchIndexProgressState.class)); final List noSearchContextSearchRequests = searchRequestArgumentCaptor.getAllValues(); assertThat(noSearchContextSearchRequests.size(), equalTo(2)); @@ -283,7 +283,7 @@ void run_with_getNextPartition_with_acknowledgments_processes_and_closes_that_pa assertThat(executorService.awaitTermination(100, TimeUnit.MILLISECONDS), equalTo(true)); verify(searchAccessor, times(2)).searchWithoutSearchContext(any(NoSearchContextSearchRequest.class)); - verify(sourceCoordinator, times(2)).saveProgressStateForPartition(eq(partitionKey), any(OpenSearchIndexProgressState.class)); + verify(sourceCoordinator, times(0)).saveProgressStateForPartition(eq(partitionKey), any(OpenSearchIndexProgressState.class)); final List noSearchContextSearchRequests = searchRequestArgumentCaptor.getAllValues(); assertThat(noSearchContextSearchRequests.size(), equalTo(2)); diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorkerTest.java b/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorkerTest.java index 7784f7ddff..a333925137 100644 --- a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorkerTest.java +++ b/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorkerTest.java @@ -191,7 +191,7 @@ void run_with_getNextPartition_with_non_empty_partition_creates_and_deletes_pit_ assertThat(createPointInTimeRequest.getKeepAlive(), equalTo(STARTING_KEEP_ALIVE)); verify(searchAccessor, times(2)).searchWithPit(any(SearchPointInTimeRequest.class)); - verify(sourceCoordinator, times(2)).saveProgressStateForPartition(eq(partitionKey), any(OpenSearchIndexProgressState.class)); + verify(sourceCoordinator, times(0)).saveProgressStateForPartition(eq(partitionKey), any(OpenSearchIndexProgressState.class)); final List searchPointInTimeRequestList = searchPointInTimeRequestArgumentCaptor.getAllValues(); assertThat(searchPointInTimeRequestList.size(), equalTo(2)); @@ -292,7 +292,7 @@ void run_with_acknowledgments_enabled_creates_and_deletes_pit_and_closes_that_pa assertThat(createPointInTimeRequest.getKeepAlive(), equalTo(STARTING_KEEP_ALIVE)); verify(searchAccessor, times(2)).searchWithPit(any(SearchPointInTimeRequest.class)); - verify(sourceCoordinator, times(2)).saveProgressStateForPartition(eq(partitionKey), any(OpenSearchIndexProgressState.class)); + verify(sourceCoordinator, times(0)).saveProgressStateForPartition(eq(partitionKey), any(OpenSearchIndexProgressState.class)); final List searchPointInTimeRequestList = searchPointInTimeRequestArgumentCaptor.getAllValues(); assertThat(searchPointInTimeRequestList.size(), equalTo(2)); @@ -378,7 +378,7 @@ void run_with_getNextPartition_with_valid_existing_point_in_time_does_not_create verify(searchAccessor, never()).createPit(any(CreatePointInTimeRequest.class)); verify(searchAccessor, times(2)).searchWithPit(any(SearchPointInTimeRequest.class)); - verify(sourceCoordinator, times(2)).saveProgressStateForPartition(eq(partitionKey), eq(openSearchIndexProgressState)); + verify(sourceCoordinator, times(0)).saveProgressStateForPartition(eq(partitionKey), eq(openSearchIndexProgressState)); verify(sourceCoordinator, times(0)).updatePartitionForAcknowledgmentWait(anyString(), any(Duration.class)); verify(documentsProcessedCounter, times(3)).increment(); diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorkerTest.java b/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorkerTest.java index 63f88c272c..ffd24b5972 100644 --- a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorkerTest.java +++ b/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorkerTest.java @@ -191,7 +191,7 @@ void run_with_getNextPartition_with_non_empty_partition_creates_and_deletes_scro assertThat(createScrollRequest.getScrollTime(), equalTo(SCROLL_TIME_PER_BATCH)); verify(searchAccessor, times(2)).searchWithScroll(any(SearchScrollRequest.class)); - verify(sourceCoordinator, times(2)).saveProgressStateForPartition(eq(partitionKey), eq(null)); + verify(sourceCoordinator, times(0)).saveProgressStateForPartition(eq(partitionKey), eq(null)); final List searchScrollRequests = searchScrollRequestArgumentCaptor.getAllValues(); assertThat(searchScrollRequests.size(), equalTo(2)); @@ -286,7 +286,7 @@ void run_with_getNextPartition_with_acknowledgments_creates_and_deletes_scroll_a assertThat(createScrollRequest.getScrollTime(), equalTo(SCROLL_TIME_PER_BATCH)); verify(searchAccessor, times(2)).searchWithScroll(any(SearchScrollRequest.class)); - verify(sourceCoordinator, times(2)).saveProgressStateForPartition(eq(partitionKey), eq(null)); + verify(sourceCoordinator, times(0)).saveProgressStateForPartition(eq(partitionKey), eq(null)); final List searchScrollRequests = searchScrollRequestArgumentCaptor.getAllValues(); assertThat(searchScrollRequests.size(), equalTo(2)); diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java index 9f8f501fb6..793029f6ef 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java @@ -259,7 +259,7 @@ private void handleFailures(final AccumulatingBulkRequest> actions, final Expres for (final Map actionMap: actions) { String action = (String)actionMap.get("type"); if (action != null) { - checkArgument((EnumUtils.isValidEnumIgnoreCase(OpenSearchBulkActions.class, action) || JacksonEvent.isValidFormatExpressions(action, expressionEvaluator)), "action must be one of the following: " + OpenSearchBulkActions.values()); + checkArgument((EnumUtils.isValidEnumIgnoreCase(OpenSearchBulkActions.class, action) || + (action.contains("${") && JacksonEvent.isValidFormatExpressions(action, expressionEvaluator))), "Invalid action \"" + action + "\". action must be one of the following: " + Arrays.stream(OpenSearchBulkActions.values()).collect(Collectors.toList())); } } this.actions = actions; diff --git a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3ObjectWorker.java b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3ObjectWorker.java index 024cb05903..0397183877 100644 --- a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3ObjectWorker.java +++ b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3ObjectWorker.java @@ -22,6 +22,7 @@ import java.time.Duration; import java.util.concurrent.Callable; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; import java.util.function.BiConsumer; /** @@ -30,7 +31,8 @@ */ class S3ObjectWorker implements S3ObjectHandler { private static final Logger LOG = LoggerFactory.getLogger(S3ObjectWorker.class); - static final int RECORDS_TO_ACCUMULATE_TO_SAVE_STATE = 10_000; + private static final long DEFAULT_CHECKPOINT_INTERVAL_MILLS = 5 * 60_000; + private final S3Client s3Client; private final Buffer> buffer; @@ -83,6 +85,7 @@ private void doParseObject(final AcknowledgementSet acknowledgementSet, final long totalBytesRead; LOG.info("Read S3 object: {}", s3ObjectReference); + AtomicLong lastCheckpointTime = new AtomicLong(System.currentTimeMillis()); final S3InputFile inputFile = new S3InputFile(s3Client, s3ObjectReference, bucketOwnerProvider, s3ObjectPluginMetrics); @@ -104,10 +107,12 @@ private void doParseObject(final AcknowledgementSet acknowledgementSet, acknowledgementSet.add(record.getData()); } bufferAccumulator.add(record); - int recordsWrittenAfterLastSaveState = bufferAccumulator.getTotalWritten() - saveStateCounter.get() * RECORDS_TO_ACCUMULATE_TO_SAVE_STATE; - // Saving state to renew source coordination ownership for every 10,000 records, ownership time is 10 minutes - if (recordsWrittenAfterLastSaveState >= RECORDS_TO_ACCUMULATE_TO_SAVE_STATE && sourceCoordinator != null && partitionKey != null) { + + if (sourceCoordinator != null && partitionKey != null && + (System.currentTimeMillis() - lastCheckpointTime.get() > DEFAULT_CHECKPOINT_INTERVAL_MILLS)) { + LOG.debug("Renew partition ownership for the object {}", partitionKey); sourceCoordinator.saveProgressStateForPartition(partitionKey, null); + lastCheckpointTime.set(System.currentTimeMillis()); saveStateCounter.getAndIncrement(); } } catch (final Exception e) { diff --git a/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ObjectWorkerTest.java b/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ObjectWorkerTest.java index 2cbd5bbf47..eadc9a7b8b 100644 --- a/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ObjectWorkerTest.java +++ b/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ObjectWorkerTest.java @@ -59,10 +59,10 @@ 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.verifyNoInteractions; import static org.mockito.Mockito.when; -import static org.opensearch.dataprepper.plugins.source.s3.S3ObjectWorker.RECORDS_TO_ACCUMULATE_TO_SAVE_STATE; @ExtendWith(MockitoExtension.class) class S3ObjectWorkerTest { @@ -306,14 +306,12 @@ void parseS3Object_calls_Codec_parse_with_Consumer_that_adds_to_BufferAccumulato final Record record = mock(Record.class); final Event event = mock(Event.class); when(record.getData()).thenReturn(event); - when(bufferAccumulator.getTotalWritten()).thenReturn(RECORDS_TO_ACCUMULATE_TO_SAVE_STATE + 1); - consumerUnderTest.accept(record); final InOrder inOrder = inOrder(eventConsumer, bufferAccumulator, sourceCoordinator); inOrder.verify(eventConsumer).accept(event, s3ObjectReference); inOrder.verify(bufferAccumulator).add(record); - inOrder.verify(sourceCoordinator).saveProgressStateForPartition(testPartitionKey, null); + inOrder.verify(sourceCoordinator, times(0)).saveProgressStateForPartition(testPartitionKey, null); } @Test From 60f84ad85e90d946c0f0d44f7b41d2d5d4521301 Mon Sep 17 00:00:00 2001 From: Taylor Gray Date: Sun, 5 Nov 2023 10:49:05 -0600 Subject: [PATCH 27/53] Add dynamodb_event_name metadata attribute, change mapping for Ddb INSERT and MODIFY to be index bulk action (#3585) Signed-off-by: Taylor Gray --- .../source/dynamodb/converter/MetadataKeyAttributes.java | 6 ++++-- .../plugins/source/dynamodb/converter/RecordConverter.java | 5 +++-- .../dynamodb/converter/ExportRecordConverterTest.java | 3 +++ .../dynamodb/converter/StreamRecordConverterTest.java | 4 +++- 4 files changed, 13 insertions(+), 5 deletions(-) diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java index 204ec82e02..03c92c09e0 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java @@ -12,9 +12,11 @@ public class MetadataKeyAttributes { static final String SORT_KEY_METADATA_ATTRIBUTE = "sort_key"; - static final String EVENT_TIMESTAMP_METADATA_ATTRIBUTE = "ts"; + static final String EVENT_TIMESTAMP_METADATA_ATTRIBUTE = "dynamodb_timestamp"; - static final String EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE = "op"; + static final String EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE = "opensearch_action"; + + static final String DDB_STREAM_EVENT_NAME_METADATA_ATTRIBUTE = "dynamodb_event_name"; static final String EVENT_TABLE_NAME_METADATA_ATTRIBUTE = "table_name"; } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java index de79feadfe..d15a18b989 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java @@ -16,6 +16,7 @@ import java.util.Map; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.DDB_STREAM_EVENT_NAME_METADATA_ATTRIBUTE; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_TABLE_NAME_METADATA_ATTRIBUTE; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_TIMESTAMP_METADATA_ATTRIBUTE; @@ -80,6 +81,7 @@ public void addToBuffer(final AcknowledgementSet acknowledgementSet, Map Date: Mon, 6 Nov 2023 09:57:04 -0800 Subject: [PATCH 28/53] Enable acknowledgements for KafkaBuffer plugin (#3584) Signed-off-by: Hardeep --- .../dataprepper/plugins/kafka/buffer/KafkaBufferConfig.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferConfig.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferConfig.java index 0e0eea2837..9574ae94e2 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferConfig.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferConfig.java @@ -125,7 +125,7 @@ public String getClientDnsLookup() { @Override public boolean getAcknowledgementsEnabled() { - return false; + return true; } public Duration getDrainTimeout() { From a40f4b8f1c075c6d6e42101e53dc1858a6a32df2 Mon Sep 17 00:00:00 2001 From: kkondaka <41027584+kkondaka@users.noreply.github.com> Date: Tue, 7 Nov 2023 09:42:58 -0800 Subject: [PATCH 29/53] Add Internal and external latency to OpenSearch and S3 sinks. (#3583) Add Internal and external latency to OpenSearch and S3 sinks Signed-off-by: Krishna Kondaka Co-authored-by: Krishna Kondaka --- .../model/event/DefaultEventHandle.java | 13 ++++ .../dataprepper/model/event/EventHandle.java | 2 + .../dataprepper/model/sink/AbstractSink.java | 17 ++++ .../dataprepper/model/sink/Sink.java | 8 ++ .../model/sink/SinkLatencyMetrics.java | 33 ++++++++ .../model/event/DefaultEventHandleTests.java | 13 ++++ .../model/sink/AbstractSinkTest.java | 64 +++++++++++++++ .../model/sink/SinkLatencyMetricsTest.java | 71 +++++++++++++++++ .../dataprepper/model/sink/SinkTest.java | 43 ++++++++++ .../dataprepper/pipeline/Pipeline.java | 7 +- data-prepper-plugins/date-processor/README.md | 2 + .../date-processor/build.gradle | 1 + .../plugins/processor/date/DateProcessor.java | 20 +++-- .../processor/date/DateProcessorConfig.java | 8 ++ .../date/DateProcessorConfigTest.java | 7 +- .../processor/date/DateProcessorTests.java | 32 +++++++- .../sink/opensearch/BulkOperationWrapper.java | 13 +--- .../opensearch/BulkRetryStrategyTests.java | 78 +++++++++++++------ ...cumulatingUncompressedBulkRequestTest.java | 2 +- 19 files changed, 388 insertions(+), 46 deletions(-) create mode 100644 data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetrics.java create mode 100644 data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetricsTest.java create mode 100644 data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/SinkTest.java diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java index 86278bea74..cb7c12da84 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java @@ -8,6 +8,9 @@ import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; import java.lang.ref.WeakReference; +import java.util.ArrayList; +import java.util.List; +import java.util.function.BiConsumer; import java.time.Instant; import java.io.Serializable; @@ -15,11 +18,13 @@ public class DefaultEventHandle implements EventHandle, InternalEventHandle, Ser private Instant externalOriginationTime; private final Instant internalOriginationTime; private WeakReference acknowledgementSetRef; + private List> releaseConsumers; public DefaultEventHandle(final Instant internalOriginationTime) { this.acknowledgementSetRef = null; this.externalOriginationTime = null; this.internalOriginationTime = internalOriginationTime; + this.releaseConsumers = new ArrayList<>(); } @Override @@ -51,9 +56,17 @@ public Instant getExternalOriginationTime() { @Override public void release(boolean result) { + for (final BiConsumer consumer: releaseConsumers) { + consumer.accept(this, result); + } AcknowledgementSet acknowledgementSet = getAcknowledgementSet(); if (acknowledgementSet != null) { acknowledgementSet.release(this, result); } } + + @Override + public void onRelease(BiConsumer releaseConsumer) { + releaseConsumers.add(releaseConsumer); + } } diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java index 64ef3be574..6049a196e3 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java @@ -6,6 +6,7 @@ package org.opensearch.dataprepper.model.event; import java.time.Instant; +import java.util.function.BiConsumer; public interface EventHandle { /** @@ -41,5 +42,6 @@ public interface EventHandle { */ Instant getInternalOriginationTime(); + void onRelease(BiConsumer releaseConsumer); } diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/AbstractSink.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/AbstractSink.java index f116d89e9e..1c3e596265 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/AbstractSink.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/AbstractSink.java @@ -11,6 +11,7 @@ import org.opensearch.dataprepper.metrics.PluginMetrics; 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; @@ -22,6 +23,7 @@ public abstract class AbstractSink> implements Sink { protected static final int DEFAULT_WAIT_TIME_MS = 1000; protected final PluginMetrics pluginMetrics; private final Counter recordsInCounter; + private final SinkLatencyMetrics latencyMetrics; private final Timer timeElapsedTimer; private Thread retryThread; private int maxRetries; @@ -31,6 +33,7 @@ public AbstractSink(final PluginSetting pluginSetting, int numRetries, int waitT this.pluginMetrics = PluginMetrics.fromPluginSetting(pluginSetting); recordsInCounter = pluginMetrics.counter(MetricNames.RECORDS_IN); timeElapsedTimer = pluginMetrics.timer(MetricNames.TIME_ELAPSED); + this.latencyMetrics = new SinkLatencyMetrics(pluginMetrics); retryThread = null; this.maxRetries = numRetries; this.waitTimeMs = waitTimeMs; @@ -77,6 +80,20 @@ public void shutdown() { } } + @Override + public void updateLatencyMetrics(Collection records) { + for (final Record record : records) { + if (record.getData() instanceof Event) { + Event event = (Event)record.getData(); + event.getEventHandle().onRelease((eventHandle, result) -> { + if (result) { + latencyMetrics.update(eventHandle); + } + }); + } + } + } + Thread.State getRetryThreadState() { if (retryThread != null) { return retryThread.getState(); diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/Sink.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/Sink.java index 0ce6fa5ac1..178566ba5b 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/Sink.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/Sink.java @@ -38,4 +38,12 @@ public interface Sink> { */ boolean isReady(); + /** + * updates latency metrics of sink + * + * @param events list of events used for updating the latency metrics + */ + default void updateLatencyMetrics(final Collection events) { + } + } diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetrics.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetrics.java new file mode 100644 index 0000000000..31b584271f --- /dev/null +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetrics.java @@ -0,0 +1,33 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.model.sink; + +import org.opensearch.dataprepper.metrics.PluginMetrics; +import io.micrometer.core.instrument.DistributionSummary; +import org.opensearch.dataprepper.model.event.EventHandle; + +import java.time.Duration; +import java.time.Instant; + +public class SinkLatencyMetrics { + public static final String INTERNAL_LATENCY = "internalLatency"; + public static final String EXTERNAL_LATENCY = "externalLatency"; + private final DistributionSummary internalLatencySummary; + private final DistributionSummary externalLatencySummary; + + public SinkLatencyMetrics(PluginMetrics pluginMetrics) { + internalLatencySummary = pluginMetrics.summary(INTERNAL_LATENCY); + externalLatencySummary = pluginMetrics.summary(EXTERNAL_LATENCY); + } + public void update(final EventHandle eventHandle) { + Instant now = Instant.now(); + internalLatencySummary.record(Duration.between(eventHandle.getInternalOriginationTime(), now).toMillis()); + if (eventHandle.getExternalOriginationTime() == null) { + return; + } + externalLatencySummary.record(Duration.between(eventHandle.getExternalOriginationTime(), now).toMillis()); + } +} diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventHandleTests.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventHandleTests.java index 3f1af4fb26..b2a66b2d1d 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventHandleTests.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventHandleTests.java @@ -20,6 +20,7 @@ class DefaultEventHandleTests { @Mock private AcknowledgementSet acknowledgementSet; + private int count; @Test void testBasic() { @@ -56,4 +57,16 @@ void testWithExternalOriginationTime() { assertThat(eventHandle.getExternalOriginationTime(), equalTo(now.minusSeconds(60))); eventHandle.release(true); } + + @Test + void testWithOnReleaseHandler() { + Instant now = Instant.now(); + count = 0; + DefaultEventHandle eventHandle = new DefaultEventHandle(now); + eventHandle.onRelease((handle, result) -> {if (result) count++; }); + eventHandle.release(true); + assertThat(count, equalTo(1)); + + } + } diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/AbstractSinkTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/AbstractSinkTest.java index e4b19cf7ca..6d58f7dd71 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/AbstractSinkTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/AbstractSinkTest.java @@ -13,6 +13,12 @@ import org.opensearch.dataprepper.metrics.MetricsTestUtil; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.event.JacksonEvent; +import org.opensearch.dataprepper.model.event.EventHandle; + +import static org.mockito.Mockito.when; +import static org.mockito.Mockito.mock; import java.time.Duration; import java.util.Arrays; @@ -25,6 +31,7 @@ import static org.awaitility.Awaitility.await; public class AbstractSinkTest { + private int count; @Test public void testMetrics() { final String sinkName = "testSink"; @@ -35,6 +42,8 @@ public void testMetrics() { AbstractSink> abstractSink = new AbstractSinkImpl(pluginSetting); abstractSink.initialize(); Assert.assertEquals(abstractSink.isReady(), true); + abstractSink.updateLatencyMetrics(Arrays.asList( + new Record<>(UUID.randomUUID().toString()))); abstractSink.output(Arrays.asList( new Record<>(UUID.randomUUID().toString()), new Record<>(UUID.randomUUID().toString()), @@ -80,6 +89,61 @@ public void testSinkNotReady() { abstractSink.shutdown(); } + @Test + public void testSinkWithRegisterEventReleaseHandler() { + final String sinkName = "testSink"; + final String pipelineName = "pipelineName"; + MetricsTestUtil.initMetrics(); + PluginSetting pluginSetting = new PluginSetting(sinkName, Collections.emptyMap()); + pluginSetting.setPipelineName(pipelineName); + AbstractSink> abstractSink = new AbstractEventSinkImpl(pluginSetting); + abstractSink.initialize(); + Assert.assertEquals(abstractSink.isReady(), true); + count = 0; + Event event = JacksonEvent.builder() + .withEventType("event") + .build(); + Record record = mock(Record.class); + EventHandle eventHandle = mock(EventHandle.class); + when(record.getData()).thenReturn(event); + + abstractSink.updateLatencyMetrics(Arrays.asList(record)); + abstractSink.output(Arrays.asList(record)); + await().atMost(Duration.ofSeconds(5)) + .until(abstractSink::isReady); + abstractSink.shutdown(); + } + + private static class AbstractEventSinkImpl extends AbstractSink> { + + public AbstractEventSinkImpl(PluginSetting pluginSetting) { + super(pluginSetting, 10, 1000); + } + + @Override + public void doOutput(Collection> records) { + for (final Record record: records) { + Event event = record.getData(); + event.getEventHandle().release(true); + } + } + + @Override + public void shutdown() { + super.shutdown(); + } + + @Override + public void doInitialize() { + } + + @Override + public boolean isReady() { + return true; + } + } + + private static class AbstractSinkImpl extends AbstractSink> { public AbstractSinkImpl(PluginSetting pluginSetting) { diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetricsTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetricsTest.java new file mode 100644 index 0000000000..4cf5043cae --- /dev/null +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetricsTest.java @@ -0,0 +1,71 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.model.sink; + +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.event.EventHandle; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import io.micrometer.core.instrument.DistributionSummary; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.greaterThanOrEqualTo; + +import java.time.Instant; + +class SinkLatencyMetricsTest { + + private PluginMetrics pluginMetrics; + private EventHandle eventHandle; + private SinkLatencyMetrics latencyMetrics; + private DistributionSummary internalLatencySummary; + private DistributionSummary externalLatencySummary; + + public SinkLatencyMetrics createObjectUnderTest() { + return new SinkLatencyMetrics(pluginMetrics); + } + + @BeforeEach + void setup() { + pluginMetrics = mock(PluginMetrics.class); + SimpleMeterRegistry registry = new SimpleMeterRegistry(); + internalLatencySummary = DistributionSummary + .builder("internalLatency") + .baseUnit("milliseconds") + .register(registry); + externalLatencySummary = DistributionSummary + .builder("externalLatency") + .baseUnit("milliseconds") + .register(registry); + when(pluginMetrics.summary(SinkLatencyMetrics.INTERNAL_LATENCY)).thenReturn(internalLatencySummary); + when(pluginMetrics.summary(SinkLatencyMetrics.EXTERNAL_LATENCY)).thenReturn(externalLatencySummary); + eventHandle = mock(EventHandle.class); + when(eventHandle.getInternalOriginationTime()).thenReturn(Instant.now()); + latencyMetrics = createObjectUnderTest(); + } + + @Test + public void testInternalOriginationTime() { + latencyMetrics.update(eventHandle); + assertThat(internalLatencySummary.count(), equalTo(1L)); + } + + @Test + public void testExternalOriginationTime() { + when(eventHandle.getExternalOriginationTime()).thenReturn(Instant.now().minusMillis(10)); + latencyMetrics.update(eventHandle); + assertThat(internalLatencySummary.count(), equalTo(1L)); + assertThat(externalLatencySummary.count(), equalTo(1L)); + assertThat(externalLatencySummary.max(), greaterThanOrEqualTo(10.0)); + } +} + + diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/SinkTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/SinkTest.java new file mode 100644 index 0000000000..5f66a623aa --- /dev/null +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/sink/SinkTest.java @@ -0,0 +1,43 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.model.sink; + +import org.opensearch.dataprepper.model.record.Record; +import org.junit.jupiter.api.Test; + +import java.util.Collection; +import java.util.Collections; + +public class SinkTest { + private static class SinkTestClass implements Sink> { + + @Override + public boolean isReady() { + return true; + } + + @Override + public void shutdown() { + } + + @Override + public void initialize() { + } + + @Override + public void output(Collection> records) { + } + + }; + + SinkTestClass sink; + + @Test + public void testSinkUpdateLatencyMetrics() { + sink = new SinkTestClass(); + sink.updateLatencyMetrics(Collections.emptyList()); + } +} diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/Pipeline.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/Pipeline.java index 0bd5d469ca..a71f8d14b1 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/Pipeline.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/Pipeline.java @@ -344,8 +344,11 @@ List> publishToSinks(final Collection records) { InactiveAcknowledgementSetManager.getInstance(), sinks); router.route(records, sinks, getRecordStrategy, (sink, events) -> - sinkFutures.add(sinkExecutorService.submit(() -> sink.output(events), null)) - ); + sinkFutures.add(sinkExecutorService.submit(() -> { + sink.updateLatencyMetrics(events); + sink.output(events); + }, null)) + ); return sinkFutures; } } diff --git a/data-prepper-plugins/date-processor/README.md b/data-prepper-plugins/date-processor/README.md index a5c147ecd6..c4d9acd4f5 100644 --- a/data-prepper-plugins/date-processor/README.md +++ b/data-prepper-plugins/date-processor/README.md @@ -104,6 +104,8 @@ processor: * Type: String * Default: `Locale.ROOT` +* `to_origination_metadata` (Optional): When this option is used, matched time is put into the event's metadata as an instance of `Instant`. + ## Metrics * `dateProcessingMatchSuccessCounter`: Number of records that match with at least one pattern specified in match configuration option. diff --git a/data-prepper-plugins/date-processor/build.gradle b/data-prepper-plugins/date-processor/build.gradle index 6433af9b6d..743761201d 100644 --- a/data-prepper-plugins/date-processor/build.gradle +++ b/data-prepper-plugins/date-processor/build.gradle @@ -12,5 +12,6 @@ dependencies { implementation project(':data-prepper-test-common') implementation 'com.fasterxml.jackson.core:jackson-databind' implementation 'io.micrometer:micrometer-core' + implementation libs.commons.lang3 testImplementation libs.commons.lang3 } 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 b9094926ee..d0808053b2 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 @@ -27,6 +27,7 @@ import java.util.List; import java.util.Objects; 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> { @@ -71,7 +72,16 @@ public Collection> doExecute(Collection> records) { zonedDateTime = getDateTimeFromTimeReceived(record); else if (keyToParse != null && !keyToParse.isEmpty()) { - zonedDateTime = getDateTimeFromMatch(record); + 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); } @@ -119,7 +129,7 @@ private String getDateTimeFromTimeReceived(final Record record) { return timeReceived.atZone(dateProcessorConfig.getDestinationZoneId()).format(getOutputFormatter()); } - private String getDateTimeFromMatch(final Record record) { + private Pair getDateTimeFromMatch(final Record record) { final String sourceTimestamp = getSourceTimestamp(record); if (sourceTimestamp == null) return null; @@ -136,12 +146,12 @@ private String getSourceTimestamp(final Record record) { } } - private String getFormattedDateTimeString(final String sourceTimestamp) { + private Pair getFormattedDateTimeString(final String sourceTimestamp) { for (DateTimeFormatter formatter : dateTimeFormatters) { try { - return ZonedDateTime.parse(sourceTimestamp, formatter).format(getOutputFormatter().withZone(dateProcessorConfig.getDestinationZoneId())); + ZonedDateTime tmp = ZonedDateTime.parse(sourceTimestamp, formatter); + return Pair.of(tmp.format(getOutputFormatter().withZone(dateProcessorConfig.getDestinationZoneId())), tmp.toInstant()); } catch (Exception ignored) { - } } 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 fd6c1f25be..5e06b48cbb 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 @@ -15,6 +15,7 @@ public class DateProcessorConfig { static final Boolean DEFAULT_FROM_TIME_RECEIVED = false; + static final Boolean DEFAULT_TO_ORIGINATION_METADATA = false; static final String DEFAULT_DESTINATION = "@timestamp"; static final String DEFAULT_SOURCE_TIMEZONE = ZoneId.systemDefault().toString(); static final String DEFAULT_DESTINATION_TIMEZONE = ZoneId.systemDefault().toString(); @@ -45,6 +46,9 @@ public List getPatterns() { @JsonProperty("from_time_received") private Boolean fromTimeReceived = DEFAULT_FROM_TIME_RECEIVED; + @JsonProperty("to_origination_metadata") + private Boolean toOriginationMetadata = DEFAULT_TO_ORIGINATION_METADATA; + @JsonProperty("match") private List match; @@ -76,6 +80,10 @@ public Boolean getFromTimeReceived() { return fromTimeReceived; } + public Boolean getToOriginationMetadata() { + return toOriginationMetadata; + } + public List getMatch() { return match; } 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 0959c9db2f..db604039fa 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 @@ -51,8 +51,9 @@ void isValidMatchAndFromTimestampReceived_should_return_true_if_from_time_receiv } @Test - void isValidMatchAndFromTimestampReceived_should_return_false_if_from_time_received_and_match_are_not_configured() { - assertThat(dateProcessorConfig.isValidMatchAndFromTimestampReceived(), equalTo(false)); + void testToOriginationMetadata_should_return_true() throws NoSuchFieldException, IllegalAccessException { + reflectivelySetField(dateProcessorConfig, "toOriginationMetadata", true); + assertThat(dateProcessorConfig.getToOriginationMetadata(), equalTo(true)); } @Test @@ -178,4 +179,4 @@ private void reflectivelySetField(final DateProcessorConfig dateProcessorConfig, field.setAccessible(false); } } -} \ No newline at end of file +} 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 a45daeb56b..ce3906a635 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 @@ -83,7 +83,8 @@ void setup() { lenient().when(pluginMetrics.counter(DateProcessor.DATE_PROCESSING_MATCH_SUCCESS)).thenReturn(dateProcessingMatchSuccessCounter); lenient().when(pluginMetrics.counter(DateProcessor.DATE_PROCESSING_MATCH_FAILURE)).thenReturn(dateProcessingMatchFailureCounter); when(mockDateProcessorConfig.getDateWhen()).thenReturn(null); - expectedDateTime = LocalDateTime.now(); + expectedInstant = Instant.now(); + expectedDateTime = LocalDateTime.ofInstant(expectedInstant, ZoneId.systemDefault()); } @AfterEach @@ -361,6 +362,35 @@ void match_with_different_year_formats_test(String pattern) { verify(dateProcessingMatchSuccessCounter, times(1)).increment(); } + @ParameterizedTest + @ValueSource(strings = {"yyyy MM dd HH mm ss"}) + void match_with_to_origination_metadata(String pattern) { + when(mockDateMatch.getKey()).thenReturn("logDate"); + when(mockDateMatch.getPatterns()).thenReturn(Collections.singletonList(pattern)); + + List dateMatches = Collections.singletonList(mockDateMatch); + when(mockDateProcessorConfig.getMatch()).thenReturn(dateMatches); + when(mockDateProcessorConfig.getSourceZoneId()).thenReturn(ZoneId.systemDefault()); + when(mockDateProcessorConfig.getDestinationZoneId()).thenReturn(ZoneId.systemDefault()); + when(mockDateProcessorConfig.getSourceLocale()).thenReturn(Locale.ROOT); + when(mockDateProcessorConfig.getToOriginationMetadata()).thenReturn(true); + + dateProcessor = createObjectUnderTest(); + + Map testData = getTestData(); + testData.put("logDate", expectedDateTime.format(DateTimeFormatter.ofPattern(pattern))); + + final Record record = buildRecordWithEvent(testData); + final List> processedRecords = (List>) dateProcessor.doExecute(Collections.singletonList(record)); + + Event event = (Event)processedRecords.get(0).getData(); + Assertions.assertTrue(event.getMetadata().getExternalOriginationTime() != null); + Assertions.assertTrue(event.getEventHandle().getExternalOriginationTime() != null); + ZonedDateTime expectedZonedDatetime = expectedDateTime.atZone(mockDateProcessorConfig.getSourceZoneId()).truncatedTo(ChronoUnit.SECONDS); + Assertions.assertTrue(expectedZonedDatetime.equals(event.getMetadata().getExternalOriginationTime().atZone(mockDateProcessorConfig.getSourceZoneId()))); + verify(dateProcessingMatchSuccessCounter, times(1)).increment(); + } + @ParameterizedTest @ValueSource(strings = {"MMM/dd", "MM dd"}) void match_without_year_test(String pattern) { diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkOperationWrapper.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkOperationWrapper.java index f4c1ebb0b9..7f11db2234 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkOperationWrapper.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkOperationWrapper.java @@ -47,9 +47,7 @@ public class BulkOperationWrapper { private final SerializedJson jsonNode; public BulkOperationWrapper(final BulkOperation bulkOperation) { - this.bulkOperation = bulkOperation; - this.eventHandle = null; - this.jsonNode = null; + this(bulkOperation, null, null); } public BulkOperationWrapper(final BulkOperation bulkOperation, final EventHandle eventHandle, final SerializedJson jsonNode) { @@ -60,10 +58,7 @@ public BulkOperationWrapper(final BulkOperation bulkOperation, final EventHandle } public BulkOperationWrapper(final BulkOperation bulkOperation, final EventHandle eventHandle) { - checkNotNull(bulkOperation); - this.bulkOperation = bulkOperation; - this.eventHandle = eventHandle; - this.jsonNode = null; + this(bulkOperation, eventHandle, null); } public BulkOperation getBulkOperation() { @@ -75,9 +70,7 @@ public EventHandle getEventHandle() { } public void releaseEventHandle(boolean result) { - if (eventHandle != null) { - eventHandle.release(result); - } + eventHandle.release(result); } public Object getDocument() { diff --git a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategyTests.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategyTests.java index 7c05b99aa9..b2556e3192 100644 --- a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategyTests.java +++ b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategyTests.java @@ -37,6 +37,7 @@ import java.util.StringJoiner; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiConsumer; +import java.util.function.Supplier; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.notNullValue; @@ -106,13 +107,42 @@ public void setUp() { } } + public BulkRetryStrategy createObjectUnderTest( + final RequestFunction, BulkResponse> requestFunction, + final BiConsumer, Throwable> logFailure, + final Supplier bulkRequestSupplier +) { + return new BulkRetryStrategy( + requestFunction, + logFailure, + pluginMetrics, + Integer.MAX_VALUE, + bulkRequestSupplier, + pluginSetting); + } + + public BulkRetryStrategy createObjectUnderTest( + final RequestFunction, BulkResponse> requestFunction, + final BiConsumer, Throwable> logFailure, + final int maxRetries, + final Supplier bulkRequestSupplier +) { + return new BulkRetryStrategy( + requestFunction, + logFailure, + pluginMetrics, + maxRetries, + bulkRequestSupplier, + pluginSetting); + } + @Test public void testCanRetry() { AccumulatingBulkRequest accumulatingBulkRequest = mock(AccumulatingBulkRequest.class); - final BulkRetryStrategy bulkRetryStrategy = new BulkRetryStrategy( + final BulkRetryStrategy bulkRetryStrategy = createObjectUnderTest( bulkRequest -> mock(BulkResponse.class), - (docWriteRequest, throwable) -> {}, pluginMetrics, Integer.MAX_VALUE, - () -> mock(AccumulatingBulkRequest.class), pluginSetting); + (docWriteRequest, throwable) -> {}, + () -> mock(AccumulatingBulkRequest.class)); final String testIndex = "foo"; final BulkResponseItem bulkItemResponse1 = successItemResponse(testIndex); final BulkResponseItem bulkItemResponse2 = badRequestItemResponse(testIndex); @@ -139,9 +169,9 @@ public void testExecuteSuccessOnFirstAttempt() throws Exception { numEventsSucceeded = 0; numEventsFailed = 0; - final BulkRetryStrategy bulkRetryStrategy = new BulkRetryStrategy( - client::bulk, logFailureConsumer, pluginMetrics, Integer.MAX_VALUE, - () -> new JavaClientAccumulatingUncompressedBulkRequest(new BulkRequest.Builder()), pluginSetting); + final BulkRetryStrategy bulkRetryStrategy = createObjectUnderTest( + client::bulk, logFailureConsumer, + () -> new JavaClientAccumulatingUncompressedBulkRequest(new BulkRequest.Builder())); final IndexOperation indexOperation1 = new IndexOperation.Builder().index(testIndex).id("1").document(arbitraryDocument()).build(); final IndexOperation indexOperation2 = new IndexOperation.Builder().index(testIndex).id("2").document(arbitraryDocument()).build(); @@ -178,9 +208,9 @@ public void testExecuteRetryable() throws Exception { numEventsSucceeded = 0; numEventsFailed = 0; - final BulkRetryStrategy bulkRetryStrategy = new BulkRetryStrategy( - client::bulk, logFailureConsumer, pluginMetrics, Integer.MAX_VALUE, - () -> new JavaClientAccumulatingUncompressedBulkRequest(new BulkRequest.Builder()), pluginSetting); + final BulkRetryStrategy bulkRetryStrategy = createObjectUnderTest( + client::bulk, logFailureConsumer, + () -> new JavaClientAccumulatingUncompressedBulkRequest(new BulkRequest.Builder())); final IndexOperation indexOperation1 = new IndexOperation.Builder().index(testIndex).id("1").document(arbitraryDocument()).build(); final IndexOperation indexOperation2 = new IndexOperation.Builder().index(testIndex).id("2").document(arbitraryDocument()).build(); final IndexOperation indexOperation3 = new IndexOperation.Builder().index(testIndex).id("3").document(arbitraryDocument()).build(); @@ -241,9 +271,9 @@ public void testExecuteNonRetryableException() throws Exception { numEventsSucceeded = 0; numEventsFailed = 0; - final BulkRetryStrategy bulkRetryStrategy = new BulkRetryStrategy( - client::bulk, logFailureConsumer, pluginMetrics, Integer.MAX_VALUE, - () -> new JavaClientAccumulatingUncompressedBulkRequest(new BulkRequest.Builder()), pluginSetting); + final BulkRetryStrategy bulkRetryStrategy = createObjectUnderTest( + client::bulk, logFailureConsumer, + () -> new JavaClientAccumulatingUncompressedBulkRequest(new BulkRequest.Builder())); final IndexOperation indexOperation1 = new IndexOperation.Builder().index(testIndex).id("1").document(arbitraryDocument()).build(); final IndexOperation indexOperation2 = new IndexOperation.Builder().index(testIndex).id("2").document(arbitraryDocument()).build(); final IndexOperation indexOperation3 = new IndexOperation.Builder().index(testIndex).id("3").document(arbitraryDocument()).build(); @@ -306,9 +336,9 @@ public void testExecuteWithMaxRetries() throws Exception { maxRetriesLimitReached = false; client.maxRetriesTestValue = MAX_RETRIES; logFailureConsumer = this::logFailureMaxRetries; - final BulkRetryStrategy bulkRetryStrategy = new BulkRetryStrategy( - client::bulk, logFailureConsumer, pluginMetrics, MAX_RETRIES, - () -> new JavaClientAccumulatingUncompressedBulkRequest(new BulkRequest.Builder()), pluginSetting); + final BulkRetryStrategy bulkRetryStrategy = createObjectUnderTest( + client::bulk, logFailureConsumer, MAX_RETRIES, + () -> new JavaClientAccumulatingUncompressedBulkRequest(new BulkRequest.Builder())); final IndexOperation indexOperation1 = new IndexOperation.Builder().index(testIndex).id("1").document(arbitraryDocument()).build(); final IndexOperation indexOperation2 = new IndexOperation.Builder().index(testIndex).id("2").document(arbitraryDocument()).build(); final IndexOperation indexOperation3 = new IndexOperation.Builder().index(testIndex).id("3").document(arbitraryDocument()).build(); @@ -336,9 +366,9 @@ public void testExecuteWithMaxRetriesWithException() throws Exception { client.maxRetriesTestValue = MAX_RETRIES; client.maxRetriesWithException = true; logFailureConsumer = this::logFailureMaxRetries; - final BulkRetryStrategy bulkRetryStrategy = new BulkRetryStrategy( - client::bulk, logFailureConsumer, pluginMetrics, MAX_RETRIES, - () -> new JavaClientAccumulatingUncompressedBulkRequest(new BulkRequest.Builder()), pluginSetting); + final BulkRetryStrategy bulkRetryStrategy = createObjectUnderTest( + client::bulk, logFailureConsumer, MAX_RETRIES, + () -> new JavaClientAccumulatingUncompressedBulkRequest(new BulkRequest.Builder())); final IndexOperation indexOperation1 = new IndexOperation.Builder().index(testIndex).id("1").document(arbitraryDocument()).build(); final IndexOperation indexOperation2 = new IndexOperation.Builder().index(testIndex).id("2").document(arbitraryDocument()).build(); final IndexOperation indexOperation3 = new IndexOperation.Builder().index(testIndex).id("3").document(arbitraryDocument()).build(); @@ -366,9 +396,9 @@ public void testExecuteWithMaxRetriesAndSuccesses() throws Exception { client.maxRetriesTestValue = MAX_RETRIES; client.maxRetriesWithSuccesses = true; logFailureConsumer = this::logFailureMaxRetries; - final BulkRetryStrategy bulkRetryStrategy = new BulkRetryStrategy( - client::bulk, logFailureConsumer, pluginMetrics, MAX_RETRIES, - () -> new JavaClientAccumulatingUncompressedBulkRequest(new BulkRequest.Builder()), pluginSetting); + final BulkRetryStrategy bulkRetryStrategy = createObjectUnderTest( + client::bulk, logFailureConsumer, MAX_RETRIES, + () -> new JavaClientAccumulatingUncompressedBulkRequest(new BulkRequest.Builder())); final IndexOperation indexOperation1 = new IndexOperation.Builder().index(testIndex).id("1").document(arbitraryDocument()).build(); final IndexOperation indexOperation2 = new IndexOperation.Builder().index(testIndex).id("2").document(arbitraryDocument()).build(); final IndexOperation indexOperation3 = new IndexOperation.Builder().index(testIndex).id("3").document(arbitraryDocument()).build(); @@ -393,9 +423,9 @@ public void testExecuteNonRetryableResponse() throws Exception { numEventsSucceeded = 0; numEventsFailed = 0; - final BulkRetryStrategy bulkRetryStrategy = new BulkRetryStrategy( - client::bulk, logFailureConsumer, pluginMetrics, Integer.MAX_VALUE, - () -> new JavaClientAccumulatingUncompressedBulkRequest(new BulkRequest.Builder()), pluginSetting); + final BulkRetryStrategy bulkRetryStrategy = createObjectUnderTest( + client::bulk, logFailureConsumer, + () -> new JavaClientAccumulatingUncompressedBulkRequest(new BulkRequest.Builder())); final IndexOperation indexOperation1 = new IndexOperation.Builder().index(testIndex).id("1").document(arbitraryDocument()).build(); final IndexOperation indexOperation2 = new IndexOperation.Builder().index(testIndex).id("2").document(arbitraryDocument()).build(); final IndexOperation indexOperation3 = new IndexOperation.Builder().index(testIndex).id("3").document(arbitraryDocument()).build(); diff --git a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/bulk/JavaClientAccumulatingUncompressedBulkRequestTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/bulk/JavaClientAccumulatingUncompressedBulkRequestTest.java index a1ea5159f0..5d6996059b 100644 --- a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/bulk/JavaClientAccumulatingUncompressedBulkRequestTest.java +++ b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/bulk/JavaClientAccumulatingUncompressedBulkRequestTest.java @@ -219,4 +219,4 @@ private SizedDocument generateDocumentWithLength(long documentLength) { when(sizedDocument.getDocumentSize()).thenReturn(documentLength); return sizedDocument; } -} \ No newline at end of file +} From 15d8247205765fb39c41354a2b7214d98ac3ba83 Mon Sep 17 00:00:00 2001 From: kkondaka <41027584+kkondaka@users.noreply.github.com> Date: Tue, 7 Nov 2023 12:18:43 -0800 Subject: [PATCH 30/53] Javadoc fixes (#3594) Signed-off-by: Krishna Kondaka Co-authored-by: Krishna Kondaka --- .../dataprepper/model/configuration/PluginSetting.java | 1 + .../opensearch/dataprepper/model/event/EventHandle.java | 5 +++++ .../dataprepper/model/plugin/PluginConfigObservable.java | 3 +++ .../dataprepper/model/plugin/PluginConfigPublisher.java | 3 +++ .../dataprepper/model/plugin/PluginFactory.java | 1 + .../coordinator/enhanced/EnhancedSourcePartition.java | 8 ++++++++ 6 files changed, 21 insertions(+) diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/configuration/PluginSetting.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/configuration/PluginSetting.java index 6924cb35bc..a8ea4a3ee1 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/configuration/PluginSetting.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/configuration/PluginSetting.java @@ -161,6 +161,7 @@ public List getTypedList(final String attribute, final Class type) { * Returns the value of the specified {@literal List>}, or {@code defaultValue} if this settings contains no value for * the attribute. * + * @param attribute attribute to be looked up * @param keyType key type of the Map * @param valueType value type stored in the Map * @param The key type diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java index 6049a196e3..d05dd8e36c 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventHandle.java @@ -42,6 +42,11 @@ public interface EventHandle { */ Instant getInternalOriginationTime(); + /** + * registers onRelease consumer with event handle + * + * @param releaseConsumer consumer to be calledback when event handle is released. + */ void onRelease(BiConsumer releaseConsumer); } diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginConfigObservable.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginConfigObservable.java index 6a62e2647a..3c436fc68a 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginConfigObservable.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginConfigObservable.java @@ -10,6 +10,9 @@ public interface PluginConfigObservable { /** * Onboard a new {@link PluginConfigObserver} within the plugin. + * + * @param pluginConfigObserver plugin config observer + * @return returns true if the opration is successful, false otherwise */ boolean addPluginConfigObserver(PluginConfigObserver pluginConfigObserver); diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginConfigPublisher.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginConfigPublisher.java index 66982a8db6..ae3ff0ccb4 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginConfigPublisher.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginConfigPublisher.java @@ -7,6 +7,9 @@ public interface PluginConfigPublisher { /** * Onboard a new {@link PluginConfigObservable}. + * + * @param pluginConfigObservable observable plugin configuration + * @return true if the operation is successful, false otherwise */ boolean addPluginConfigObservable(PluginConfigObservable pluginConfigObservable); diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginFactory.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginFactory.java index b233d89043..155df776f4 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginFactory.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/PluginFactory.java @@ -22,6 +22,7 @@ public interface PluginFactory { * * @param baseClass The class type that the plugin is supporting. * @param pluginSetting The {@link PluginSetting} to configure this plugin + * @param args variable number of arguments * @param The type * @return A new instance of your plugin, configured * @since 1.2 diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/enhanced/EnhancedSourcePartition.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/enhanced/EnhancedSourcePartition.java index ac52ace478..0fa683b2a5 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/enhanced/EnhancedSourcePartition.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/source/coordinator/enhanced/EnhancedSourcePartition.java @@ -48,6 +48,11 @@ public void setSourcePartitionStoreItem(SourcePartitionStoreItem sourcePartition /** * Helper method to convert progress state. * This is because the state is currently stored as a String in the coordination store. + * + * @param progressStateClass class of progress state + * @param serializedPartitionProgressState serialized value of the partition progress state + * + * @return returns the converted value of the progress state */ public T convertStringToPartitionProgressState(Class progressStateClass, final String serializedPartitionProgressState) { if (Objects.isNull(serializedPartitionProgressState)) { @@ -69,6 +74,9 @@ public T convertStringToPartitionProgressState(Class progressStateClass, fina /** * Helper method to convert progress state to String * This is because the state is currently stored as a String in the coordination store. + * + * @param partitionProgressState optional parameter indicating the partition progress state + * @return returns the progress state as string */ public String convertPartitionProgressStatetoString(Optional partitionProgressState) { if (partitionProgressState.isEmpty()) { From 97cd930be30b74730dffc122b3f45092d6a8c6c3 Mon Sep 17 00:00:00 2001 From: Aiden Dai <68811299+daixba@users.noreply.github.com> Date: Wed, 8 Nov 2023 06:26:24 +0800 Subject: [PATCH 31/53] Add CMK encryption support to DynamoDB export (#3592) Signed-off-by: Aiden Dai --- .../dynamodb-source/README.md | 2 ++ .../source/dynamodb/DynamoDBService.java | 11 ++++++-- .../dynamodb/configuration/ExportConfig.java | 17 ++++++++++- .../state/ExportProgressState.java | 13 ++++++++- .../dynamodb/export/DataFileLoader.java | 15 ++++------ .../dynamodb/export/ExportScheduler.java | 6 ++-- .../dynamodb/export/ExportTaskManager.java | 9 ++++-- .../dynamodb/export/ManifestFileReader.java | 28 ++++++++----------- .../source/dynamodb/model/TableMetadata.java | 15 +++++++++- .../dynamodb/export/DataFileLoaderTest.java | 6 ++-- 10 files changed, 81 insertions(+), 41 deletions(-) diff --git a/data-prepper-plugins/dynamodb-source/README.md b/data-prepper-plugins/dynamodb-source/README.md index 1585e66e12..db620a6406 100644 --- a/data-prepper-plugins/dynamodb-source/README.md +++ b/data-prepper-plugins/dynamodb-source/README.md @@ -52,6 +52,8 @@ source: * 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 diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java index 1916a5344f..fed9e83f79 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java @@ -171,7 +171,12 @@ public void init() { Instant startTime = Instant.now(); if (tableInfo.getMetadata().isExportRequired()) { - createExportPartition(tableInfo.getTableArn(), startTime, tableInfo.getMetadata().getExportBucket(), tableInfo.getMetadata().getExportPrefix()); + createExportPartition( + tableInfo.getTableArn(), + startTime, + tableInfo.getMetadata().getExportBucket(), + tableInfo.getMetadata().getExportPrefix(), + tableInfo.getMetadata().getExportKmsKeyId()); } if (tableInfo.getMetadata().isStreamRequired()) { @@ -209,11 +214,12 @@ public void init() { * @param bucket Export bucket * @param prefix Export Prefix */ - private void createExportPartition(String tableArn, Instant exportTime, String bucket, String prefix) { + private void createExportPartition(String tableArn, Instant exportTime, String bucket, String prefix, String kmsKeyId) { ExportProgressState exportProgressState = new ExportProgressState(); exportProgressState.setBucket(bucket); exportProgressState.setPrefix(prefix); exportProgressState.setExportTime(exportTime.toString()); // information purpose + exportProgressState.setKmsKeyId(kmsKeyId); ExportPartition exportPartition = new ExportPartition(tableArn, exportTime, Optional.of(exportProgressState)); coordinator.createPartition(exportPartition); } @@ -310,6 +316,7 @@ private TableInfo getTableInfo(TableConfig tableConfig) { .streamStartPosition(streamStartPosition) .exportBucket(tableConfig.getExportConfig() == null ? null : tableConfig.getExportConfig().getS3Bucket()) .exportPrefix(tableConfig.getExportConfig() == null ? null : tableConfig.getExportConfig().getS3Prefix()) + .exportKmsKeyId(tableConfig.getExportConfig() == null ? null : tableConfig.getExportConfig().getS3SseKmsKeyId()) .build(); return new TableInfo(tableConfig.getTableArn(), metadata); } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/configuration/ExportConfig.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/configuration/ExportConfig.java index d34017599a..dd1aac12a7 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/configuration/ExportConfig.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/configuration/ExportConfig.java @@ -6,7 +6,9 @@ package org.opensearch.dataprepper.plugins.source.dynamodb.configuration; import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.constraints.AssertTrue; import jakarta.validation.constraints.NotBlank; +import software.amazon.awssdk.arns.Arn; import software.amazon.awssdk.regions.Region; public class ExportConfig { @@ -14,13 +16,15 @@ public class ExportConfig { @JsonProperty("s3_bucket") @NotBlank(message = "Bucket Name is required for export") private String s3Bucket; - @JsonProperty("s3_prefix") private String s3Prefix; @JsonProperty("s3_region") private String s3Region; + @JsonProperty("s3_sse_kms_key_id") + private String s3SseKmsKeyId; + public String getS3Bucket() { return s3Bucket; } @@ -33,4 +37,15 @@ public Region getAwsRegion() { return s3Region != null ? Region.of(s3Region) : null; } + public String getS3SseKmsKeyId() { + return s3SseKmsKeyId; + } + + @AssertTrue(message = "KMS Key ID must be a valid one.") + boolean isKmsKeyIdValid() { + // If key id is provided, it should be in a format like + // arn:aws:kms:us-west-2:123456789012:key/0a4bc22f-bb96-4ad3-80ca-63b12b3ec147 + return s3SseKmsKeyId == null || Arn.fromString(s3SseKmsKeyId).resourceAsString() != null; + } + } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/coordination/state/ExportProgressState.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/coordination/state/ExportProgressState.java index 4aa15cacf5..e7dadde99c 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/coordination/state/ExportProgressState.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/coordination/state/ExportProgressState.java @@ -21,10 +21,13 @@ public class ExportProgressState { @JsonProperty("prefix") private String prefix; + @JsonProperty("kmsKeyId") + private String kmsKeyId; + @JsonProperty("exportTime") private String exportTime; - + public String getExportArn() { return exportArn; } @@ -64,4 +67,12 @@ public String getStatus() { public void setStatus(String status) { this.status = status; } + + public String getKmsKeyId() { + return kmsKeyId; + } + + public void setKmsKeyId(String kmsKeyId) { + this.kmsKeyId = kmsKeyId; + } } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoader.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoader.java index 56a2065055..6093d5d688 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoader.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoader.java @@ -7,10 +7,10 @@ import org.opensearch.dataprepper.buffer.common.BufferAccumulator; import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; 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.acknowledgements.AcknowledgementSet; import org.opensearch.dataprepper.plugins.source.dynamodb.converter.ExportRecordConverter; import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo; import org.slf4j.Logger; @@ -169,14 +169,12 @@ public void run() { int lineCount = 0; int lastLineProcessed = 0; - try { - InputStream inputStream = objectReader.readFile(bucketName, key); - GZIPInputStream gzipInputStream = new GZIPInputStream(inputStream); - BufferedReader reader = new BufferedReader(new InputStreamReader(gzipInputStream)); + try (InputStream inputStream = objectReader.readFile(bucketName, key); + GZIPInputStream gzipInputStream = new GZIPInputStream(inputStream); + BufferedReader reader = new BufferedReader(new InputStreamReader(gzipInputStream))) { String line; while ((line = reader.readLine()) != null) { - if (shouldStop) { checkpointer.checkpoint(lastLineProcessed); LOG.debug("Should Stop flag is set to True, looks like shutdown has triggered"); @@ -213,9 +211,6 @@ public void run() { } lines.clear(); - reader.close(); - gzipInputStream.close(); - inputStream.close(); LOG.info("Completed loading s3://{}/{} to buffer", bucketName, key); @@ -226,7 +221,7 @@ public void run() { } catch (Exception e) { checkpointer.checkpoint(lineCount); - String errorMessage = String.format("Loading of s3://%s/%s completed with Exception: %S", bucketName, key, e.getMessage()); + String errorMessage = String.format("Loading of s3://%s/%s completed with Exception: %s", bucketName, key, e.getMessage()); throw new RuntimeException(errorMessage); } } diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java index 9ace3aa6ae..c150643d56 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java @@ -156,8 +156,7 @@ private BiConsumer completeExport(ExportPartition exportParti ExportProgressState state = exportPartition.getProgressState().get(); String bucketName = state.getBucket(); String exportArn = state.getExportArn(); - - + String manifestKey = exportTaskManager.getExportManifest(exportArn); LOG.debug("Export manifest summary file is " + manifestKey); @@ -189,6 +188,7 @@ private void createDataFilePartitions(String exportArn, String bucketName, Map parseDataFile(String bucket, String key) { LOG.info("Try to read the manifest data file"); Map result = new HashMap<>(); - InputStream object = objectReader.readFile(bucket, key); - BufferedReader reader = new BufferedReader(new InputStreamReader(object)); - - String line; - try { + + try (InputStream object = objectReader.readFile(bucket, key); + BufferedReader reader = new BufferedReader(new InputStreamReader(object))) { + String line; while ((line = reader.readLine()) != null) { // An example line as below: // {"itemCount":46331,"md5Checksum":"a0k21IY3eelgr2PuWJLjJw==","etag":"51f9f394903c5d682321c6211aae8b6a-1","dataFileS3Key":"test-table-export/AWSDynamoDB/01692350182719-6de2c037/data/fpgzwz7ome3s7a5gqn2mu3ogtq.json.gz"} @@ -76,8 +69,9 @@ public Map parseDataFile(String bucket, String key) { result.put(map.get(DATA_FILE_S3_KEY), Integer.valueOf(map.get(DATA_FILE_ITEM_COUNT_KEY))); } - } catch (IOException e) { - LOG.error("IO Exception due to {}", e.getMessage()); + } catch (Exception e) { + LOG.error("Exception due to {}", e.getMessage()); + throw new RuntimeException(e); } return result; diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/TableMetadata.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/TableMetadata.java index 8415a41001..ba21304d31 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/TableMetadata.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/TableMetadata.java @@ -19,7 +19,6 @@ public class TableMetadata { private static final String REQUIRE_EXPORT_KEY = "export"; private static final String REQUIRE_STREAM_KEY = "stream"; - private final String partitionKeyAttributeName; private final String sortKeyAttributeName; @@ -36,6 +35,8 @@ public class TableMetadata { private final String exportPrefix; + private final String exportKmsKeyId; + private TableMetadata(Builder builder) { this.partitionKeyAttributeName = builder.partitionKeyAttributeName; this.sortKeyAttributeName = builder.sortKeyAttributeName; @@ -45,6 +46,7 @@ private TableMetadata(Builder builder) { this.exportBucket = builder.exportBucket; this.exportPrefix = builder.exportPrefix; this.streamStartPosition = builder.streamStartPosition; + this.exportKmsKeyId = builder.exportKmsKeyId; } @@ -70,6 +72,8 @@ public static class Builder { private String exportPrefix; + private String exportKmsKeyId; + private StreamStartPosition streamStartPosition; @@ -108,6 +112,11 @@ public Builder exportPrefix(String exportPrefix) { return this; } + public Builder exportKmsKeyId(String exportKmsKeyId) { + this.exportKmsKeyId = exportKmsKeyId; + return this; + } + public Builder streamStartPosition(StreamStartPosition streamStartPosition) { this.streamStartPosition = streamStartPosition; return this; @@ -173,4 +182,8 @@ public String getExportBucket() { public String getExportPrefix() { return exportPrefix; } + + public String getExportKmsKeyId() { + return exportKmsKeyId; + } } diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java index a92b80645b..162cef81b3 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java @@ -93,7 +93,7 @@ class DataFileLoaderTest { private final Random random = new Random(); - private final int total = random.nextInt(10); + private final int total = random.nextInt(10) + 1; @BeforeEach void setup() { @@ -153,8 +153,8 @@ void test_run_loadFile_correctly() { try ( final MockedStatic bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class); final MockedConstruction recordConverterMockedConstruction = mockConstruction(ExportRecordConverter.class, (mock, context) -> { - exportRecordConverter = mock; - })) { + exportRecordConverter = mock; + })) { bufferAccumulatorMockedStatic.when(() -> BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT)).thenReturn(bufferAccumulator); loader = DataFileLoader.builder(objectReader, pluginMetrics, buffer) .bucketName(bucketName) From 206752311174c8211c1c883cf9042bd5e581455b Mon Sep 17 00:00:00 2001 From: Aiden Dai <68811299+daixba@users.noreply.github.com> Date: Wed, 8 Nov 2023 07:55:18 +0800 Subject: [PATCH 32/53] Optimize idle item between GetRecords call (#3593) Signed-off-by: Aiden Dai --- .../source/dynamodb/stream/ShardConsumer.java | 24 ++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java index 48e093f7b6..c5da709775 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumer.java @@ -7,10 +7,10 @@ import org.opensearch.dataprepper.buffer.common.BufferAccumulator; import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet; 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.acknowledgements.AcknowledgementSet; import org.opensearch.dataprepper.plugins.source.dynamodb.converter.StreamRecordConverter; import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo; import org.slf4j.Logger; @@ -47,6 +47,16 @@ public class ShardConsumer implements Runnable { */ private static final int GET_RECORD_INTERVAL_MILLS = 300; + /** + * Idle Time between GetRecords Reads + */ + private static final int MINIMUM_GET_RECORD_INTERVAL_MILLS = 10; + + /** + * Minimum Idle Time between GetRecords Reads + */ + private static final long GET_RECORD_DELAY_THRESHOLD_MILLS = 15_000; + /** * Default interval to check if export is completed. */ @@ -209,12 +219,15 @@ public void run() { shardIterator = response.nextShardIterator(); + int interval; + if (!response.records().isEmpty()) { // Always use the last sequence number for checkpoint sequenceNumber = response.records().get(response.records().size() - 1).dynamodb().sequenceNumber(); + Instant lastEventTime = response.records().get(response.records().size() - 1).dynamodb().approximateCreationDateTime(); if (waitForExport) { - Instant lastEventTime = response.records().get(response.records().size() - 1).dynamodb().approximateCreationDateTime(); + if (lastEventTime.compareTo(startTime) <= 0) { LOG.debug("Get {} events before start time, ignore...", response.records().size()); continue; @@ -230,10 +243,15 @@ public void run() { records = response.records(); } recordConverter.writeToBuffer(acknowledgementSet, records); + long delay = System.currentTimeMillis() - lastEventTime.toEpochMilli(); + interval = delay > GET_RECORD_DELAY_THRESHOLD_MILLS ? MINIMUM_GET_RECORD_INTERVAL_MILLS : GET_RECORD_INTERVAL_MILLS; + + } else { + interval = GET_RECORD_INTERVAL_MILLS; } try { // Idle between get records call. - Thread.sleep(GET_RECORD_INTERVAL_MILLS); + Thread.sleep(interval); } catch (InterruptedException e) { throw new RuntimeException(e); } From edb13ef309aefcfc84fcdee4d9823c14efe0b7fa Mon Sep 17 00:00:00 2001 From: Taylor Gray Date: Tue, 7 Nov 2023 18:38:46 -0600 Subject: [PATCH 33/53] =?UTF-8?q?Add=20document=5Fversion=20and=20document?= =?UTF-8?q?=5Fversion=5Ftype=20parameters=20to=20the=20open=E2=80=A6=20(#3?= =?UTF-8?q?591)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add document_version and document_version_type parameters to the opensearch sink for conditional indexing of documents Signed-off-by: Taylor Gray --- .../expression/ExpressionEvaluator.java | 4 +- .../dataprepper/model/event/JacksonEvent.java | 38 +++------- .../expression/ExpressionEvaluatorTest.java | 5 ++ .../model/event/JacksonEventTest.java | 24 ++---- .../GenericExpressionEvaluator.java | 22 ++++++ .../GenericExpressionEvaluatorTest.java | 27 ++++++- .../sink/opensearch/OpenSearchSinkIT.java | 8 +- .../sink/opensearch/BulkRetryStrategy.java | 29 +++++-- .../sink/opensearch/OpenSearchSink.java | 76 +++++++++++++++---- .../opensearch/index/IndexConfiguration.java | 65 +++++++++++++++- .../opensearch/BulkRetryStrategyTests.java | 67 ++++++++++------ .../OpenSearchSinkConfigurationTests.java | 5 +- .../index/IndexConfigurationTests.java | 38 ++++++++++ 13 files changed, 311 insertions(+), 97 deletions(-) diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/expression/ExpressionEvaluator.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/expression/ExpressionEvaluator.java index c006a31cae..ea786ce5bd 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/expression/ExpressionEvaluator.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/expression/ExpressionEvaluator.java @@ -34,4 +34,6 @@ default Boolean evaluateConditional(final String statement, final Event context) } Boolean isValidExpressionStatement(final String statement); -} + + Boolean isValidFormatExpressions(final String format); +} \ No newline at end of file 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 1f36e01e25..341d5277bc 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 @@ -318,29 +318,6 @@ public String formatString(final String format, final ExpressionEvaluator expres return formatStringInternal(format, expressionEvaluator); } - public static boolean isValidFormatExpressions(final String format, final ExpressionEvaluator expressionEvaluator) { - if (Objects.isNull(expressionEvaluator)) { - return false; - } - - int fromIndex = 0; - int position = 0; - while ((position = format.indexOf("${", fromIndex)) != -1) { - int endPosition = format.indexOf("}", position + 1); - if (endPosition == -1) { - return false; - } - String name = format.substring(position + 2, endPosition); - - Object val; - if (!expressionEvaluator.isValidExpressionStatement(name)) { - return false; - } - fromIndex = endPosition + 1; - } - return true; - } - private String formatStringInternal(final String format, final ExpressionEvaluator expressionEvaluator) { int fromIndex = 0; String result = ""; @@ -362,7 +339,7 @@ private String formatStringInternal(final String format, final ExpressionEvaluat } if (val == null) { - if (Objects.nonNull(expressionEvaluator) && expressionEvaluator.isValidExpressionStatement(name)) { + if (expressionEvaluator != null && expressionEvaluator.isValidExpressionStatement(name)) { val = expressionEvaluator.evaluate(name, this); } else { throw new EventKeyNotFoundException(String.format("The key %s could not be found in the Event when formatting", name)); @@ -410,12 +387,21 @@ public Map toMap() { return mapper.convertValue(jsonNode, MAP_TYPE_REFERENCE); } + + public static boolean isValidEventKey(final String key) { + try { + checkKey(key); + return true; + } catch (final Exception e) { + return false; + } + } private String checkAndTrimKey(final String key) { checkKey(key); return trimKey(key); } - private void checkKey(final String 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.length() > MAX_KEY_LENGTH) { @@ -432,7 +418,7 @@ private String trimKey(final String key) { return trimmedLeadingSlash.endsWith(SEPARATOR) ? trimmedLeadingSlash.substring(0, trimmedLeadingSlash.length() - 2) : trimmedLeadingSlash; } - private boolean isValidKey(final String key) { + private static boolean isValidKey(final String key) { for (int i = 0; i < key.length(); i++) { char c = key.charAt(i); diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/expression/ExpressionEvaluatorTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/expression/ExpressionEvaluatorTest.java index 36a60ac447..6ce1af8660 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/expression/ExpressionEvaluatorTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/expression/ExpressionEvaluatorTest.java @@ -23,6 +23,11 @@ public Object evaluate(final String statement, final Event event) { public Boolean isValidExpressionStatement(final String statement) { return true; } + + @Override + public Boolean isValidFormatExpressions(String format) { + return true; + } } @Test 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 deb68f83bd..4fe8f272cc 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 @@ -535,25 +535,6 @@ public void testBuild_withFormatString(String formattedString, String finalStrin assertThat(event.formatString(formattedString), is(equalTo(finalString))); } - @ParameterizedTest - @CsvSource({ - "abc-${/foo, false", - "abc-${/foo}, true", - "abc-${getMetadata(\"key\")}, true", - "abc-${getXYZ(\"key\")}, false" - }) - public void testBuild_withIsValidFormatExpressions(final String format, final Boolean expectedResult) { - final ExpressionEvaluator expressionEvaluator = mock(ExpressionEvaluator.class); - when(expressionEvaluator.isValidExpressionStatement("/foo")).thenReturn(true); - when(expressionEvaluator.isValidExpressionStatement("getMetadata(\"key\")")).thenReturn(true); - assertThat(JacksonEvent.isValidFormatExpressions(format, expressionEvaluator), equalTo(expectedResult)); - } - - @Test - public void testBuild_withIsValidFormatExpressionsWithNullEvaluator() { - assertThat(JacksonEvent.isValidFormatExpressions("${}", null), equalTo(false)); - } - @Test public void formatString_with_expression_evaluator_catches_exception_when_Event_get_throws_exception() { @@ -818,6 +799,11 @@ void testJsonStringBuilderWithExcludeKeys() { } + @ParameterizedTest + @CsvSource(value = {"test_key, true", "/test_key, true", "inv(alid, false", "getMetadata(\"test_key\"), false"}) + void isValidEventKey_returns_expected_result(final String key, final boolean isValid) { + assertThat(JacksonEvent.isValidEventKey(key), equalTo(isValid)); + } private static Map createComplexDataMap() { final Map dataObject = new HashMap<>(); diff --git a/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluator.java b/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluator.java index 072b3a7393..89d8b15a84 100644 --- a/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluator.java +++ b/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluator.java @@ -7,6 +7,7 @@ import org.antlr.v4.runtime.tree.ParseTree; import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.event.JacksonEvent; import javax.inject.Inject; import javax.inject.Named; @@ -52,4 +53,25 @@ public Boolean isValidExpressionStatement(final String statement) { return false; } } + + @Override + public Boolean isValidFormatExpressions(final String format) { + int fromIndex = 0; + int position = 0; + while ((position = format.indexOf("${", fromIndex)) != -1) { + int endPosition = format.indexOf("}", position + 1); + if (endPosition == -1) { + return false; + } + String name = format.substring(position + 2, endPosition); + + Object val; + // Invalid if it is not a valid key and not a valid expression statement + if (!JacksonEvent.isValidEventKey(name) && !isValidExpressionStatement(name)) { + return false; + } + fromIndex = endPosition + 1; + } + return true; + } } diff --git a/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluatorTest.java b/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluatorTest.java index d68808dc85..0b8ac2ee03 100644 --- a/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluatorTest.java +++ b/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluatorTest.java @@ -5,22 +5,26 @@ package org.opensearch.dataprepper.expression; -import org.opensearch.dataprepper.model.event.Event; import org.antlr.v4.runtime.tree.ParseTree; 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.junit.jupiter.params.provider.ValueSource; import org.mockito.InjectMocks; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.model.event.Event; -import java.util.UUID; import java.util.Random; +import java.util.UUID; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; 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.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.doThrow; @@ -119,5 +123,24 @@ void isValidExpressionStatement_returns_false_when_parse_throws() { assertThat(result, equalTo(false)); } + @ParameterizedTest + @CsvSource({ + "abc-${/foo, false", + "abc-${/foo}, true", + "abc-${getMetadata(\"key\")}, true", + "abc-${getXYZ(\"key\")}, true", + "abc-${invalid, false" + }) + void isValidFormatExpressionsReturnsCorrectResult(final String format, final Boolean expectedResult) { + assertThat(statementEvaluator.isValidFormatExpressions(format), equalTo(expectedResult)); + } + + @ParameterizedTest + @ValueSource(strings = {"abc-${anyS(=tring}"}) + void isValidFormatExpressionsReturnsFalseWhenIsValidKeyAndValidExpressionIsFalse(final String format) { + doThrow(RuntimeException.class).when(parser).parse(anyString()); + assertThat(statementEvaluator.isValidFormatExpressions(format), equalTo(false)); + } + } 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 755ccb5283..fbe17763a0 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 @@ -646,9 +646,11 @@ public void testOpenSearchBulkActionsCreateWithExpression() throws IOException, pluginSetting.getSettings().put(IndexConfiguration.DOCUMENT_ID_FIELD, testIdField); Event event = (Event)testRecords.get(0).getData(); event.getMetadata().setAttribute("action", "create"); + final String actionFormatExpression = "${getMetadata(\"action\")}"; + when(expressionEvaluator.isValidFormatExpressions(actionFormatExpression)).thenReturn(true); when(expressionEvaluator.isValidExpressionStatement("getMetadata(\"action\")")).thenReturn(true); when(expressionEvaluator.evaluate("getMetadata(\"action\")", event)).thenReturn(event.getMetadata().getAttribute("action")); - pluginSetting.getSettings().put(IndexConfiguration.ACTION, "${getMetadata(\"action\")}"); + pluginSetting.getSettings().put(IndexConfiguration.ACTION, actionFormatExpression); final OpenSearchSink sink = createObjectUnderTest(pluginSetting, true); sink.output(testRecords); final List> retSources = getSearchResponseDocSources(testIndexAlias); @@ -677,9 +679,11 @@ public void testOpenSearchBulkActionsCreateWithInvalidExpression() throws IOExce pluginSetting.getSettings().put(IndexConfiguration.DOCUMENT_ID_FIELD, testIdField); Event event = (Event)testRecords.get(0).getData(); event.getMetadata().setAttribute("action", "unknown"); + final String actionFormatExpression = "${getMetadata(\"action\")}"; + when(expressionEvaluator.isValidFormatExpressions(actionFormatExpression)).thenReturn(true); when(expressionEvaluator.isValidExpressionStatement("getMetadata(\"action\")")).thenReturn(true); when(expressionEvaluator.evaluate("getMetadata(\"action\")", event)).thenReturn(event.getMetadata().getAttribute("action")); - pluginSetting.getSettings().put(IndexConfiguration.ACTION, "${getMetadata(\"action\")}"); + pluginSetting.getSettings().put(IndexConfiguration.ACTION, actionFormatExpression); final OpenSearchSink sink = createObjectUnderTest(pluginSetting, true); sink.output(testRecords); final List> retSources = getSearchResponseDocSources(testIndexAlias); diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java index 793029f6ef..5f6271010b 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java @@ -43,8 +43,10 @@ public final class BulkRetryStrategy { public static final String BULK_REQUEST_NOT_FOUND_ERRORS = "bulkRequestNotFoundErrors"; public static final String BULK_REQUEST_TIMEOUT_ERRORS = "bulkRequestTimeoutErrors"; public static final String BULK_REQUEST_SERVER_ERRORS = "bulkRequestServerErrors"; + public static final String DOCUMENTS_VERSION_CONFLICT_ERRORS = "documentsVersionConflictErrors"; static final long INITIAL_DELAY_MS = 50; static final long MAXIMUM_DELAY_MS = Duration.ofMinutes(10).toMillis(); + static final String VERSION_CONFLICT_EXCEPTION_TYPE = "version_conflict_engine_exception"; private static final Set NON_RETRY_STATUS = new HashSet<>( Arrays.asList( @@ -116,6 +118,7 @@ public final class BulkRetryStrategy { private final Counter bulkRequestNotFoundErrors; private final Counter bulkRequestTimeoutErrors; private final Counter bulkRequestServerErrors; + private final Counter documentsVersionConflictErrors; private static final Logger LOG = LoggerFactory.getLogger(BulkRetryStrategy.class); static class BulkOperationRequestResponse { @@ -160,6 +163,7 @@ public BulkRetryStrategy(final RequestFunction bulkRequest, final BulkResponse bulkResponse, final Throwable failure) { - LOG.warn("Bulk Operation Failed.", failure); if (failure == null) { for (final BulkResponseItem bulkItemResponse : bulkResponse.items()) { - if (Objects.nonNull(bulkItemResponse.error())) { + // Skip logging the error for version conflicts + if (bulkItemResponse.error() != null && !VERSION_CONFLICT_EXCEPTION_TYPE.equals(bulkItemResponse.error().type())) { LOG.warn("operation = {}, error = {}", bulkItemResponse.operationType(), bulkItemResponse.error().reason()); } } handleFailures(bulkRequest, bulkResponse.items()); } else { + LOG.warn("Bulk Operation Failed.", failure); handleFailures(bulkRequest, failure); } bulkRequestFailedCounter.increment(); @@ -313,6 +318,10 @@ private AccumulatingBulkRequest createBulkReq if (bulkItemResponse.error() != null) { if (!NON_RETRY_STATUS.contains(bulkItemResponse.status())) { requestToReissue.addOperation(bulkOperation); + } else if (VERSION_CONFLICT_EXCEPTION_TYPE.equals(bulkItemResponse.error().type())) { + documentsVersionConflictErrors.increment(); + LOG.debug("Received version conflict from OpenSearch: {}", bulkItemResponse.error().reason()); + bulkOperation.releaseEventHandle(true); } else { nonRetryableFailures.add(FailedBulkOperation.builder() .withBulkOperation(bulkOperation) @@ -338,10 +347,16 @@ private void handleFailures(final AccumulatingBulkRequest> { public static final String INVALID_ACTION_ERRORS = "invalidActionErrors"; public static final String BULKREQUEST_SIZE_BYTES = "bulkRequestSizeBytes"; public static final String DYNAMIC_INDEX_DROPPED_EVENTS = "dynamicIndexDroppedEvents"; + public static final String INVALID_VERSION_EXPRESSION_DROPPED_EVENTS = "dynamicDocumentVersionDroppedEvents"; private static final Logger LOG = LoggerFactory.getLogger(OpenSearchSink.class); private static final int INITIALIZE_RETRY_WAIT_TIME_MS = 5000; @@ -112,12 +114,15 @@ public class OpenSearchSink extends AbstractSink> { private final String documentRootKey; private String configuredIndexAlias; private final ReentrantLock lock; + private final VersionType versionType; + private final String versionExpression; private final Timer bulkRequestTimer; private final Counter bulkRequestErrorsCounter; private final Counter invalidActionErrorsCounter; private final Counter dynamicIndexDroppedEvents; private final DistributionSummary bulkRequestSizeBytesSummary; + private final Counter dynamicDocumentVersionDroppedEvents; private OpenSearchClient openSearchClient; private ObjectMapper objectMapper; private volatile boolean initialized; @@ -146,6 +151,7 @@ public OpenSearchSink(final PluginSetting pluginSetting, invalidActionErrorsCounter = pluginMetrics.counter(INVALID_ACTION_ERRORS); dynamicIndexDroppedEvents = pluginMetrics.counter(DYNAMIC_INDEX_DROPPED_EVENTS); bulkRequestSizeBytesSummary = pluginMetrics.summary(BULKREQUEST_SIZE_BYTES); + dynamicDocumentVersionDroppedEvents = pluginMetrics.counter(INVALID_VERSION_EXPRESSION_DROPPED_EVENTS); this.openSearchSinkConfig = OpenSearchSinkConfiguration.readESConfig(pluginSetting, expressionEvaluator); this.bulkSize = ByteSizeUnit.MB.toBytes(openSearchSinkConfig.getIndexConfiguration().getBulkSize()); @@ -157,6 +163,8 @@ public OpenSearchSink(final PluginSetting pluginSetting, this.action = openSearchSinkConfig.getIndexConfiguration().getAction(); this.actions = openSearchSinkConfig.getIndexConfiguration().getActions(); this.documentRootKey = openSearchSinkConfig.getIndexConfiguration().getDocumentRootKey(); + this.versionType = openSearchSinkConfig.getIndexConfiguration().getVersionType(); + this.versionExpression = openSearchSinkConfig.getIndexConfiguration().getVersionExpression(); this.indexManagerFactory = new IndexManagerFactory(new ClusterSettingsParser()); this.failedBulkOperationConverter = new FailedBulkOperationConverter(pluginSetting.getPipelineName(), pluginSetting.getName(), pluginSetting.getName()); @@ -260,7 +268,11 @@ public boolean isReady() { return initialized; } - private BulkOperation getBulkOperationForAction(final String action, final SerializedJson document, final String indexName, final JsonNode jsonNode) { + private BulkOperation getBulkOperationForAction(final String action, + final SerializedJson document, + final Long version, + final String indexName, + final JsonNode jsonNode) { BulkOperation bulkOperation; final Optional docId = document.getDocumentId(); final Optional routing = document.getRoutingField(); @@ -283,10 +295,14 @@ private BulkOperation getBulkOperationForAction(final String action, final Seria new UpdateOperation.Builder<>() .index(indexName) .document(jsonNode) - .upsert(jsonNode) : + .upsert(jsonNode) + .versionType(versionType) + .version(version) : new UpdateOperation.Builder<>() .index(indexName) - .document(jsonNode); + .document(jsonNode) + .versionType(versionType) + .version(version); docId.ifPresent(updateOperationBuilder::id); routing.ifPresent(updateOperationBuilder::routing); bulkOperation = new BulkOperation.Builder() @@ -300,13 +316,20 @@ private BulkOperation getBulkOperationForAction(final String action, final Seria docId.ifPresent(deleteOperationBuilder::id); routing.ifPresent(deleteOperationBuilder::routing); bulkOperation = new BulkOperation.Builder() - .delete(deleteOperationBuilder.build()) + .delete(deleteOperationBuilder + .versionType(versionType) + .version(version) + .build()) .build(); return bulkOperation; } // Default to "index" final IndexOperation.Builder indexOperationBuilder = - new IndexOperation.Builder<>().index(indexName).document(document); + new IndexOperation.Builder<>() + .index(indexName) + .document(document) + .version(version) + .versionType(versionType); docId.ifPresent(indexOperationBuilder::id); routing.ifPresent(indexOperationBuilder::routing); bulkOperation = new BulkOperation.Builder() @@ -337,16 +360,27 @@ public void doOutput(final Collection> records) { } catch (IOException | EventKeyNotFoundException e) { LOG.error("There was an exception when constructing the index name. Check the dlq if configured to see details about the affected Event: {}", e.getMessage()); dynamicIndexDroppedEvents.increment(); - logFailureForDlqObjects(List.of(DlqObject.builder() - .withEventHandle(event.getEventHandle()) - .withFailedData(FailedDlqData.builder().withDocument(event.toJsonString()).withIndex(indexName).withMessage(e.getMessage()).build()) - .withPluginName(pluginSetting.getName()) - .withPipelineName(pluginSetting.getPipelineName()) - .withPluginId(pluginSetting.getName()) - .build()), e); + logFailureForDlqObjects(List.of(createDlqObjectFromEvent(event, indexName, e.getMessage())), e); continue; } + Long version = null; + String versionExpressionEvaluationResult = null; + if (versionExpression != null) { + try { + versionExpressionEvaluationResult = event.formatString(versionExpression, expressionEvaluator); + version = Long.valueOf(event.formatString(versionExpression, expressionEvaluator)); + } catch (final NumberFormatException e) { + LOG.warn("Unable to convert the result of evaluating document_version '{}' to Long for an Event. The evaluation result '{}' must be a valid Long type", versionExpression, versionExpressionEvaluationResult); + logFailureForDlqObjects(List.of(createDlqObjectFromEvent(event, indexName, e.getMessage())), e); + dynamicDocumentVersionDroppedEvents.increment(); + } catch (final RuntimeException e) { + LOG.error("There was an exception when evaluating the document_version '{}'. Check the dlq if configured to see details about the affected Event: {}", versionExpression, e.getMessage()); + logFailureForDlqObjects(List.of(createDlqObjectFromEvent(event, indexName, e.getMessage())), e); + dynamicDocumentVersionDroppedEvents.increment(); + } + } + String eventAction = action; if (actions != null) { for (final Map actionEntry: actions) { @@ -373,7 +407,7 @@ public void doOutput(final Collection> records) { StringUtils.equals(action, OpenSearchBulkActions.DELETE.toString())) { serializedJsonNode = SerializedJson.fromJsonNode(event.getJsonNode(), document); } - BulkOperation bulkOperation = getBulkOperationForAction(eventAction, document, indexName, event.getJsonNode()); + BulkOperation bulkOperation = getBulkOperationForAction(eventAction, document, version, indexName, event.getJsonNode()); BulkOperationWrapper bulkOperationWrapper = new BulkOperationWrapper(bulkOperation, event.getEventHandle(), serializedJsonNode); final long estimatedBytesBeforeAdd = bulkRequest.estimateSizeInBytesWithDocument(bulkOperationWrapper); @@ -524,4 +558,20 @@ private void maybeUpdateServerlessNetworkPolicy() { connectionConfiguration.getServerlessVpceId()); } } + + private DlqObject createDlqObjectFromEvent(final Event event, + final String index, + final String message) { + return DlqObject.builder() + .withEventHandle(event.getEventHandle()) + .withFailedData(FailedDlqData.builder() + .withDocument(event.toJsonString()) + .withIndex(index) + .withMessage(message) + .build()) + .withPluginName(pluginSetting.getName()) + .withPipelineName(pluginSetting.getPipelineName()) + .withPluginId(pluginSetting.getName()) + .build(); + } } diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/IndexConfiguration.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/IndexConfiguration.java index 426675e9ea..c6ad893578 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/IndexConfiguration.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/IndexConfiguration.java @@ -8,9 +8,9 @@ import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import org.apache.commons.lang3.EnumUtils; +import org.opensearch.client.opensearch._types.VersionType; import org.opensearch.dataprepper.expression.ExpressionEvaluator; import org.opensearch.dataprepper.model.configuration.PluginSetting; -import org.opensearch.dataprepper.model.event.JacksonEvent; import org.opensearch.dataprepper.model.opensearch.OpenSearchBulkActions; import org.opensearch.dataprepper.model.plugin.InvalidPluginConfigurationException; import org.opensearch.dataprepper.plugins.sink.opensearch.DistributionVersion; @@ -70,6 +70,8 @@ public class IndexConfiguration { public static final String DISTRIBUTION_VERSION = "distribution_version"; public static final String AWS_OPTION = "aws"; public static final String DOCUMENT_ROOT_KEY = "document_root_key"; + public static final String DOCUMENT_VERSION_EXPRESSION = "document_version"; + public static final String DOCUMENT_VERSION_TYPE = "document_version_type"; private IndexType indexType; private TemplateType templateType; @@ -92,6 +94,8 @@ public class IndexConfiguration { private final boolean serverless; private final DistributionVersion distributionVersion; private final String documentRootKey; + private final String versionExpression; + private final VersionType versionType; private static final String S3_PREFIX = "s3://"; private static final String DEFAULT_AWS_REGION = "us-east-1"; @@ -106,6 +110,8 @@ private IndexConfiguration(final Builder builder) { this.s3AwsStsRoleArn = builder.s3AwsStsRoleArn; this.s3AwsExternalId = builder.s3AwsStsExternalId; this.s3Client = builder.s3Client; + this.versionExpression = builder.versionExpression; + this.versionType = builder.versionType; determineTemplateType(builder); @@ -222,6 +228,15 @@ public static IndexConfiguration readIndexConfig(final PluginSetting pluginSetti final String documentIdField = pluginSetting.getStringOrDefault(DOCUMENT_ID_FIELD, null); final String documentId = pluginSetting.getStringOrDefault(DOCUMENT_ID, null); + final String versionExpression = pluginSetting.getStringOrDefault(DOCUMENT_VERSION_EXPRESSION, null); + final String versionType = pluginSetting.getStringOrDefault(DOCUMENT_VERSION_TYPE, null); + + builder = builder.withVersionExpression(versionExpression); + if (versionExpression != null && (!expressionEvaluator.isValidFormatExpressions(versionExpression))) { + throw new InvalidPluginConfigurationException("document_version {} is not a valid format expression."); + } + + builder = builder.withVersionType(versionType); if (Objects.nonNull(documentIdField) && Objects.nonNull(documentId)) { throw new InvalidPluginConfigurationException("Both document_id_field and document_id cannot be used at the same time. It is preferred to only use document_id as document_id_field is deprecated."); @@ -357,6 +372,10 @@ public String getDocumentRootKey() { return documentRootKey; } + public VersionType getVersionType() { return versionType; } + + public String getVersionExpression() { return versionExpression; } + /** * This method is used in the creation of IndexConfiguration object. It takes in the template file path * or index type and returns the index template read from the file or specific to index type or returns an @@ -437,6 +456,8 @@ public static class Builder { private boolean serverless; private DistributionVersion distributionVersion; private String documentRootKey; + private VersionType versionType; + private String versionExpression; public Builder withIndexAlias(final String indexAlias) { checkArgument(indexAlias != null, "indexAlias cannot be null."); @@ -525,7 +546,7 @@ public Builder withIsmPolicyFile(final String ismPolicyFile) { public Builder withAction(final String action, final ExpressionEvaluator expressionEvaluator) { checkArgument((EnumUtils.isValidEnumIgnoreCase(OpenSearchBulkActions.class, action) || - (action.contains("${") && JacksonEvent.isValidFormatExpressions(action, expressionEvaluator))), "Invalid action \"" + action + "\". action must be one of the following: " + Arrays.stream(OpenSearchBulkActions.values()).collect(Collectors.toList())); + (action.contains("${") && expressionEvaluator.isValidFormatExpressions(action))), "action \"" + action + "\" is invalid. action must be one of the following: " + Arrays.stream(OpenSearchBulkActions.values()).collect(Collectors.toList())); this.action = action; return this; } @@ -535,7 +556,7 @@ public Builder withActions(final List> actions, final Expres String action = (String)actionMap.get("type"); if (action != null) { checkArgument((EnumUtils.isValidEnumIgnoreCase(OpenSearchBulkActions.class, action) || - (action.contains("${") && JacksonEvent.isValidFormatExpressions(action, expressionEvaluator))), "Invalid action \"" + action + "\". action must be one of the following: " + Arrays.stream(OpenSearchBulkActions.values()).collect(Collectors.toList())); + (action.contains("${") && expressionEvaluator.isValidFormatExpressions(action))), "action \"" + action + "\". action must be one of the following: " + Arrays.stream(OpenSearchBulkActions.values()).collect(Collectors.toList())); } } this.actions = actions; @@ -591,6 +612,44 @@ public Builder withDocumentRootKey(final String documentRootKey) { return this; } + public Builder withVersionType(final String versionType) { + if (versionType != null) { + try { + this.versionType = getVersionType(versionType); + } catch (final IllegalArgumentException e) { + throw new InvalidPluginConfigurationException( + String.format("version_type %s is invalid. version_type must be one of: %s", + versionType, Arrays.stream(VersionType.values()).collect(Collectors.toList()))); + } + } + + return this; + } + + private VersionType getVersionType(final String versionType) { + switch (versionType.toLowerCase()) { + case "internal": + return VersionType.Internal; + case "external": + return VersionType.External; + case "external_gte": + return VersionType.ExternalGte; + default: + throw new IllegalArgumentException(); + } + } + + public Builder withVersionExpression(final String versionExpression) { + if (versionExpression != null && !versionExpression.contains("${")) { + throw new InvalidPluginConfigurationException( + String.format("document_version %s is invalid. It must be in the format of \"${/key}\" or \"${expression}\"", versionExpression)); + } + + this.versionExpression = versionExpression; + + return this; + } + public IndexConfiguration build() { return new IndexConfiguration(this); } diff --git a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategyTests.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategyTests.java index b2556e3192..e131e5bb7e 100644 --- a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategyTests.java +++ b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategyTests.java @@ -5,31 +5,31 @@ package org.opensearch.dataprepper.plugins.sink.opensearch; -import org.opensearch.client.opensearch._types.ErrorResponse; -import org.opensearch.dataprepper.metrics.MetricNames; -import org.opensearch.dataprepper.metrics.MetricsTestUtil; -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.event.EventHandle; -import org.opensearch.dataprepper.model.configuration.PluginSetting; import io.micrometer.core.instrument.Measurement; import org.hamcrest.MatcherAssert; 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.opensearch.client.opensearch._types.OpenSearchException; +import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.client.opensearch._types.ErrorCause; +import org.opensearch.client.opensearch._types.ErrorResponse; +import org.opensearch.client.opensearch._types.OpenSearchException; import org.opensearch.client.opensearch.core.BulkRequest; import org.opensearch.client.opensearch.core.BulkResponse; import org.opensearch.client.opensearch.core.bulk.BulkOperation; import org.opensearch.client.opensearch.core.bulk.BulkResponseItem; import org.opensearch.client.opensearch.core.bulk.IndexOperation; +import org.opensearch.dataprepper.metrics.MetricNames; +import org.opensearch.dataprepper.metrics.MetricsTestUtil; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.configuration.PluginSetting; +import org.opensearch.dataprepper.model.event.EventHandle; import org.opensearch.dataprepper.plugins.sink.opensearch.bulk.AccumulatingBulkRequest; import org.opensearch.dataprepper.plugins.sink.opensearch.bulk.JavaClientAccumulatingUncompressedBulkRequest; import org.opensearch.dataprepper.plugins.sink.opensearch.bulk.SerializedJson; import org.opensearch.dataprepper.plugins.sink.opensearch.dlq.FailedBulkOperation; import org.opensearch.rest.RestStatus; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.lenient; import java.io.IOException; import java.util.Arrays; @@ -44,13 +44,14 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import static org.mockito.ArgumentMatchers.eq; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.junit.jupiter.MockitoExtension; +import static org.opensearch.dataprepper.plugins.sink.opensearch.BulkRetryStrategy.VERSION_CONFLICT_EXCEPTION_TYPE; @ExtendWith(MockitoExtension.class) public class BulkRetryStrategyTests { @@ -66,6 +67,7 @@ public class BulkRetryStrategyTests { private EventHandle eventHandle2; private EventHandle eventHandle3; private EventHandle eventHandle4; + private EventHandle eventHandle5; @BeforeEach public void setUp() { @@ -75,6 +77,7 @@ public void setUp() { eventHandle2 = mock(EventHandle.class); eventHandle3 = mock(EventHandle.class); eventHandle4 = mock(EventHandle.class); + eventHandle5 = mock(EventHandle.class); lenient().doAnswer(a -> { List l = a.getArgument(0); @@ -220,7 +223,6 @@ public void testExecuteRetryable() throws Exception { accumulatingBulkRequest.addOperation(new BulkOperationWrapper(new BulkOperation.Builder().index(indexOperation2).build(), eventHandle2)); accumulatingBulkRequest.addOperation(new BulkOperationWrapper(new BulkOperation.Builder().index(indexOperation3).build(), eventHandle3)); accumulatingBulkRequest.addOperation(new BulkOperationWrapper(new BulkOperation.Builder().index(indexOperation4).build(), eventHandle4)); - bulkRetryStrategy.execute(accumulatingBulkRequest); assertEquals(3, client.attempt); @@ -228,8 +230,8 @@ public void testExecuteRetryable() throws Exception { assertFalse(client.finalResponse.errors()); assertEquals("3", client.finalRequest.operations().get(0).index().id()); assertEquals("4", client.finalRequest.operations().get(1).index().id()); - assertEquals(numEventsSucceeded, 3); - assertEquals(numEventsFailed, 1); + assertEquals(3, numEventsSucceeded); + assertEquals(1, numEventsFailed); final ArgumentCaptor> failedBulkOperationsCaptor = ArgumentCaptor.forClass(List.class); ArgumentCaptor throwableArgCaptor = ArgumentCaptor.forClass(Throwable.class); @@ -238,12 +240,11 @@ public void testExecuteRetryable() throws Exception { final List failedBulkOperations = failedBulkOperationsCaptor.getValue(); MatcherAssert.assertThat(failedBulkOperations.size(), equalTo(1)); - failedBulkOperations.forEach(failedBulkOperation -> { - final BulkOperationWrapper bulkOperationWithHandle = failedBulkOperation.getBulkOperation(); - final BulkOperation bulkOperation = bulkOperationWithHandle.getBulkOperation(); - MatcherAssert.assertThat(bulkOperation.index().index(), equalTo(testIndex)); - MatcherAssert.assertThat(bulkOperation.index().id(), equalTo(String.valueOf("2"))); - }); + + final BulkOperationWrapper bulkOperationWithHandle = failedBulkOperations.get(0).getBulkOperation(); + final BulkOperation bulkOperation = bulkOperationWithHandle.getBulkOperation(); + MatcherAssert.assertThat(bulkOperation.index().index(), equalTo(testIndex)); + MatcherAssert.assertThat(bulkOperation.index().id(), equalTo("2")); // verify metrics final List documentsSuccessFirstAttemptMeasurements = MetricsTestUtil.getMeasurementList( @@ -403,15 +404,23 @@ public void testExecuteWithMaxRetriesAndSuccesses() throws Exception { final IndexOperation indexOperation2 = new IndexOperation.Builder().index(testIndex).id("2").document(arbitraryDocument()).build(); final IndexOperation indexOperation3 = new IndexOperation.Builder().index(testIndex).id("3").document(arbitraryDocument()).build(); final IndexOperation indexOperation4 = new IndexOperation.Builder().index(testIndex).id("4").document(arbitraryDocument()).build(); + final IndexOperation indexOperation5 = new IndexOperation.Builder().index(testIndex).id("5").document(arbitraryDocument()).build(); final AccumulatingBulkRequest accumulatingBulkRequest = new JavaClientAccumulatingUncompressedBulkRequest(new BulkRequest.Builder()); accumulatingBulkRequest.addOperation(new BulkOperationWrapper(new BulkOperation.Builder().index(indexOperation1).build(), eventHandle1)); accumulatingBulkRequest.addOperation(new BulkOperationWrapper(new BulkOperation.Builder().index(indexOperation2).build(), eventHandle2)); accumulatingBulkRequest.addOperation(new BulkOperationWrapper(new BulkOperation.Builder().index(indexOperation3).build(), eventHandle3)); accumulatingBulkRequest.addOperation(new BulkOperationWrapper(new BulkOperation.Builder().index(indexOperation4).build(), eventHandle4)); + accumulatingBulkRequest.addOperation(new BulkOperationWrapper(new BulkOperation.Builder().index(indexOperation5).build(), eventHandle5)); bulkRetryStrategy.execute(accumulatingBulkRequest); MatcherAssert.assertThat(maxRetriesLimitReached, equalTo(true)); assertEquals(numEventsSucceeded, 2); assertEquals(numEventsFailed, 2); + + final List documentVersionConflictMeasurement = MetricsTestUtil.getMeasurementList( + new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) + .add(BulkRetryStrategy.DOCUMENTS_VERSION_CONFLICT_ERRORS).toString()); + assertEquals(1, documentVersionConflictMeasurement.size()); + assertEquals(1.0, documentVersionConflictMeasurement.get(0).getValue(), 0); } @Test @@ -488,6 +497,10 @@ private static BulkResponseItem internalServerErrorItemResponse(final String ind return customBulkFailureResponse(index, RestStatus.INTERNAL_SERVER_ERROR); } + private static BulkResponseItem versionConflictErrorItemResponse() { + return customBulkFailureResponse(RestStatus.CONFLICT, VERSION_CONFLICT_EXCEPTION_TYPE); + } + private static BulkResponseItem customBulkFailureResponse(final String index, final RestStatus restStatus) { final ErrorCause errorCause = mock(ErrorCause.class); final BulkResponseItem badResponse = mock(BulkResponseItem.class); @@ -496,6 +509,15 @@ private static BulkResponseItem customBulkFailureResponse(final String index, fi return badResponse; } + private static BulkResponseItem customBulkFailureResponse(final RestStatus restStatus, final String errorType) { + final ErrorCause errorCause = mock(ErrorCause.class); + lenient().when(errorCause.type()).thenReturn(errorType); + final BulkResponseItem badResponse = mock(BulkResponseItem.class); + lenient().when(badResponse.status()).thenReturn(restStatus.getStatus()); + lenient().when(badResponse.error()).thenReturn(errorCause); + return badResponse; + } + private SerializedJson arbitraryDocument() { return SerializedJson.fromStringAndOptionals("{}", null, null); } @@ -578,7 +600,8 @@ private BulkResponse bulkMaxRetriesResponseWithSuccesses(final BulkRequest bulkR internalServerErrorItemResponse(index), successItemResponse(index), successItemResponse(index), - tooManyRequestItemResponse(index)); + tooManyRequestItemResponse(index), + versionConflictErrorItemResponse()); return new BulkResponse.Builder().items(bulkItemResponses).errors(true).took(10).build(); } diff --git a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkConfigurationTests.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkConfigurationTests.java index 8789bd872c..78f3efc516 100644 --- a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkConfigurationTests.java +++ b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkConfigurationTests.java @@ -132,16 +132,17 @@ public void testReadESConfigWithBulkActionCreate() { @Test public void testReadESConfigWithBulkActionCreateExpression() { + final String actionFormatExpression = "${getMetadata(\"action\")}"; final Map metadata = new HashMap<>(); metadata.put(IndexConfiguration.INDEX_TYPE, IndexType.TRACE_ANALYTICS_RAW.getValue()); - metadata.put(IndexConfiguration.ACTION, "${getMetadata(\"action\")}"); + metadata.put(IndexConfiguration.ACTION, actionFormatExpression); metadata.put(ConnectionConfiguration.HOSTS, TEST_HOSTS); final PluginSetting pluginSetting = new PluginSetting(PLUGIN_NAME, metadata); pluginSetting.setPipelineName(PIPELINE_NAME); expressionEvaluator = mock(ExpressionEvaluator.class); - when(expressionEvaluator.isValidExpressionStatement(anyString())).thenReturn(true); + when(expressionEvaluator.isValidFormatExpressions(actionFormatExpression)).thenReturn(true); final OpenSearchSinkConfiguration openSearchSinkConfiguration = OpenSearchSinkConfiguration.readESConfig(pluginSetting, expressionEvaluator); diff --git a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/IndexConfigurationTests.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/IndexConfigurationTests.java index 46de83df7e..b0878f53e9 100644 --- a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/IndexConfigurationTests.java +++ b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/IndexConfigurationTests.java @@ -12,6 +12,8 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.ValueSource; +import org.opensearch.dataprepper.expression.ExpressionEvaluator; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.plugin.InvalidPluginConfigurationException; import org.opensearch.dataprepper.plugins.sink.opensearch.DistributionVersion; @@ -49,6 +51,7 @@ import static org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexConfiguration.AWS_OPTION; import static org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexConfiguration.DISTRIBUTION_VERSION; import static org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexConfiguration.DOCUMENT_ROOT_KEY; +import static org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexConfiguration.DOCUMENT_VERSION_EXPRESSION; import static org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexConfiguration.SERVERLESS; import static org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexConfiguration.TEMPLATE_TYPE; import static org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexConstants.RAW_DEFAULT_TEMPLATE_FILE; @@ -473,6 +476,41 @@ public void testReadIndexConfig_emptyDocumentRootKey() { assertThrows(IllegalArgumentException.class, () -> IndexConfiguration.readIndexConfig(pluginSetting)); } + @ParameterizedTest + @ValueSource(strings = {"${key}", "${getMetadata(\"key\")}"}) + public void testReadIndexConfig_withValidDocumentVersionExpression(final String versionExpression) { + + final ExpressionEvaluator expressionEvaluator = mock(ExpressionEvaluator.class); + when(expressionEvaluator.isValidFormatExpressions(versionExpression)).thenReturn(true); + + final Map metadata = initializeConfigMetaData( + IndexType.CUSTOM.getValue(), "foo", null, null, null, null, null); + metadata.put(DOCUMENT_VERSION_EXPRESSION, versionExpression); + + final PluginSetting pluginSetting = getPluginSetting(metadata); + + final IndexConfiguration indexConfiguration = IndexConfiguration.readIndexConfig(pluginSetting, expressionEvaluator); + + assertThat(indexConfiguration, notNullValue()); + assertThat(indexConfiguration.getVersionExpression(), equalTo(versionExpression)); + } + + @Test + public void testReadIndexConfig_withInvalidDocumentVersionExpression_throws_InvalidPluginConfigurationException() { + final String versionExpression = UUID.randomUUID().toString(); + + final ExpressionEvaluator expressionEvaluator = mock(ExpressionEvaluator.class); + when(expressionEvaluator.isValidFormatExpressions(versionExpression)).thenReturn(false); + + final Map metadata = initializeConfigMetaData( + IndexType.CUSTOM.getValue(), "foo", null, null, null, null, null); + metadata.put(DOCUMENT_VERSION_EXPRESSION, versionExpression); + + final PluginSetting pluginSetting = getPluginSetting(metadata); + + assertThrows(InvalidPluginConfigurationException.class, () -> IndexConfiguration.readIndexConfig(pluginSetting, expressionEvaluator)); + } + @Test void getTemplateType_defaults_to_V1() { final Map metadata = initializeConfigMetaData( From a1c56d06fa9091867b346561f3797a76531911aa Mon Sep 17 00:00:00 2001 From: Taylor Gray Date: Tue, 7 Nov 2023 19:16:16 -0600 Subject: [PATCH 34/53] =?UTF-8?q?Add=20dynamodb=5Fitem=5Fversion=20metadat?= =?UTF-8?q?a=20that=20is=20derived=20from=20timestamp=20for=E2=80=A6=20(#3?= =?UTF-8?q?596)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add dynamodb_item_version metadata that is derived from timestamp for stream events Signed-off-by: Taylor Gray --- .../converter/MetadataKeyAttributes.java | 2 + .../dynamodb/converter/RecordConverter.java | 23 +++- .../converter/StreamRecordConverter.java | 22 ++- .../converter/ExportRecordConverterTest.java | 5 + .../converter/StreamRecordConverterTest.java | 129 +++++++++++++++--- 5 files changed, 159 insertions(+), 22 deletions(-) diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java index 03c92c09e0..51c85bbefb 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/MetadataKeyAttributes.java @@ -14,6 +14,8 @@ public class MetadataKeyAttributes { static final String EVENT_TIMESTAMP_METADATA_ATTRIBUTE = "dynamodb_timestamp"; + static final String EVENT_DYNAMODB_ITEM_VERSION = "dynamodb_item_version"; + static final String EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE = "opensearch_action"; static final String DDB_STREAM_EVENT_NAME_METADATA_ATTRIBUTE = "dynamodb_event_name"; diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java index d15a18b989..7123b6825e 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/RecordConverter.java @@ -14,9 +14,11 @@ import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo; +import java.time.Instant; import java.util.Map; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.DDB_STREAM_EVENT_NAME_METADATA_ATTRIBUTE; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_DYNAMODB_ITEM_VERSION; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_TABLE_NAME_METADATA_ATTRIBUTE; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_TIMESTAMP_METADATA_ATTRIBUTE; @@ -72,17 +74,31 @@ void flushBuffer() throws Exception { * @param eventName Event name * @throws Exception Exception if failed to write to buffer. */ - public void addToBuffer(final AcknowledgementSet acknowledgementSet, Map data, Map keys, long eventCreationTimeMillis, String eventName) throws Exception { + public void addToBuffer(final AcknowledgementSet acknowledgementSet, + final Map data, + final Map keys, + final long eventCreationTimeMillis, + final long eventVersionNumber, + final String eventName) throws Exception { Event event = JacksonEvent.builder() .withEventType(getEventType()) .withData(data) .build(); + + // Only set external origination time for stream events, not export + if (eventName != null) { + final Instant externalOriginationTime = Instant.ofEpochMilli(eventCreationTimeMillis); + event.getEventHandle().setExternalOriginationTime(externalOriginationTime); + event.getMetadata().setExternalOriginationTime(externalOriginationTime); + } EventMetadata eventMetadata = event.getMetadata(); eventMetadata.setAttribute(EVENT_TABLE_NAME_METADATA_ATTRIBUTE, tableInfo.getTableName()); eventMetadata.setAttribute(EVENT_TIMESTAMP_METADATA_ATTRIBUTE, eventCreationTimeMillis); eventMetadata.setAttribute(DDB_STREAM_EVENT_NAME_METADATA_ATTRIBUTE, eventName); eventMetadata.setAttribute(EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE, mapStreamEventNameToBulkAction(eventName)); + eventMetadata.setAttribute(EVENT_DYNAMODB_ITEM_VERSION, eventVersionNumber); + String partitionKey = getAttributeValue(keys, tableInfo.getMetadata().getPartitionKeyAttributeName()); eventMetadata.setAttribute(PARTITION_KEY_METADATA_ATTRIBUTE, partitionKey); @@ -101,8 +117,9 @@ public void addToBuffer(final AcknowledgementSet acknowledgementSet, Map data) throws Exception { // Export data doesn't have an event timestamp - // Default to current timestamp when the event is added to buffer - addToBuffer(acknowledgementSet, data, data, System.currentTimeMillis(), null); + // We consider this to be time of 0, meaning stream records will always be considered as newer + // than export records + addToBuffer(acknowledgementSet, data, data, System.currentTimeMillis(), 0L, null); } private String mapStreamEventNameToBulkAction(final String streamEventName) { diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java index 17c4e41e93..942e4d6a81 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverter.java @@ -20,6 +20,7 @@ import software.amazon.awssdk.services.dynamodb.model.AttributeValue; import software.amazon.awssdk.services.dynamodb.model.Record; +import java.time.Instant; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -41,6 +42,9 @@ public class StreamRecordConverter extends RecordConverter { private final Counter changeEventSuccessCounter; private final Counter changeEventErrorCounter; + private Instant currentSecond; + private int recordsSeenThisSecond = 0; + public StreamRecordConverter(final BufferAccumulator> bufferAccumulator, TableInfo tableInfo, PluginMetrics pluginMetrics) { super(bufferAccumulator, tableInfo); this.pluginMetrics = pluginMetrics; @@ -64,7 +68,8 @@ public void writeToBuffer(final AcknowledgementSet acknowledgementSet, List keys = convertKeys(record.dynamodb().keys()); try { - addToBuffer(acknowledgementSet, data, keys, record.dynamodb().approximateCreationDateTime().toEpochMilli(), record.eventNameAsString()); + final long eventCreationTimeMillis = calculateTieBreakingVersionFromTimestamp(record.dynamodb().approximateCreationDateTime()); + addToBuffer(acknowledgementSet, data, keys, record.dynamodb().approximateCreationDateTime().toEpochMilli(), eventCreationTimeMillis, record.eventNameAsString()); eventCount++; } catch (Exception e) { // will this cause too many logs? @@ -116,4 +121,19 @@ private Map convertKeys(Map keys) { return result; } + + private long calculateTieBreakingVersionFromTimestamp(final Instant eventTimeInSeconds) { + if (currentSecond == null) { + currentSecond = eventTimeInSeconds; + } else if (currentSecond.isAfter(eventTimeInSeconds)) { + return eventTimeInSeconds.getEpochSecond() * 1_000_000; + } else if (currentSecond.isBefore(eventTimeInSeconds)) { + recordsSeenThisSecond = 0; + currentSecond = eventTimeInSeconds; + } else { + recordsSeenThisSecond++; + } + + return eventTimeInSeconds.getEpochSecond() * 1_000_000 + recordsSeenThisSecond; + } } diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java index 8f7c0a8f23..093536bbbd 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/ExportRecordConverterTest.java @@ -41,6 +41,7 @@ import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.ExportRecordConverter.EXPORT_RECORDS_PROCESSED_COUNT; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.ExportRecordConverter.EXPORT_RECORDS_PROCESSING_ERROR_COUNT; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.DDB_STREAM_EVENT_NAME_METADATA_ATTRIBUTE; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_DYNAMODB_ITEM_VERSION; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_TIMESTAMP_METADATA_ATTRIBUTE; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.PARTITION_KEY_METADATA_ATTRIBUTE; @@ -144,5 +145,9 @@ void test_writeSingleRecordToBuffer() throws Exception { assertThat(event.getMetadata().getAttribute(EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE), equalTo(OpenSearchBulkActions.INDEX.toString())); assertThat(event.getMetadata().getAttribute(EVENT_TIMESTAMP_METADATA_ATTRIBUTE), notNullValue()); assertThat(event.getMetadata().getAttribute(DDB_STREAM_EVENT_NAME_METADATA_ATTRIBUTE), nullValue()); + assertThat(event.getMetadata().getAttribute(EVENT_TIMESTAMP_METADATA_ATTRIBUTE), notNullValue()); + assertThat(event.getMetadata().getAttribute(EVENT_DYNAMODB_ITEM_VERSION), equalTo(0L)); + assertThat(event.getEventHandle(), notNullValue()); + assertThat(event.getEventHandle().getExternalOriginationTime(), nullValue()); } } \ No newline at end of file diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java index 75e5274dfc..d8977b7875 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/converter/StreamRecordConverterTest.java @@ -42,6 +42,7 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.DDB_STREAM_EVENT_NAME_METADATA_ATTRIBUTE; +import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_DYNAMODB_ITEM_VERSION; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.EVENT_TIMESTAMP_METADATA_ATTRIBUTE; import static org.opensearch.dataprepper.plugins.source.dynamodb.converter.MetadataKeyAttributes.PARTITION_KEY_METADATA_ATTRIBUTE; @@ -99,7 +100,7 @@ void test_writeToBuffer() throws Exception { int numberOfRecords = random.nextInt(10); - List records = buildRecords(numberOfRecords); + List records = buildRecords(numberOfRecords, Instant.now()); StreamRecordConverter recordConverter = new StreamRecordConverter(bufferAccumulator, tableInfo, pluginMetrics); @@ -115,7 +116,7 @@ void test_writeToBuffer() throws Exception { @Test void test_writeSingleRecordToBuffer() throws Exception { - List records = buildRecords(1); + List records = buildRecords(1, Instant.now()); final ArgumentCaptor recordArgumentCaptor = ArgumentCaptor.forClass(Record.class); software.amazon.awssdk.services.dynamodb.model.Record record = records.get(0); StreamRecordConverter recordConverter = new StreamRecordConverter(bufferAccumulator, tableInfo, pluginMetrics); @@ -142,27 +143,119 @@ void test_writeSingleRecordToBuffer() throws Exception { verifyNoInteractions(changeEventErrorCounter); } - private List buildRecords(int count) { + @Test + void writingToBuffer_with_nth_event_in_that_second_returns_expected_that_timestamp() throws Exception { + final long currentSecond = 1699336310; + final Instant timestamp = Instant.ofEpochSecond(currentSecond); + final Instant olderSecond = Instant.ofEpochSecond(currentSecond - 1); + final Instant newerSecond = Instant.ofEpochSecond(currentSecond + 1); + + List records = buildRecords(2, timestamp); + records.add(buildRecord(olderSecond)); + records.add(buildRecord(newerSecond)); + + final ArgumentCaptor recordArgumentCaptor = ArgumentCaptor.forClass(Record.class); + StreamRecordConverter recordConverter = new StreamRecordConverter(bufferAccumulator, tableInfo, pluginMetrics); + doNothing().when(bufferAccumulator).add(recordArgumentCaptor.capture()); + + recordConverter.writeToBuffer(null, records); + + verify(bufferAccumulator, times(4)).add(any(Record.class)); + verify(bufferAccumulator).flush(); + verify(changeEventSuccessCounter).increment(anyDouble()); + assertThat(recordArgumentCaptor.getValue().getData(), notNullValue()); + + final List createdEvents = recordArgumentCaptor.getAllValues(); + + assertThat(createdEvents.size(), equalTo(records.size())); + + JacksonEvent firstEventForSecond = (JacksonEvent) createdEvents.get(0).getData(); + + assertThat(firstEventForSecond.getMetadata(), notNullValue()); + String partitionKey = records.get(0).dynamodb().keys().get(partitionKeyAttrName).s(); + String sortKey = records.get(0).dynamodb().keys().get(sortKeyAttrName).s(); + assertThat(firstEventForSecond.getMetadata().getAttribute(PARTITION_KEY_METADATA_ATTRIBUTE), equalTo(partitionKey)); + assertThat(firstEventForSecond.getMetadata().getAttribute(SORT_KEY_METADATA_ATTRIBUTE), equalTo(sortKey)); + assertThat(firstEventForSecond.getMetadata().getAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE), equalTo(partitionKey + "|" + sortKey)); + assertThat(firstEventForSecond.getMetadata().getAttribute(EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE), equalTo(OpenSearchBulkActions.INDEX.toString())); + assertThat(firstEventForSecond.getMetadata().getAttribute(DDB_STREAM_EVENT_NAME_METADATA_ATTRIBUTE), equalTo("INSERT")); + assertThat(firstEventForSecond.getMetadata().getAttribute(EVENT_TIMESTAMP_METADATA_ATTRIBUTE), equalTo(timestamp.toEpochMilli())); + assertThat(firstEventForSecond.getMetadata().getAttribute(EVENT_DYNAMODB_ITEM_VERSION), equalTo(timestamp.toEpochMilli() * 1000)); + assertThat(firstEventForSecond.getEventHandle(), notNullValue()); + assertThat(firstEventForSecond.getEventHandle().getExternalOriginationTime(), equalTo(timestamp)); + + JacksonEvent secondEventForSameSecond = (JacksonEvent) createdEvents.get(1).getData(); + + assertThat(secondEventForSameSecond.getMetadata(), notNullValue()); + String secondPartitionKey = records.get(1).dynamodb().keys().get(partitionKeyAttrName).s(); + String secondSortKey = records.get(1).dynamodb().keys().get(sortKeyAttrName).s(); + assertThat(secondEventForSameSecond.getMetadata().getAttribute(PARTITION_KEY_METADATA_ATTRIBUTE), equalTo(secondPartitionKey)); + assertThat(secondEventForSameSecond.getMetadata().getAttribute(SORT_KEY_METADATA_ATTRIBUTE), equalTo(secondSortKey)); + assertThat(secondEventForSameSecond.getMetadata().getAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE), equalTo(secondPartitionKey + "|" + secondSortKey)); + assertThat(secondEventForSameSecond.getMetadata().getAttribute(EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE), equalTo(OpenSearchBulkActions.INDEX.toString())); + assertThat(secondEventForSameSecond.getMetadata().getAttribute(DDB_STREAM_EVENT_NAME_METADATA_ATTRIBUTE), equalTo("INSERT")); + assertThat(secondEventForSameSecond.getMetadata().getAttribute(EVENT_TIMESTAMP_METADATA_ATTRIBUTE), equalTo(timestamp.toEpochMilli())); + assertThat(secondEventForSameSecond.getMetadata().getAttribute(EVENT_DYNAMODB_ITEM_VERSION), equalTo(timestamp.toEpochMilli() * 1000 + 1)); + assertThat(secondEventForSameSecond.getEventHandle(), notNullValue()); + assertThat(secondEventForSameSecond.getEventHandle().getExternalOriginationTime(), equalTo(timestamp)); + + JacksonEvent thirdEventWithOlderSecond = (JacksonEvent) createdEvents.get(2).getData(); + + assertThat(thirdEventWithOlderSecond.getMetadata(), notNullValue()); + String thirdPartitionKey = records.get(2).dynamodb().keys().get(partitionKeyAttrName).s(); + String thirdSortKey = records.get(2).dynamodb().keys().get(sortKeyAttrName).s(); + assertThat(thirdEventWithOlderSecond.getMetadata().getAttribute(PARTITION_KEY_METADATA_ATTRIBUTE), equalTo(thirdPartitionKey)); + assertThat(thirdEventWithOlderSecond.getMetadata().getAttribute(SORT_KEY_METADATA_ATTRIBUTE), equalTo(thirdSortKey)); + assertThat(thirdEventWithOlderSecond.getMetadata().getAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE), equalTo(thirdPartitionKey + "|" + thirdSortKey)); + assertThat(thirdEventWithOlderSecond.getMetadata().getAttribute(EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE), equalTo(OpenSearchBulkActions.INDEX.toString())); + assertThat(thirdEventWithOlderSecond.getMetadata().getAttribute(DDB_STREAM_EVENT_NAME_METADATA_ATTRIBUTE), equalTo("INSERT")); + assertThat(thirdEventWithOlderSecond.getMetadata().getAttribute(EVENT_TIMESTAMP_METADATA_ATTRIBUTE), equalTo(olderSecond.toEpochMilli())); + assertThat(thirdEventWithOlderSecond.getMetadata().getAttribute(EVENT_DYNAMODB_ITEM_VERSION), equalTo(olderSecond.toEpochMilli() * 1000)); + assertThat(thirdEventWithOlderSecond.getEventHandle(), notNullValue()); + assertThat(thirdEventWithOlderSecond.getEventHandle().getExternalOriginationTime(), equalTo(olderSecond)); + + JacksonEvent fourthEventWithNewerSecond = (JacksonEvent) createdEvents.get(3).getData(); + + assertThat(fourthEventWithNewerSecond.getMetadata(), notNullValue()); + String fourthPartitionKey = records.get(3).dynamodb().keys().get(partitionKeyAttrName).s(); + String fourthSortKey = records.get(3).dynamodb().keys().get(sortKeyAttrName).s(); + assertThat(fourthEventWithNewerSecond.getMetadata().getAttribute(PARTITION_KEY_METADATA_ATTRIBUTE), equalTo(fourthPartitionKey)); + assertThat(fourthEventWithNewerSecond.getMetadata().getAttribute(SORT_KEY_METADATA_ATTRIBUTE), equalTo(fourthSortKey)); + assertThat(fourthEventWithNewerSecond.getMetadata().getAttribute(PRIMARY_KEY_DOCUMENT_ID_METADATA_ATTRIBUTE), equalTo(fourthPartitionKey + "|" + fourthSortKey)); + assertThat(fourthEventWithNewerSecond.getMetadata().getAttribute(EVENT_NAME_BULK_ACTION_METADATA_ATTRIBUTE), equalTo(OpenSearchBulkActions.INDEX.toString())); + assertThat(fourthEventWithNewerSecond.getMetadata().getAttribute(DDB_STREAM_EVENT_NAME_METADATA_ATTRIBUTE), equalTo("INSERT")); + assertThat(fourthEventWithNewerSecond.getMetadata().getAttribute(EVENT_TIMESTAMP_METADATA_ATTRIBUTE), equalTo(newerSecond.toEpochMilli())); + assertThat(fourthEventWithNewerSecond.getMetadata().getAttribute(EVENT_DYNAMODB_ITEM_VERSION), equalTo(newerSecond.toEpochMilli() * 1000)); + assertThat(fourthEventWithNewerSecond.getEventHandle(), notNullValue()); + assertThat(fourthEventWithNewerSecond.getEventHandle().getExternalOriginationTime(), equalTo(newerSecond)); + + verifyNoInteractions(changeEventErrorCounter); + } + + private List buildRecords(int count, final Instant creationTime) { List records = new ArrayList<>(); for (int i = 0; i < count; i++) { - Map data = Map.of( - partitionKeyAttrName, AttributeValue.builder().s(UUID.randomUUID().toString()).build(), - sortKeyAttrName, AttributeValue.builder().s(UUID.randomUUID().toString()).build()); - - StreamRecord streamRecord = StreamRecord.builder() - .newImage(data) - .keys(data) - .sequenceNumber(UUID.randomUUID().toString()) - .approximateCreationDateTime(Instant.now()) - .build(); - software.amazon.awssdk.services.dynamodb.model.Record record = software.amazon.awssdk.services.dynamodb.model.Record.builder() - .dynamodb(streamRecord) - .eventName(OperationType.INSERT) - .build(); - records.add(record); + records.add(buildRecord(creationTime)); } return records; } + private software.amazon.awssdk.services.dynamodb.model.Record buildRecord(final Instant creationTime) { + Map data = Map.of( + partitionKeyAttrName, AttributeValue.builder().s(UUID.randomUUID().toString()).build(), + sortKeyAttrName, AttributeValue.builder().s(UUID.randomUUID().toString()).build()); + StreamRecord streamRecord = StreamRecord.builder() + .newImage(data) + .keys(data) + .sequenceNumber(UUID.randomUUID().toString()) + .approximateCreationDateTime(creationTime) + .build(); + software.amazon.awssdk.services.dynamodb.model.Record record = software.amazon.awssdk.services.dynamodb.model.Record.builder() + .dynamodb(streamRecord) + .eventName(OperationType.INSERT) + .build(); + return record; + } + } \ No newline at end of file From c8d02f9ec997048d1984d2ded7ed9c06cea1f643 Mon Sep 17 00:00:00 2001 From: Dinu John <86094133+dinujoh@users.noreply.github.com> Date: Tue, 7 Nov 2023 22:45:19 -0600 Subject: [PATCH 35/53] Add Kafka Producer Metrics and Kafka Buffer Metrics (#3598) Signed-off-by: Dinu John <86094133+dinujoh@users.noreply.github.com> --- .../kafka/sink/KafkaSinkAvroTypeIT.java | 6 +- .../kafka/sink/KafkaSinkJsonTypeIT.java | 6 +- .../kafka/sink/KafkaSinkPlainTextTypeIT.java | 6 +- .../plugins/kafka/buffer/KafkaBuffer.java | 9 +- .../kafka/buffer/KafkaBufferConfig.java | 7 + .../kafka/consumer/KafkaCustomConsumer.java | 8 +- .../consumer/KafkaCustomConsumerFactory.java | 7 +- .../kafka/producer/KafkaCustomProducer.java | 26 ++- .../producer/KafkaCustomProducerFactory.java | 8 +- .../plugins/kafka/sink/KafkaSink.java | 8 +- .../plugins/kafka/source/KafkaSource.java | 4 +- ...cs.java => KafkaTopicConsumerMetrics.java} | 70 ++++--- .../kafka/util/KafkaTopicProducerMetrics.java | 133 ++++++++++++ .../plugins/kafka/buffer/KafkaBufferTest.java | 3 +- .../consumer/KafkaCustomConsumerTest.java | 6 +- .../producer/KafkaCustomProducerTest.java | 41 ++-- .../plugins/kafka/sink/KafkaSinkTest.java | 6 +- ...va => KafkaTopicConsumerMetricsTests.java} | 9 +- .../util/KafkaTopicProducerMetricsTests.java | 190 ++++++++++++++++++ 19 files changed, 472 insertions(+), 81 deletions(-) rename data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/{KafkaTopicMetrics.java => KafkaTopicConsumerMetrics.java} (78%) create mode 100644 data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicProducerMetrics.java rename data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/{KafkaTopicMetricsTests.java => KafkaTopicConsumerMetricsTests.java} (97%) create mode 100644 data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicProducerMetricsTests.java diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkAvroTypeIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkAvroTypeIT.java index da901ed241..69dc2c9d5e 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkAvroTypeIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkAvroTypeIT.java @@ -26,6 +26,7 @@ import org.mockito.Mock; 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.event.Event; import org.opensearch.dataprepper.model.event.JacksonEvent; @@ -81,6 +82,9 @@ public class KafkaSinkAvroTypeIT { @Mock private PluginFactory pluginFactory; + @Mock + private PluginMetrics pluginMetrics; + private SinkContext sinkContext; private String registryUrl; @@ -104,7 +108,7 @@ public class KafkaSinkAvroTypeIT { public KafkaSink createObjectUnderTest() { - return new KafkaSink(pluginSetting, kafkaSinkConfig, pluginFactory, evaluator, sinkContext, awsCredentialsSupplier); + return new KafkaSink(pluginSetting, kafkaSinkConfig, pluginFactory, pluginMetrics, evaluator, sinkContext, awsCredentialsSupplier); } @BeforeEach diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkJsonTypeIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkJsonTypeIT.java index bf2d1cb3df..51b74ae316 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkJsonTypeIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkJsonTypeIT.java @@ -24,6 +24,7 @@ import org.mockito.Mock; 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.event.Event; import org.opensearch.dataprepper.model.event.JacksonEvent; @@ -73,6 +74,9 @@ public class KafkaSinkJsonTypeIT { @Mock private PluginFactory pluginFactory; + @Mock + private PluginMetrics pluginMetrics; + private SinkContext sinkContext; @Mock @@ -95,7 +99,7 @@ public class KafkaSinkJsonTypeIT { public KafkaSink createObjectUnderTest() { - return new KafkaSink(pluginSetting, kafkaSinkConfig, pluginFactory, evaluator, sinkContext, awsCredentialsSupplier); + return new KafkaSink(pluginSetting, kafkaSinkConfig, pluginFactory, pluginMetrics, evaluator, sinkContext, awsCredentialsSupplier); } @BeforeEach diff --git a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkPlainTextTypeIT.java b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkPlainTextTypeIT.java index c00304d4e8..ea855b0ec3 100644 --- a/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkPlainTextTypeIT.java +++ b/data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkPlainTextTypeIT.java @@ -22,6 +22,7 @@ import org.mockito.Mock; 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.event.Event; import org.opensearch.dataprepper.model.event.JacksonEvent; @@ -71,6 +72,9 @@ public class KafkaSinkPlainTextTypeIT { @Mock private PluginFactory pluginFactory; + @Mock + private PluginMetrics pluginMetrics; + private SinkContext sinkContext; @Mock @@ -93,7 +97,7 @@ public class KafkaSinkPlainTextTypeIT { public KafkaSink createObjectUnderTest() { - return new KafkaSink(pluginSetting, kafkaSinkConfig, pluginFactory, evaluator, sinkContext, awsCredentialsSupplier); + return new KafkaSink(pluginSetting, kafkaSinkConfig, pluginFactory, pluginMetrics, evaluator, sinkContext, awsCredentialsSupplier); } @BeforeEach 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 86a934e972..3a358c23fe 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 @@ -43,6 +43,8 @@ public class KafkaBuffer> extends AbstractBuffer { static final long EXECUTOR_SERVICE_SHUTDOWN_TIMEOUT = 30L; public static final int INNER_BUFFER_CAPACITY = 1000000; public static final int INNER_BUFFER_BATCH_SIZE = 250000; + static final String WRITE = "Write"; + static final String READ = "Read"; private final KafkaCustomProducer producer; private final AbstractBuffer innerBuffer; private final ExecutorService executorService; @@ -58,12 +60,15 @@ public KafkaBuffer(final PluginSetting pluginSetting, final KafkaBufferConfig ka SerializationFactory serializationFactory = new SerializationFactory(); final KafkaCustomProducerFactory kafkaCustomProducerFactory = new KafkaCustomProducerFactory(serializationFactory, awsCredentialsSupplier); this.byteDecoder = byteDecoder; - producer = kafkaCustomProducerFactory.createProducer(kafkaBufferConfig, pluginFactory, pluginSetting, null, null); + 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, 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); + final PluginMetrics consumerMetrics = PluginMetrics.fromNames(metricPrefixName + READ, pluginSetting.getPipelineName()); final List consumers = kafkaCustomConsumerFactory.createConsumersForTopic(kafkaBufferConfig, kafkaBufferConfig.getTopic(), - innerBuffer, pluginMetrics, acknowledgementSetManager, byteDecoder, shutdownInProgress); + innerBuffer, consumerMetrics, acknowledgementSetManager, byteDecoder, shutdownInProgress, false); this.executorService = Executors.newFixedThreadPool(consumers.size()); consumers.forEach(this.executorService::submit); diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferConfig.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferConfig.java index 9574ae94e2..581d0bcdbd 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferConfig.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferConfig.java @@ -53,6 +53,9 @@ class KafkaBufferConfig implements KafkaProducerConfig, KafkaConsumerConfig { @JsonProperty("drain_timeout") private Duration drainTimeout = DEFAULT_DRAIN_TIMEOUT; + @JsonProperty("custom_metric_prefix") + private String customMetricPrefix; + public List getBootstrapServers() { if (Objects.nonNull(bootStrapServers)) { @@ -131,4 +134,8 @@ public boolean getAcknowledgementsEnabled() { public Duration getDrainTimeout() { return drainTimeout; } + + public Optional getCustomMetricPrefix() { + return Optional.ofNullable(customMetricPrefix); + } } 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 e36e93e2d9..d8d577e49a 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 @@ -32,7 +32,7 @@ import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaConsumerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaKeyMode; -import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicMetrics; +import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicConsumerMetrics; import org.opensearch.dataprepper.plugins.kafka.util.LogRateLimiter; import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; import org.slf4j.Logger; @@ -86,7 +86,7 @@ public class KafkaCustomConsumer implements Runnable, ConsumerRebalanceListener private List> acknowledgedOffsets; private final boolean acknowledgementsEnabled; private final Duration acknowledgementsTimeout; - private final KafkaTopicMetrics topicMetrics; + private final KafkaTopicConsumerMetrics topicMetrics; private long metricsUpdatedTime; private final AtomicInteger numberOfAcksPending; private long numRecordsCommitted = 0; @@ -101,7 +101,7 @@ public KafkaCustomConsumer(final KafkaConsumer consumer, final String schemaType, final AcknowledgementSetManager acknowledgementSetManager, final ByteDecoder byteDecoder, - KafkaTopicMetrics topicMetrics) { + final KafkaTopicConsumerMetrics topicMetrics) { this.topicName = topicConfig.getName(); this.topicConfig = topicConfig; this.shutdownInProgress = shutdownInProgress; @@ -127,7 +127,7 @@ public KafkaCustomConsumer(final KafkaConsumer consumer, this.errLogRateLimiter = new LogRateLimiter(2, System.currentTimeMillis()); } - KafkaTopicMetrics getTopicMetrics() { + KafkaTopicConsumerMetrics getTopicMetrics() { return topicMetrics; } 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 de553afc99..5954e91c16 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 @@ -40,7 +40,7 @@ import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.opensearch.dataprepper.plugins.kafka.util.ClientDNSLookupType; import org.opensearch.dataprepper.plugins.kafka.util.KafkaSecurityConfigurer; -import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicMetrics; +import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicConsumerMetrics; import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -72,10 +72,11 @@ public List createConsumersForTopic(final KafkaConsumerConf final Buffer> buffer, final PluginMetrics pluginMetrics, final AcknowledgementSetManager acknowledgementSetManager, final ByteDecoder byteDecoder, - final AtomicBoolean shutdownInProgress) { + final AtomicBoolean shutdownInProgress, + final boolean topicNameInMetrics) { Properties authProperties = new Properties(); KafkaSecurityConfigurer.setAuthProperties(authProperties, kafkaConsumerConfig, LOG); - KafkaTopicMetrics topicMetrics = new KafkaTopicMetrics(topic.getName(), pluginMetrics); + KafkaTopicConsumerMetrics topicMetrics = new KafkaTopicConsumerMetrics(topic.getName(), pluginMetrics, topicNameInMetrics); Properties consumerProperties = getConsumerProperties(kafkaConsumerConfig, topic, authProperties); MessageFormat schema = MessageFormat.getByMessageFormatByName(schemaType); 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 ddca58ae2e..ab745b9ebf 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 @@ -18,7 +18,7 @@ import org.apache.avro.generic.GenericRecord; import org.apache.commons.lang3.ObjectUtils; import org.apache.kafka.clients.producer.Callback; -import org.apache.kafka.clients.producer.Producer; +import org.apache.kafka.clients.producer.KafkaProducer; import org.apache.kafka.clients.producer.ProducerRecord; import org.opensearch.dataprepper.expression.ExpressionEvaluator; import org.opensearch.dataprepper.model.event.Event; @@ -28,6 +28,7 @@ import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerConfig; import org.opensearch.dataprepper.plugins.kafka.service.SchemaService; import org.opensearch.dataprepper.plugins.kafka.sink.DLQSink; +import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicProducerMetrics; import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -49,7 +50,7 @@ public class KafkaCustomProducer { private static final Logger LOG = LoggerFactory.getLogger(KafkaCustomProducer.class); - private final Producer producer; + private final KafkaProducer producer; private final KafkaProducerConfig kafkaProducerConfig; @@ -69,28 +70,39 @@ public class KafkaCustomProducer { private final SchemaService schemaService; + private final KafkaTopicProducerMetrics topicMetrics; - public KafkaCustomProducer(final Producer producer, + + public KafkaCustomProducer(final KafkaProducer producer, final KafkaProducerConfig kafkaProducerConfig, final DLQSink dlqSink, final ExpressionEvaluator expressionEvaluator, - final String tagTargetKey + final String tagTargetKey, + final KafkaTopicProducerMetrics topicMetrics ) { this.producer = producer; this.kafkaProducerConfig = kafkaProducerConfig; this.dlqSink = dlqSink; - bufferedEventHandles = new LinkedList<>(); + this.bufferedEventHandles = new LinkedList<>(); this.expressionEvaluator = expressionEvaluator; this.tagTargetKey = tagTargetKey; this.topicName = ObjectUtils.isEmpty(kafkaProducerConfig.getTopic()) ? null : kafkaProducerConfig.getTopic().getName(); this.serdeFormat = ObjectUtils.isEmpty(kafkaProducerConfig.getSerdeFormat()) ? null : kafkaProducerConfig.getSerdeFormat(); - schemaService = new SchemaService.SchemaServiceBuilder().getFetchSchemaService(topicName, kafkaProducerConfig.getSchemaConfig()).build(); + this.schemaService = new SchemaService.SchemaServiceBuilder().getFetchSchemaService(topicName, kafkaProducerConfig.getSchemaConfig()).build(); + this.topicMetrics = topicMetrics; + this.topicMetrics.register(this.producer); + } + + KafkaTopicProducerMetrics getTopicMetrics() { + return topicMetrics; } public void produceRawData(final byte[] bytes, final String key) { try { send(topicName, key, bytes).get(); + topicMetrics.update(producer); } catch (Exception e) { + // TODO: Metrics LOG.error("Error occurred while publishing {}", e.getMessage()); } } @@ -107,7 +119,9 @@ public void produceRecords(final Record record) { } else { publishPlaintextMessage(record, key); } + topicMetrics.update(producer); } catch (Exception e) { + // TODO: Metrics LOG.error("Error occurred while publishing {}", e.getMessage()); releaseEventHandles(false); } 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 ccd28eba35..fa378bfc8c 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 @@ -9,6 +9,7 @@ import org.apache.kafka.common.serialization.Serializer; 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; @@ -27,6 +28,7 @@ import org.opensearch.dataprepper.plugins.kafka.service.TopicService; import org.opensearch.dataprepper.plugins.kafka.sink.DLQSink; import org.opensearch.dataprepper.plugins.kafka.util.KafkaSecurityConfigurer; +import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicProducerMetrics; import org.opensearch.dataprepper.plugins.kafka.util.RestUtils; import org.opensearch.dataprepper.plugins.kafka.util.SinkPropertyConfigurer; import org.slf4j.Logger; @@ -46,7 +48,8 @@ public KafkaCustomProducerFactory(final SerializationFactory serializationFactor } public KafkaCustomProducer createProducer(final KafkaProducerConfig kafkaProducerConfig, final PluginFactory pluginFactory, final PluginSetting pluginSetting, - final ExpressionEvaluator expressionEvaluator, final SinkContext sinkContext) { + final ExpressionEvaluator expressionEvaluator, final SinkContext sinkContext, final PluginMetrics pluginMetrics, + final boolean topicNameInMetrics) { AwsContext awsContext = new AwsContext(kafkaProducerConfig, awsCredentialsSupplier); KeyFactory keyFactory = new KeyFactory(awsContext); prepareTopicAndSchema(kafkaProducerConfig); @@ -59,9 +62,10 @@ public KafkaCustomProducer createProducer(final KafkaProducerConfig kafkaProduce Serializer valueSerializer = (Serializer) serializationFactory.getSerializer(dataConfig); final KafkaProducer producer = new KafkaProducer<>(properties, keyDeserializer, valueSerializer); final DLQSink dlqSink = new DLQSink(pluginFactory, kafkaProducerConfig, pluginSetting); + final KafkaTopicProducerMetrics topicMetrics = new KafkaTopicProducerMetrics(topic.getName(), pluginMetrics, topicNameInMetrics); return new KafkaCustomProducer(producer, kafkaProducerConfig, dlqSink, - expressionEvaluator, Objects.nonNull(sinkContext) ? sinkContext.getTagsTargetKey() : null); + expressionEvaluator, Objects.nonNull(sinkContext) ? sinkContext.getTagsTargetKey() : null, topicMetrics); } private void prepareTopicAndSchema(final KafkaProducerConfig kafkaProducerConfig) { checkTopicCreationCriteriaAndCreateTopic(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 5cf46ea71a..4303b9fb07 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 @@ -7,6 +7,7 @@ import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; 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.configuration.PluginSetting; @@ -60,6 +61,8 @@ public class KafkaSink extends AbstractSink> { private final PluginSetting pluginSetting; + private final PluginMetrics pluginMetrics; + private final ExpressionEvaluator expressionEvaluator; private final Lock reentrantLock; @@ -69,10 +72,11 @@ public class KafkaSink extends AbstractSink> { @DataPrepperPluginConstructor public KafkaSink(final PluginSetting pluginSetting, final KafkaSinkConfig kafkaSinkConfig, final PluginFactory pluginFactory, - final ExpressionEvaluator expressionEvaluator, final SinkContext sinkContext, + final PluginMetrics pluginMetrics, final ExpressionEvaluator expressionEvaluator, final SinkContext sinkContext, AwsCredentialsSupplier awsCredentialsSupplier) { super(pluginSetting); this.pluginSetting = pluginSetting; + this.pluginMetrics = pluginMetrics; this.kafkaSinkConfig = kafkaSinkConfig; this.pluginFactory = pluginFactory; this.expressionEvaluator = expressionEvaluator; @@ -156,7 +160,7 @@ private void checkTopicCreationCriteriaAndCreateTopic() { public KafkaCustomProducer createProducer() { // TODO: Add the DLQSink here. new DLQSink(pluginFactory, kafkaSinkConfig, pluginSetting) - return kafkaCustomProducerFactory.createProducer(kafkaSinkConfig, pluginFactory, pluginSetting, expressionEvaluator, sinkContext); + return kafkaCustomProducerFactory.createProducer(kafkaSinkConfig, pluginFactory, pluginSetting, expressionEvaluator, sinkContext, pluginMetrics, 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 76219bf951..5d625439cf 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 @@ -40,7 +40,7 @@ import org.opensearch.dataprepper.plugins.kafka.extension.KafkaClusterConfigSupplier; import org.opensearch.dataprepper.plugins.kafka.util.ClientDNSLookupType; import org.opensearch.dataprepper.plugins.kafka.util.KafkaSecurityConfigurer; -import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicMetrics; +import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicConsumerMetrics; import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -109,7 +109,7 @@ public void start(Buffer> buffer) { KafkaSecurityConfigurer.setAuthProperties(authProperties, sourceConfig, LOG); sourceConfig.getTopics().forEach(topic -> { consumerGroupID = topic.getGroupId(); - KafkaTopicMetrics topicMetrics = new KafkaTopicMetrics(topic.getName(), pluginMetrics); + KafkaTopicConsumerMetrics topicMetrics = new KafkaTopicConsumerMetrics(topic.getName(), pluginMetrics, true); Properties consumerProperties = getConsumerProperties(topic, authProperties); MessageFormat schema = MessageFormat.getByMessageFormatByName(schemaType); try { diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicMetrics.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetrics.java similarity index 78% rename from data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicMetrics.java rename to data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetrics.java index df4b22a61f..aaa81b39b5 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicMetrics.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetrics.java @@ -16,7 +16,7 @@ import java.util.Map; import java.util.HashMap; -public class KafkaTopicMetrics { +public class KafkaTopicConsumerMetrics { static final String NUMBER_OF_POSITIVE_ACKNOWLEDGEMENTS = "numberOfPositiveAcknowledgements"; static final String NUMBER_OF_NEGATIVE_ACKNOWLEDGEMENTS = "numberOfNegativeAcknowledgements"; static final String NUMBER_OF_RECORDS_FAILED_TO_PARSE = "numberOfRecordsFailedToParse"; @@ -42,39 +42,40 @@ public class KafkaTopicMetrics { private final Counter numberOfRecordsConsumed; private final Counter numberOfBytesConsumed; - public KafkaTopicMetrics(final String topicName, final PluginMetrics pluginMetrics) { + public KafkaTopicConsumerMetrics(final String topicName, final PluginMetrics pluginMetrics, + final boolean topicNameInMetrics) { this.pluginMetrics = pluginMetrics; this.topicName = topicName; this.updateTime = Instant.now().getEpochSecond(); this.metricValues = new HashMap<>(); - initializeMetricNamesMap(); - this.numberOfRecordsConsumed = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_RECORDS_CONSUMED)); - this.numberOfBytesConsumed = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_BYTES_CONSUMED)); - this.numberOfRecordsCommitted = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_RECORDS_COMMITTED)); - this.numberOfRecordsFailedToParse = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_RECORDS_FAILED_TO_PARSE)); - this.numberOfDeserializationErrors = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_DESERIALIZATION_ERRORS)); - this.numberOfBufferSizeOverflows = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_BUFFER_SIZE_OVERFLOWS)); - this.numberOfPollAuthErrors = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_POLL_AUTH_ERRORS)); - this.numberOfPositiveAcknowledgements = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_POSITIVE_ACKNOWLEDGEMENTS)); - this.numberOfNegativeAcknowledgements = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_NEGATIVE_ACKNOWLEDGEMENTS)); - } - - private void initializeMetricNamesMap() { + initializeMetricNamesMap(topicNameInMetrics); + this.numberOfRecordsConsumed = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_RECORDS_CONSUMED, topicNameInMetrics)); + this.numberOfBytesConsumed = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_BYTES_CONSUMED, topicNameInMetrics)); + this.numberOfRecordsCommitted = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_RECORDS_COMMITTED, topicNameInMetrics)); + this.numberOfRecordsFailedToParse = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_RECORDS_FAILED_TO_PARSE, topicNameInMetrics)); + this.numberOfDeserializationErrors = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_DESERIALIZATION_ERRORS, topicNameInMetrics)); + this.numberOfBufferSizeOverflows = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_BUFFER_SIZE_OVERFLOWS, topicNameInMetrics)); + this.numberOfPollAuthErrors = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_POLL_AUTH_ERRORS, topicNameInMetrics)); + this.numberOfPositiveAcknowledgements = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_POSITIVE_ACKNOWLEDGEMENTS, topicNameInMetrics)); + this.numberOfNegativeAcknowledgements = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_NEGATIVE_ACKNOWLEDGEMENTS, topicNameInMetrics)); + } + + private void initializeMetricNamesMap(final boolean topicNameInMetrics) { this.metricsNameMap = new HashMap<>(); - metricsNameMap.put("bytes-consumed-total", "bytesConsumedTotal"); - metricsNameMap.put("records-consumed-total", "recordsConsumedTotal"); - metricsNameMap.put("bytes-consumed-rate", "bytesConsumedRate"); - metricsNameMap.put("records-consumed-rate", "recordsConsumedRate"); - metricsNameMap.put("records-lag-max", "recordsLagMax"); - metricsNameMap.put("records-lead-min", "recordsLeadMin"); - metricsNameMap.put("commit-rate", "commitRate"); - metricsNameMap.put("join-rate", "joinRate"); - metricsNameMap.put("incoming-byte-rate", "incomingByteRate"); - metricsNameMap.put("outgoing-byte-rate", "outgoingByteRate"); - metricsNameMap.put("assigned-partitions", "numberOfNonConsumers"); - metricsNameMap.forEach((metricName, camelCaseName) -> { + this.metricsNameMap.put("bytes-consumed-total", "bytesConsumedTotal"); + this.metricsNameMap.put("records-consumed-total", "recordsConsumedTotal"); + this.metricsNameMap.put("bytes-consumed-rate", "bytesConsumedRate"); + this.metricsNameMap.put("records-consumed-rate", "recordsConsumedRate"); + this.metricsNameMap.put("records-lag-max", "recordsLagMax"); + this.metricsNameMap.put("records-lead-min", "recordsLeadMin"); + this.metricsNameMap.put("commit-rate", "commitRate"); + this.metricsNameMap.put("join-rate", "joinRate"); + this.metricsNameMap.put("incoming-byte-rate", "incomingByteRate"); + this.metricsNameMap.put("outgoing-byte-rate", "outgoingByteRate"); + this.metricsNameMap.put("assigned-partitions", "numberOfNonConsumers"); + this.metricsNameMap.forEach((metricName, camelCaseName) -> { if (metricName.equals("records-lag-max")) { - pluginMetrics.gauge(getTopicMetricName(camelCaseName), metricValues, metricValues -> { + pluginMetrics.gauge(getTopicMetricName(camelCaseName, topicNameInMetrics), metricValues, metricValues -> { double max = 0.0; for (Map.Entry> entry : metricValues.entrySet()) { Map consumerMetrics = entry.getValue(); @@ -85,7 +86,7 @@ private void initializeMetricNamesMap() { return max; }); } else if (metricName.equals("records-lead-min")) { - pluginMetrics.gauge(getTopicMetricName(camelCaseName), metricValues, metricValues -> { + pluginMetrics.gauge(getTopicMetricName(camelCaseName, topicNameInMetrics), metricValues, metricValues -> { double min = Double.MAX_VALUE; for (Map.Entry> entry : metricValues.entrySet()) { Map consumerMetrics = entry.getValue(); @@ -96,7 +97,7 @@ private void initializeMetricNamesMap() { return min; }); } else if (!metricName.contains("-total")) { - pluginMetrics.gauge(getTopicMetricName(camelCaseName), metricValues, metricValues -> { + pluginMetrics.gauge(getTopicMetricName(camelCaseName, topicNameInMetrics), metricValues, metricValues -> { double sum = 0; for (Map.Entry> entry : metricValues.entrySet()) { Map consumerMetrics = entry.getValue(); @@ -154,8 +155,12 @@ public Counter getNumberOfPositiveAcknowledgements() { return numberOfPositiveAcknowledgements; } - private String getTopicMetricName(final String metricName) { - return "topic."+topicName+"."+metricName; + private String getTopicMetricName(final String metricName, final boolean topicNameInMetrics) { + if (topicNameInMetrics) { + return "topic." + topicName + "." + metricName; + } else { + return metricName; + } } private String getCamelCaseName(final String name) { @@ -172,7 +177,6 @@ Map> getMetricValues() { public void update(final KafkaConsumer consumer) { Map consumerMetrics = metricValues.get(consumer); - Map metrics = consumer.metrics(); for (Map.Entry entry : metrics.entrySet()) { MetricName metric = entry.getKey(); diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicProducerMetrics.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicProducerMetrics.java new file mode 100644 index 0000000000..525b85d0d5 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicProducerMetrics.java @@ -0,0 +1,133 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.util; + +import io.micrometer.core.instrument.Counter; +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.HashMap; +import java.util.Map; +import java.util.Objects; + +public class KafkaTopicProducerMetrics { + static final String NUMBER_OF_RECORDS_SENT = "numberOfRecordsSent"; + static final String NUMBER_OF_BYTES_SENT = "numberOfBytesSent"; + static final String RECORD_SEND_TOTAL = "record-send-total"; + static final String BYTE_TOTAL = "byte-total"; + static final String TOPIC = "topic"; + static final String METRIC_SUFFIX_TOTAL = "-total"; + static final String BYTE_RATE = "byte-rate"; + static final String RECORD_SEND_RATE = "record-send-rate"; + static final String BYTE_SEND_RATE = "byteSendRate"; + static final String BYTE_SEND_TOTAL = "byteSendTotal"; + static final String RECORD_SEND_RATE_MAP_VALUE = "recordSendRate"; + static final String RECORD_SEND_TOTAL_MAP_VALUE = "recordSendTotal"; + + private final String topicName; + private Map metricsNameMap; + private Map> metricValues; + private final PluginMetrics pluginMetrics; + private final Counter numberOfRecordsSent; + private final Counter numberOfBytesSent; + + public KafkaTopicProducerMetrics(final String topicName, final PluginMetrics pluginMetrics, + final boolean topicNameInMetrics) { + this.pluginMetrics = pluginMetrics; + this.topicName = topicName; + this.metricValues = new HashMap<>(); + initializeMetricNamesMap(topicNameInMetrics); + this.numberOfRecordsSent = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_RECORDS_SENT, topicNameInMetrics)); + this.numberOfBytesSent = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_BYTES_SENT, topicNameInMetrics)); + } + + private void initializeMetricNamesMap(final boolean topicNameInMetrics) { + this.metricsNameMap = new HashMap<>(); + metricsNameMap.put(BYTE_RATE, BYTE_SEND_RATE); + metricsNameMap.put(BYTE_TOTAL, BYTE_SEND_TOTAL); + metricsNameMap.put(RECORD_SEND_RATE, RECORD_SEND_RATE_MAP_VALUE); + metricsNameMap.put(RECORD_SEND_TOTAL, RECORD_SEND_TOTAL_MAP_VALUE); + + metricsNameMap.forEach((metricName, camelCaseName) -> { + if (!metricName.contains(METRIC_SUFFIX_TOTAL)) { + pluginMetrics.gauge(getTopicMetricName(camelCaseName, topicNameInMetrics), metricValues, metricValues -> { + double sum = 0; + for (Map.Entry> entry : metricValues.entrySet()) { + Map producerMetrics = entry.getValue(); + synchronized(producerMetrics) { + sum += producerMetrics.get(metricName); + } + } + return sum; + }); + } + }); + } + + public void register(final KafkaProducer producer) { + metricValues.put(producer, new HashMap<>()); + final Map producerMetrics = metricValues.get(producer); + metricsNameMap.forEach((k, name) -> { + producerMetrics.put(k, 0.0); + }); + } + + Counter getNumberOfRecordsSent() { + return numberOfRecordsSent; + } + + Counter getNumberOfBytesSent() { + return numberOfBytesSent; + } + + private String getTopicMetricName(final String metricName, final boolean topicNameInMetrics) { + if (topicNameInMetrics) { + return "topic." + topicName + "." + metricName; + } else { + return metricName; + } + } + + Map> getMetricValues() { + return metricValues; + } + + public void update(final KafkaProducer producer) { + Map producerMetrics = metricValues.get(producer); + + Map metrics = producer.metrics(); + for (Map.Entry entry : metrics.entrySet()) { + MetricName metric = entry.getKey(); + Metric value = entry.getValue(); + String metricName = metric.name(); + if (Objects.nonNull(metricsNameMap.get(metricName))) { + // producer metrics are emitted at topic level + if (!metric.tags().containsKey(TOPIC)) { + continue; + } + + double newValue = (Double)value.metricValue(); + if (metricName.equals(RECORD_SEND_TOTAL)) { + synchronized(producerMetrics) { + double prevValue = producerMetrics.get(metricName); + numberOfRecordsSent.increment(newValue - prevValue); + } + } else if (metricName.equals(BYTE_TOTAL)) { + synchronized(producerMetrics) { + double prevValue = producerMetrics.get(metricName); + numberOfBytesSent.increment(newValue - prevValue); + } + } + + synchronized(producerMetrics) { + producerMetrics.put(metricName, newValue); + } + } + } + } +} 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 06c44d3599..59c7641b76 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 @@ -50,6 +50,7 @@ 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.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.isNull; @@ -119,7 +120,7 @@ public KafkaBuffer> createObjectUnderTest() { final MockedConstruction producerFactoryMock = mockConstruction(KafkaCustomProducerFactory.class, (mock, context) -> { producerFactory = mock; - when(producerFactory.createProducer(any() ,any(), any(), isNull(), isNull())).thenReturn(producer); + when(producerFactory.createProducer(any() ,any(), any(), isNull(), isNull(), any(), anyBoolean())).thenReturn(producer); }); final MockedConstruction blockingBufferMock = mockConstruction(BlockingBuffer.class, (mock, context) -> { 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 a5913042de..d707cbf82d 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 @@ -35,7 +35,7 @@ import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConsumerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaConsumerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaKeyMode; -import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicMetrics; +import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicConsumerMetrics; import org.opensearch.dataprepper.plugins.kafka.util.MessageFormat; import java.time.Duration; @@ -79,7 +79,7 @@ public class KafkaCustomConsumerTest { private TopicConsumerConfig topicConfig; @Mock - private KafkaTopicMetrics topicMetrics; + private KafkaTopicConsumerMetrics topicMetrics; private KafkaCustomConsumer consumer; @@ -110,7 +110,7 @@ public class KafkaCustomConsumerTest { public void setUp() { delayTime = Duration.ofMillis(10); kafkaConsumer = mock(KafkaConsumer.class); - topicMetrics = mock(KafkaTopicMetrics.class); + topicMetrics = mock(KafkaTopicConsumerMetrics.class); counter = mock(Counter.class); posCounter = mock(Counter.class); negCounter = mock(Counter.class); diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerTest.java index 8b23975401..2a03e5d9db 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerTest.java @@ -10,13 +10,14 @@ import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; import org.apache.avro.Schema; import org.apache.avro.generic.GenericRecord; -import org.apache.kafka.clients.producer.MockProducer; -import org.apache.kafka.common.serialization.StringSerializer; -import org.apache.kafka.connect.json.JsonSerializer; +import org.apache.kafka.clients.producer.Callback; +import org.apache.kafka.clients.producer.KafkaProducer; +import org.apache.kafka.clients.producer.ProducerRecord; import org.junit.jupiter.api.Assertions; 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.mockito.junit.jupiter.MockitoSettings; @@ -29,6 +30,7 @@ import org.opensearch.dataprepper.plugins.kafka.configuration.TopicProducerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; import org.opensearch.dataprepper.plugins.kafka.sink.DLQSink; +import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicProducerMetrics; import java.io.IOException; import java.lang.reflect.InvocationTargetException; @@ -38,6 +40,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.verify; @@ -53,6 +56,9 @@ public class KafkaCustomProducerTest { @Mock private KafkaProducerConfig kafkaSinkConfig; + @Mock + private KafkaTopicProducerMetrics kafkaTopicProducerMetrics; + private Record record; KafkaCustomProducer sinkProducer; @@ -79,20 +85,23 @@ record = new Record<>(event); @Test public void producePlainTextRecordsTest() throws ExecutionException, InterruptedException { when(kafkaSinkConfig.getSerdeFormat()).thenReturn("plaintext"); - MockProducer mockProducer = new MockProducer<>(true, new StringSerializer(), new StringSerializer()); - producer = new KafkaCustomProducer(mockProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), null); + KafkaProducer kafkaProducer = mock(KafkaProducer.class); + producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), null, kafkaTopicProducerMetrics); sinkProducer = spy(producer); sinkProducer.produceRecords(record); verify(sinkProducer).produceRecords(record); - assertEquals(1, mockProducer.history().size()); + final ArgumentCaptor recordArgumentCaptor = ArgumentCaptor.forClass(ProducerRecord.class); + verify(kafkaProducer).send(recordArgumentCaptor.capture(), any(Callback.class)); + assertEquals(recordArgumentCaptor.getValue().topic(), kafkaSinkConfig.getTopic().getName()); + assertEquals(recordArgumentCaptor.getValue().value(), record.getData().toJsonString()); } @Test public void produceJsonRecordsTest() throws RestClientException, IOException { when(kafkaSinkConfig.getSerdeFormat()).thenReturn("JSON"); - MockProducer mockProducer = new MockProducer<>(true, new StringSerializer(), new JsonSerializer()); - producer = new KafkaCustomProducer(mockProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), null); + KafkaProducer kafkaProducer = mock(KafkaProducer.class); + producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), null, kafkaTopicProducerMetrics); SchemaMetadata schemaMetadata = mock(SchemaMetadata.class); String jsonSchema = "{\n" + " \"$schema\": \"http://json-schema.org/draft-07/schema#\",\n" + @@ -108,14 +117,17 @@ public void produceJsonRecordsTest() throws RestClientException, IOException { sinkProducer = spy(producer); sinkProducer.produceRecords(record); verify(sinkProducer).produceRecords(record); - assertEquals(1, mockProducer.history().size()); + final ArgumentCaptor recordArgumentCaptor = ArgumentCaptor.forClass(ProducerRecord.class); + verify(kafkaProducer).send(recordArgumentCaptor.capture(), any(Callback.class)); + assertEquals(recordArgumentCaptor.getValue().topic(), kafkaSinkConfig.getTopic().getName()); + assertEquals(recordArgumentCaptor.getValue().value(), record.getData().getJsonNode()); } @Test public void produceAvroRecordsTest() throws Exception { when(kafkaSinkConfig.getSerdeFormat()).thenReturn("AVRO"); - MockProducer mockProducer = new MockProducer<>(true, new StringSerializer(), new StringSerializer()); - producer = new KafkaCustomProducer(mockProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), null); + KafkaProducer kafkaProducer = mock(KafkaProducer.class); + producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), null, kafkaTopicProducerMetrics); SchemaMetadata schemaMetadata = mock(SchemaMetadata.class); String avroSchema = "{\"type\":\"record\",\"name\":\"MyMessage\",\"fields\":[{\"name\":\"message\",\"type\":\"string\"}]}"; when(schemaMetadata.getSchema()).thenReturn(avroSchema); @@ -127,7 +139,8 @@ public void produceAvroRecordsTest() throws Exception { @Test public void testGetGenericRecord() throws NoSuchMethodException, InvocationTargetException, IllegalAccessException { - producer = new KafkaCustomProducer(new MockProducer(), kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), null); + KafkaProducer kafkaProducer = mock(KafkaProducer.class); + producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), null, kafkaTopicProducerMetrics); final Schema schema = createMockSchema(); Method privateMethod = KafkaCustomProducer.class.getDeclaredMethod("getGenericRecord", Event.class, Schema.class); privateMethod.setAccessible(true); @@ -144,8 +157,8 @@ private Schema createMockSchema() { @Test public void validateSchema() throws IOException, ProcessingException { when(kafkaSinkConfig.getSerdeFormat()).thenReturn("avro"); - MockProducer mockProducer = new MockProducer<>(true, new StringSerializer(), new StringSerializer()); - producer = new KafkaCustomProducer(mockProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), null); + KafkaProducer kafkaProducer = mock(KafkaProducer.class); + producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), null, kafkaTopicProducerMetrics); String jsonSchema = "{\"type\": \"object\",\"properties\": {\"Year\": {\"type\": \"string\"},\"Age\": {\"type\": \"string\"},\"Ethnic\": {\"type\":\"string\",\"default\": null}}}"; String jsonSchema2 = "{\"type\": \"object\",\"properties\": {\"Year\": {\"type\": \"string\"},\"Age\": {\"type\": \"string\"},\"Ethnic\": {\"type\":\"string\",\"default\": null}}}"; assertTrue(producer.validateSchema(jsonSchema, jsonSchema2)); 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 d61eb22df0..1984e40516 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 @@ -20,6 +20,7 @@ import org.mockito.quality.Strictness; 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.event.Event; import org.opensearch.dataprepper.model.event.JacksonEvent; @@ -70,6 +71,9 @@ public class KafkaSinkTest { @Mock PluginSetting pluginSetting; + @Mock + PluginMetrics pluginMetrics; + @Mock FutureTask futureTask; @@ -114,7 +118,7 @@ void setUp() throws Exception { when(pluginSetting.getPipelineName()).thenReturn("Kafka-sink"); event = JacksonEvent.fromMessage(UUID.randomUUID().toString()); when(sinkContext.getTagsTargetKey()).thenReturn("tag"); - kafkaSink = new KafkaSink(pluginSetting, kafkaSinkConfig, pluginFactoryMock, mock(ExpressionEvaluator.class), sinkContext, awsCredentialsSupplier); + kafkaSink = new KafkaSink(pluginSetting, kafkaSinkConfig, pluginFactoryMock, pluginMetrics, mock(ExpressionEvaluator.class), sinkContext, awsCredentialsSupplier); spySink = spy(kafkaSink); executorsMockedStatic = mockStatic(Executors.class); props = new Properties(); diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicMetricsTests.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetricsTests.java similarity index 97% rename from data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicMetricsTests.java rename to data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetricsTests.java index ea31c216db..cd608ff79f 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicMetricsTests.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetricsTests.java @@ -20,7 +20,6 @@ import org.opensearch.dataprepper.metrics.PluginMetrics; import org.apache.kafka.common.Metric; import org.apache.kafka.common.MetricName; -import org.apache.kafka.common.Metric; import org.apache.kafka.clients.consumer.KafkaConsumer; import org.apache.commons.lang3.RandomStringUtils; @@ -34,7 +33,7 @@ import java.util.function.ToDoubleFunction; @ExtendWith(MockitoExtension.class) -public class KafkaTopicMetricsTests { +public class KafkaTopicConsumerMetricsTests { public final class KafkaTestMetric implements Metric { private final Double value; private final MetricName name; @@ -64,7 +63,7 @@ public Object metricValue() { private Random random; - private KafkaTopicMetrics topicMetrics; + private KafkaTopicConsumerMetrics topicMetrics; private double bytesConsumed; private double recordsConsumed; @@ -133,8 +132,8 @@ void setUp() { }).when(recordsConsumedCounter).increment(any(Double.class)); } - public KafkaTopicMetrics createObjectUnderTest() { - return new KafkaTopicMetrics(topicName, pluginMetrics); + public KafkaTopicConsumerMetrics createObjectUnderTest() { + return new KafkaTopicConsumerMetrics(topicName, pluginMetrics, true); } private KafkaTestMetric getMetric(final String name, final double value, Map tags) { diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicProducerMetricsTests.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicProducerMetricsTests.java new file mode 100644 index 0000000000..6abba8b198 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicProducerMetricsTests.java @@ -0,0 +1,190 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.util; + +import io.micrometer.core.instrument.Counter; +import org.apache.commons.lang3.RandomStringUtils; +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.extension.ExtendWith; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +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.function.ToDoubleFunction; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +public class KafkaTopicProducerMetricsTests { + public static final class KafkaTestMetric implements Metric { + private final Double value; + private final MetricName name; + + public KafkaTestMetric(final double value, final MetricName name) { + this.value = value; + this.name = name; + } + + @Override + public MetricName metricName() { + return name; + } + + @Override + public Object metricValue() { + return value; + } + } + + private String topicName; + + @Mock + private PluginMetrics pluginMetrics; + + private Map pluginMetricsMap; + + private Random random; + + private KafkaTopicProducerMetrics topicMetrics; + + private double bytesSent; + private double recordsSent; + private double bytesSendRate; + private double recordSendRate; + + @Mock + private Counter byteSendCounter; + + @Mock + private Counter recordSendCounter; + private double byteSendCount; + private double recordSendCount; + + @BeforeEach + void setUp() { + topicName = RandomStringUtils.randomAlphabetic(8); + bytesSent = 0.0; + recordsSent = 0.0; + bytesSendRate = 0.0; + recordSendRate = 0.0; + + byteSendCount = 0.0; + recordSendCount = 0.0; + + random = new Random(); + pluginMetrics = mock(PluginMetrics.class); + pluginMetricsMap = new HashMap<>(); + doAnswer((i) -> { + ToDoubleFunction f = i.getArgument(2); + Object arg = i.getArgument(1); + String name = i.getArgument(0); + pluginMetricsMap.put(name, f); + return f.applyAsDouble(arg); + }).when(pluginMetrics).gauge(any(String.class), any(Object.class), any()); + byteSendCounter = mock(Counter.class); + recordSendCounter = mock(Counter.class); + + doAnswer((i) -> { + String arg = i.getArgument(0); + if (arg.contains("Bytes")) { + return byteSendCounter; + } else { + return recordSendCounter; + } + }).when(pluginMetrics).counter(any(String.class)); + doAnswer((i) -> { + byteSendCount += (double)i.getArgument(0); + return null; + }).when(byteSendCounter).increment(any(Double.class)); + doAnswer((i) -> { + recordSendCount += (double)i.getArgument(0); + return null; + }).when(recordSendCounter).increment(any(Double.class)); + } + + public KafkaTopicProducerMetrics createObjectUnderTest() { + return new KafkaTopicProducerMetrics(topicName, pluginMetrics, true); + } + + private KafkaTestMetric getMetric(final String name, final double value, Map tags) { + MetricName metricName = new MetricName(name, "group", "metric", tags); + return new KafkaTestMetric(value, metricName); + } + + + private void populateKafkaMetrics(Map metrics, double numAssignedPartitions) { + int tmpBytesProduced = random.nextInt() % 100 + 1; + if (tmpBytesProduced < 0) { + tmpBytesProduced = -tmpBytesProduced; + } + bytesSent += tmpBytesProduced; + int tmpRecordsProduced = random.nextInt() % 10 + 1; + if (tmpRecordsProduced < 0) { + tmpRecordsProduced = -tmpRecordsProduced; + } + recordsSent += tmpRecordsProduced; + + double tmpBytesProducedRate = random.nextDouble()*100; + bytesSendRate += tmpBytesProducedRate; + + double tmpRecordsProducedRate = random.nextDouble()*10; + recordSendRate += tmpRecordsProducedRate; + + Map metricsMap = new HashMap<>(); + metricsMap.put("byte-total", (double)tmpBytesProduced); + metricsMap.put("record-send-total", (double)tmpRecordsProduced); + metricsMap.put("byte-rate", tmpBytesProducedRate); + metricsMap.put("record-send-rate", tmpRecordsProducedRate); + + metricsMap.forEach((name, value) -> { + Map tags = Map.of("topic", topicName); + KafkaTestMetric metric = getMetric(name, value, tags); + metrics.put(metric.metricName(), metric); + }); + } + + @ParameterizedTest + @ValueSource(ints = {1, 5, 10}) + public void KafkaTopicMetricTest_checkMetricUpdates(int numProducers) { + topicMetrics = createObjectUnderTest(); + for (int i = 0; i < numProducers; i++) { + KafkaProducer kafkaProducer = mock(KafkaProducer.class); + topicMetrics.register(kafkaProducer); + Map metrics = new HashMap<>(); + when(kafkaProducer.metrics()).thenReturn(metrics); + populateKafkaMetrics(metrics, (i %2 == 1) ? 0.0 : 1.0); + topicMetrics.update(kafkaProducer); + } + when(recordSendCounter.count()).thenReturn(recordSendCount); + when(byteSendCounter.count()).thenReturn(byteSendCount); + assertThat(topicMetrics.getNumberOfRecordsSent().count(), equalTo(recordsSent)); + assertThat(topicMetrics.getNumberOfBytesSent().count(), equalTo(bytesSent)); + pluginMetricsMap.forEach((k, v) -> { + double result = v.applyAsDouble(topicMetrics.getMetricValues()); + if (k.contains("byteSendRate")) { + assertEquals(result, bytesSendRate, 0.01d); + } else if (k.contains("recordSendRate")) { + assertEquals(result, recordSendRate, 0.01d); + } else { + assertThat(result, equalTo(k+": Unknown Metric")); + } + }); + } +} From 3cc2a935f9ec04796856282ed5413b6609609c2b Mon Sep 17 00:00:00 2001 From: David Venable Date: Wed, 8 Nov 2023 06:38:49 -0800 Subject: [PATCH 36/53] Support ISM correctly when using composable index templates (#3590) Correctly support custom properties in composable index templates in the OpenSearch sink. This resolves #3506. Signed-off-by: David Venable --- .../opensearch/DeclaredOpenSearchVersion.java | 2 +- .../sink/opensearch/OpenSearchSinkIT.java | 283 ++++++++++-------- .../index/ComposableIndexTemplate.java | 9 +- .../index/ComposableTemplateAPIWrapper.java | 61 ++-- .../opensearch/index/LegacyIndexTemplate.java | 2 +- .../ComposableIndexTemplateStrategyTest.java | 89 ++++++ .../ComposableTemplateAPIWrapperTest.java | 164 ++++++---- 7 files changed, 396 insertions(+), 214 deletions(-) diff --git a/data-prepper-plugins/opensearch/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/opensearch/DeclaredOpenSearchVersion.java b/data-prepper-plugins/opensearch/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/opensearch/DeclaredOpenSearchVersion.java index 8f62098746..8112f075c3 100644 --- a/data-prepper-plugins/opensearch/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/opensearch/DeclaredOpenSearchVersion.java +++ b/data-prepper-plugins/opensearch/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/opensearch/DeclaredOpenSearchVersion.java @@ -33,7 +33,7 @@ private DeclaredOpenSearchVersion(final Distribution distribution, final String } static DeclaredOpenSearchVersion parse(final String versionString) { - if(versionString == null) { + if(versionString == null || versionString.isEmpty()) { return DEFAULT; } 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 fbe17763a0..97c80fe3bc 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 @@ -8,11 +8,11 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.JsonNodeType; import io.micrometer.core.instrument.Measurement; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.RandomStringUtils; import org.apache.http.util.EntityUtils; -import org.hamcrest.MatcherAssert; import org.junit.Assert; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -85,6 +85,7 @@ import static org.hamcrest.CoreMatchers.hasItems; import static org.hamcrest.CoreMatchers.notNullValue; import static org.hamcrest.CoreMatchers.nullValue; +import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.closeTo; import static org.junit.jupiter.params.provider.Arguments.arguments; import static org.mockito.ArgumentMatchers.any; @@ -173,17 +174,17 @@ public void testInstantiateSinkRawSpanDefault() throws IOException { final String indexAlias = IndexConstants.TYPE_TO_DEFAULT_ALIAS.get(IndexType.TRACE_ANALYTICS_RAW); Request request = new Request(HttpMethod.HEAD, indexAlias); Response response = client.performRequest(request); - MatcherAssert.assertThat(response.getStatusLine().getStatusCode(), equalTo(SC_OK)); + assertThat(response.getStatusLine().getStatusCode(), equalTo(SC_OK)); final String index = String.format("%s-000001", indexAlias); final Map mappings = getIndexMappings(index); - MatcherAssert.assertThat(mappings, notNullValue()); - MatcherAssert.assertThat((boolean) mappings.get("date_detection"), equalTo(false)); + assertThat(mappings, notNullValue()); + assertThat((boolean) mappings.get("date_detection"), equalTo(false)); sink.shutdown(); if (isOSBundle()) { // Check managed index await().atMost(1, TimeUnit.SECONDS).untilAsserted(() -> { - MatcherAssert.assertThat(getIndexPolicyId(index), equalTo(IndexConstants.RAW_ISM_POLICY)); + assertThat(getIndexPolicyId(index), equalTo(IndexConstants.RAW_ISM_POLICY)); } ); } @@ -192,7 +193,7 @@ public void testInstantiateSinkRawSpanDefault() throws IOException { request = new Request(HttpMethod.POST, String.format("%s/_rollover", indexAlias)); request.setJsonEntity("{ \"conditions\" : { } }\n"); response = client.performRequest(request); - MatcherAssert.assertThat(response.getStatusLine().getStatusCode(), equalTo(SC_OK)); + assertThat(response.getStatusLine().getStatusCode(), equalTo(SC_OK)); // Instantiate sink again sink = createObjectUnderTest(pluginSetting, true); @@ -200,12 +201,12 @@ public void testInstantiateSinkRawSpanDefault() throws IOException { final String rolloverIndexName = String.format("%s-000002", indexAlias); request = new Request(HttpMethod.GET, rolloverIndexName + "/_alias"); response = client.performRequest(request); - MatcherAssert.assertThat(checkIsWriteIndex(EntityUtils.toString(response.getEntity()), indexAlias, rolloverIndexName), equalTo(true)); + assertThat(checkIsWriteIndex(EntityUtils.toString(response.getEntity()), indexAlias, rolloverIndexName), equalTo(true)); sink.shutdown(); if (isOSBundle()) { // Check managed index - MatcherAssert.assertThat(getIndexPolicyId(rolloverIndexName), equalTo(IndexConstants.RAW_ISM_POLICY)); + assertThat(getIndexPolicyId(rolloverIndexName), equalTo(IndexConstants.RAW_ISM_POLICY)); } } @@ -240,21 +241,21 @@ public void testOutputRawSpanDefault(final boolean estimateBulkSizeUsingCompress final String expIndexAlias = IndexConstants.TYPE_TO_DEFAULT_ALIAS.get(IndexType.TRACE_ANALYTICS_RAW); final List> retSources = getSearchResponseDocSources(expIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(2)); - MatcherAssert.assertThat(retSources, hasItems(expData1, expData2)); - MatcherAssert.assertThat(getDocumentCount(expIndexAlias, "_id", (String) expData1.get("spanId")), equalTo(Integer.valueOf(1))); + assertThat(retSources.size(), equalTo(2)); + assertThat(retSources, hasItems(expData1, expData2)); + assertThat(getDocumentCount(expIndexAlias, "_id", (String) expData1.get("spanId")), equalTo(Integer.valueOf(1))); sink.shutdown(); // Verify metrics final List bulkRequestErrors = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(OpenSearchSink.BULKREQUEST_ERRORS).toString()); - MatcherAssert.assertThat(bulkRequestErrors.size(), equalTo(1)); + assertThat(bulkRequestErrors.size(), equalTo(1)); Assert.assertEquals(0.0, bulkRequestErrors.get(0).getValue(), 0); final List bulkRequestLatencies = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(OpenSearchSink.BULKREQUEST_LATENCY).toString()); - MatcherAssert.assertThat(bulkRequestLatencies.size(), equalTo(3)); + assertThat(bulkRequestLatencies.size(), equalTo(3)); // COUNT Assert.assertEquals(1.0, bulkRequestLatencies.get(0).getValue(), 0); // TOTAL_TIME @@ -264,18 +265,18 @@ public void testOutputRawSpanDefault(final boolean estimateBulkSizeUsingCompress final List documentsSuccessMeasurements = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(BulkRetryStrategy.DOCUMENTS_SUCCESS).toString()); - MatcherAssert.assertThat(documentsSuccessMeasurements.size(), equalTo(1)); - MatcherAssert.assertThat(documentsSuccessMeasurements.get(0).getValue(), closeTo(2.0, 0)); + assertThat(documentsSuccessMeasurements.size(), equalTo(1)); + assertThat(documentsSuccessMeasurements.get(0).getValue(), closeTo(2.0, 0)); final List documentsSuccessFirstAttemptMeasurements = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(BulkRetryStrategy.DOCUMENTS_SUCCESS_FIRST_ATTEMPT).toString()); - MatcherAssert.assertThat(documentsSuccessFirstAttemptMeasurements.size(), equalTo(1)); - MatcherAssert.assertThat(documentsSuccessFirstAttemptMeasurements.get(0).getValue(), closeTo(2.0, 0)); + assertThat(documentsSuccessFirstAttemptMeasurements.size(), equalTo(1)); + assertThat(documentsSuccessFirstAttemptMeasurements.get(0).getValue(), closeTo(2.0, 0)); final List documentErrorsMeasurements = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(BulkRetryStrategy.DOCUMENT_ERRORS).toString()); - MatcherAssert.assertThat(documentErrorsMeasurements.size(), equalTo(1)); - MatcherAssert.assertThat(documentErrorsMeasurements.get(0).getValue(), closeTo(0.0, 0)); + assertThat(documentErrorsMeasurements.size(), equalTo(1)); + assertThat(documentErrorsMeasurements.get(0).getValue(), closeTo(0.0, 0)); /** * Metrics: Bulk Request Size in Bytes @@ -283,11 +284,11 @@ public void testOutputRawSpanDefault(final boolean estimateBulkSizeUsingCompress final List bulkRequestSizeBytesMetrics = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(OpenSearchSink.BULKREQUEST_SIZE_BYTES).toString()); - MatcherAssert.assertThat(bulkRequestSizeBytesMetrics.size(), equalTo(3)); - MatcherAssert.assertThat(bulkRequestSizeBytesMetrics.get(0).getValue(), closeTo(1.0, 0)); + assertThat(bulkRequestSizeBytesMetrics.size(), equalTo(3)); + assertThat(bulkRequestSizeBytesMetrics.get(0).getValue(), closeTo(1.0, 0)); final double expectedBulkRequestSizeBytes = isRequestCompressionEnabled && estimateBulkSizeUsingCompression ? 773.0 : 2058.0; - MatcherAssert.assertThat(bulkRequestSizeBytesMetrics.get(1).getValue(), closeTo(expectedBulkRequestSizeBytes, 0)); - MatcherAssert.assertThat(bulkRequestSizeBytesMetrics.get(2).getValue(), closeTo(expectedBulkRequestSizeBytes, 0)); + assertThat(bulkRequestSizeBytesMetrics.get(1).getValue(), closeTo(expectedBulkRequestSizeBytes, 0)); + assertThat(bulkRequestSizeBytesMetrics.get(2).getValue(), closeTo(expectedBulkRequestSizeBytes, 0)); } @DisabledIf(value = "isES6", disabledReason = TRACE_INGESTION_TEST_DISABLED_REASON) @@ -317,11 +318,11 @@ public void testOutputRawSpanWithDLQ(final boolean estimateBulkSizeUsingCompress final StringBuilder dlqContent = new StringBuilder(); Files.lines(Paths.get(expDLQFile)).forEach(dlqContent::append); final String nonPrettyJsonString = mapper.writeValueAsString(mapper.readValue(testDoc1, JsonNode.class)); - MatcherAssert.assertThat(dlqContent.toString(), containsString(nonPrettyJsonString)); + assertThat(dlqContent.toString(), containsString(nonPrettyJsonString)); final String expIndexAlias = IndexConstants.TYPE_TO_DEFAULT_ALIAS.get(IndexType.TRACE_ANALYTICS_RAW); final List> retSources = getSearchResponseDocSources(expIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); - MatcherAssert.assertThat(retSources.get(0), equalTo(expData)); + assertThat(retSources.size(), equalTo(1)); + assertThat(retSources.get(0), equalTo(expData)); // clean up temporary directory FileUtils.deleteQuietly(tempDirectory); @@ -330,13 +331,13 @@ public void testOutputRawSpanWithDLQ(final boolean estimateBulkSizeUsingCompress final List documentsSuccessMeasurements = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(BulkRetryStrategy.DOCUMENTS_SUCCESS).toString()); - MatcherAssert.assertThat(documentsSuccessMeasurements.size(), equalTo(1)); - MatcherAssert.assertThat(documentsSuccessMeasurements.get(0).getValue(), closeTo(1.0, 0)); + assertThat(documentsSuccessMeasurements.size(), equalTo(1)); + assertThat(documentsSuccessMeasurements.get(0).getValue(), closeTo(1.0, 0)); final List documentErrorsMeasurements = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(BulkRetryStrategy.DOCUMENT_ERRORS).toString()); - MatcherAssert.assertThat(documentErrorsMeasurements.size(), equalTo(1)); - MatcherAssert.assertThat(documentErrorsMeasurements.get(0).getValue(), closeTo(1.0, 0)); + assertThat(documentErrorsMeasurements.size(), equalTo(1)); + assertThat(documentErrorsMeasurements.get(0).getValue(), closeTo(1.0, 0)); /** * Metrics: Bulk Request Size in Bytes @@ -344,11 +345,11 @@ public void testOutputRawSpanWithDLQ(final boolean estimateBulkSizeUsingCompress final List bulkRequestSizeBytesMetrics = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(OpenSearchSink.BULKREQUEST_SIZE_BYTES).toString()); - MatcherAssert.assertThat(bulkRequestSizeBytesMetrics.size(), equalTo(3)); - MatcherAssert.assertThat(bulkRequestSizeBytesMetrics.get(0).getValue(), closeTo(1.0, 0)); + assertThat(bulkRequestSizeBytesMetrics.size(), equalTo(3)); + assertThat(bulkRequestSizeBytesMetrics.get(0).getValue(), closeTo(1.0, 0)); final double expectedBulkRequestSizeBytes = isRequestCompressionEnabled && estimateBulkSizeUsingCompression ? 1066.0 : 2072.0; - MatcherAssert.assertThat(bulkRequestSizeBytesMetrics.get(1).getValue(), closeTo(expectedBulkRequestSizeBytes, 0)); - MatcherAssert.assertThat(bulkRequestSizeBytesMetrics.get(2).getValue(), closeTo(expectedBulkRequestSizeBytes, 0)); + assertThat(bulkRequestSizeBytesMetrics.get(1).getValue(), closeTo(expectedBulkRequestSizeBytes, 0)); + assertThat(bulkRequestSizeBytesMetrics.get(2).getValue(), closeTo(expectedBulkRequestSizeBytes, 0)); } @@ -360,15 +361,15 @@ public void testInstantiateSinkServiceMapDefault() throws IOException { final String indexAlias = IndexConstants.TYPE_TO_DEFAULT_ALIAS.get(IndexType.TRACE_ANALYTICS_SERVICE_MAP); final Request request = new Request(HttpMethod.HEAD, indexAlias); final Response response = client.performRequest(request); - MatcherAssert.assertThat(response.getStatusLine().getStatusCode(), equalTo(SC_OK)); + assertThat(response.getStatusLine().getStatusCode(), equalTo(SC_OK)); final Map mappings = getIndexMappings(indexAlias); - MatcherAssert.assertThat(mappings, notNullValue()); - MatcherAssert.assertThat((boolean) mappings.get("date_detection"), equalTo(false)); + assertThat(mappings, notNullValue()); + assertThat((boolean) mappings.get("date_detection"), equalTo(false)); sink.shutdown(); if (isOSBundle()) { // Check managed index - MatcherAssert.assertThat(getIndexPolicyId(indexAlias), nullValue()); + assertThat(getIndexPolicyId(indexAlias), nullValue()); } } @@ -387,18 +388,18 @@ public void testOutputServiceMapDefault(final boolean estimateBulkSizeUsingCompr sink.output(testRecords); final String expIndexAlias = IndexConstants.TYPE_TO_DEFAULT_ALIAS.get(IndexType.TRACE_ANALYTICS_SERVICE_MAP); final List> retSources = getSearchResponseDocSources(expIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); - MatcherAssert.assertThat(retSources.get(0), equalTo(expData)); - MatcherAssert.assertThat(getDocumentCount(expIndexAlias, "_id", (String) expData.get("hashId")), equalTo(Integer.valueOf(1))); + assertThat(retSources.size(), equalTo(1)); + assertThat(retSources.get(0), equalTo(expData)); + assertThat(getDocumentCount(expIndexAlias, "_id", (String) expData.get("hashId")), equalTo(Integer.valueOf(1))); sink.shutdown(); // verify metrics final List bulkRequestLatencies = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(OpenSearchSink.BULKREQUEST_LATENCY).toString()); - MatcherAssert.assertThat(bulkRequestLatencies.size(), equalTo(3)); + assertThat(bulkRequestLatencies.size(), equalTo(3)); // COUNT - MatcherAssert.assertThat(bulkRequestLatencies.get(0).getValue(), closeTo(1.0, 0)); + assertThat(bulkRequestLatencies.get(0).getValue(), closeTo(1.0, 0)); /** * Metrics: Bulk Request Size in Bytes @@ -406,11 +407,11 @@ public void testOutputServiceMapDefault(final boolean estimateBulkSizeUsingCompr final List bulkRequestSizeBytesMetrics = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(OpenSearchSink.BULKREQUEST_SIZE_BYTES).toString()); - MatcherAssert.assertThat(bulkRequestSizeBytesMetrics.size(), equalTo(3)); - MatcherAssert.assertThat(bulkRequestSizeBytesMetrics.get(0).getValue(), closeTo(1.0, 0)); + assertThat(bulkRequestSizeBytesMetrics.size(), equalTo(3)); + assertThat(bulkRequestSizeBytesMetrics.get(0).getValue(), closeTo(1.0, 0)); final double expectedBulkRequestSizeBytes = isRequestCompressionEnabled && estimateBulkSizeUsingCompression ? 366.0 : 265.0; - MatcherAssert.assertThat(bulkRequestSizeBytesMetrics.get(1).getValue(), closeTo(expectedBulkRequestSizeBytes, 0)); - MatcherAssert.assertThat(bulkRequestSizeBytesMetrics.get(2).getValue(), closeTo(expectedBulkRequestSizeBytes, 0)); + assertThat(bulkRequestSizeBytesMetrics.get(1).getValue(), closeTo(expectedBulkRequestSizeBytes, 0)); + assertThat(bulkRequestSizeBytesMetrics.get(2).getValue(), closeTo(expectedBulkRequestSizeBytes, 0)); // Check restart for index already exists sink = createObjectUnderTest(pluginSetting, true); @@ -428,7 +429,7 @@ public void testInstantiateSinkCustomIndex_NoRollOver() throws IOException { OpenSearchIntegrationHelper.getVersion()) >= 0 ? INCLUDE_TYPE_NAME_FALSE_URI : ""; final Request request = new Request(HttpMethod.HEAD, testIndexAlias + extraURI); final Response response = client.performRequest(request); - MatcherAssert.assertThat(response.getStatusLine().getStatusCode(), equalTo(SC_OK)); + assertThat(response.getStatusLine().getStatusCode(), equalTo(SC_OK)); sink.shutdown(); // Check restart for index already exists @@ -436,32 +437,53 @@ public void testInstantiateSinkCustomIndex_NoRollOver() throws IOException { sink.shutdown(); } - @Test + @ParameterizedTest + @ArgumentsSource(CreateSingleWithTemplatesArgumentsProvider.class) @DisabledIf(value = "isES6", disabledReason = TRACE_INGESTION_TEST_DISABLED_REASON) - public void testInstantiateSinkCustomIndex_WithIsmPolicy() throws IOException { + public void testInstantiateSinkCustomIndex_WithIsmPolicy( + final String templateType, + final String templateFile) throws IOException { final String indexAlias = "sink-custom-index-ism-test-alias"; final String testTemplateFile = Objects.requireNonNull( - getClass().getClassLoader().getResource(TEST_TEMPLATE_V1_FILE)).getFile(); + getClass().getClassLoader().getResource(templateFile)).getFile(); final Map metadata = initializeConfigurationMetadata(null, indexAlias, testTemplateFile); metadata.put(IndexConfiguration.ISM_POLICY_FILE, TEST_CUSTOM_INDEX_POLICY_FILE); + metadata.put(IndexConfiguration.TEMPLATE_TYPE, templateType); final PluginSetting pluginSetting = generatePluginSettingByMetadata(metadata); + OpenSearchSink sink = createObjectUnderTest(pluginSetting, true); final String extraURI = DeclaredOpenSearchVersion.OPENDISTRO_0_10.compareTo( OpenSearchIntegrationHelper.getVersion()) >= 0 ? INCLUDE_TYPE_NAME_FALSE_URI : ""; Request request = new Request(HttpMethod.HEAD, indexAlias + extraURI); Response response = client.performRequest(request); - MatcherAssert.assertThat(response.getStatusLine().getStatusCode(), equalTo(SC_OK)); + assertThat(response.getStatusLine().getStatusCode(), equalTo(SC_OK)); final String index = String.format("%s-000001", indexAlias); final Map mappings = getIndexMappings(index); - MatcherAssert.assertThat(mappings, notNullValue()); - MatcherAssert.assertThat((boolean) mappings.get("date_detection"), equalTo(false)); + assertThat(mappings, notNullValue()); + assertThat((boolean) mappings.get("date_detection"), equalTo(false)); + sink.shutdown(); + JsonNode settings = getIndexSettings(index); + + assertThat(settings, notNullValue()); + JsonNode settingsIndexNode = settings.get("index"); + assertThat(settingsIndexNode, notNullValue()); + assertThat(settingsIndexNode.getNodeType(), equalTo(JsonNodeType.OBJECT)); + assertThat(settingsIndexNode.get("opendistro"), notNullValue()); + assertThat(settingsIndexNode.get("opendistro").getNodeType(), equalTo(JsonNodeType.OBJECT)); + JsonNode settingsIsmNode = settingsIndexNode.get("opendistro").get("index_state_management"); + assertThat(settingsIsmNode, notNullValue()); + assertThat(settingsIsmNode.getNodeType(), equalTo(JsonNodeType.OBJECT)); + assertThat(settingsIsmNode.get("rollover_alias"), notNullValue()); + assertThat(settingsIsmNode.get("rollover_alias").getNodeType(), equalTo(JsonNodeType.STRING)); + assertThat(settingsIsmNode.get("rollover_alias").textValue(), equalTo(indexAlias)); + final String expectedIndexPolicyName = indexAlias + "-policy"; if (isOSBundle()) { // Check managed index await().atMost(1, TimeUnit.SECONDS).untilAsserted(() -> { - MatcherAssert.assertThat(getIndexPolicyId(index), equalTo(expectedIndexPolicyName)); + assertThat(getIndexPolicyId(index), equalTo(expectedIndexPolicyName)); } ); } @@ -470,7 +492,7 @@ public void testInstantiateSinkCustomIndex_WithIsmPolicy() throws IOException { request = new Request(HttpMethod.POST, String.format("%s/_rollover", indexAlias)); request.setJsonEntity("{ \"conditions\" : { } }\n"); response = client.performRequest(request); - MatcherAssert.assertThat(response.getStatusLine().getStatusCode(), equalTo(SC_OK)); + assertThat(response.getStatusLine().getStatusCode(), equalTo(SC_OK)); // Instantiate sink again sink = createObjectUnderTest(pluginSetting, true); @@ -478,12 +500,12 @@ public void testInstantiateSinkCustomIndex_WithIsmPolicy() throws IOException { final String rolloverIndexName = String.format("%s-000002", indexAlias); request = new Request(HttpMethod.GET, rolloverIndexName + "/_alias"); response = client.performRequest(request); - MatcherAssert.assertThat(checkIsWriteIndex(EntityUtils.toString(response.getEntity()), indexAlias, rolloverIndexName), equalTo(true)); + assertThat(checkIsWriteIndex(EntityUtils.toString(response.getEntity()), indexAlias, rolloverIndexName), equalTo(true)); sink.shutdown(); if (isOSBundle()) { // Check managed index - MatcherAssert.assertThat(getIndexPolicyId(rolloverIndexName), equalTo(expectedIndexPolicyName)); + assertThat(getIndexPolicyId(rolloverIndexName), equalTo(expectedIndexPolicyName)); } } @@ -509,14 +531,14 @@ public void testInstantiateSinkDoesNotOverwriteNewerIndexTemplates( Request getTemplateRequest = new Request(HttpMethod.GET, "/" + templatePath + "/" + expectedIndexTemplateName + extraURI); Response getTemplateResponse = client.performRequest(getTemplateRequest); - MatcherAssert.assertThat(getTemplateResponse.getStatusLine().getStatusCode(), equalTo(SC_OK)); + assertThat(getTemplateResponse.getStatusLine().getStatusCode(), equalTo(SC_OK)); String responseBody = EntityUtils.toString(getTemplateResponse.getEntity()); @SuppressWarnings("unchecked") final Integer firstResponseVersion = extractVersionFunction.apply(createContentParser(XContentType.JSON.xContent(), responseBody).map(), expectedIndexTemplateName); - MatcherAssert.assertThat(firstResponseVersion, equalTo(Integer.valueOf(1))); + assertThat(firstResponseVersion, equalTo(Integer.valueOf(1))); sink.shutdown(); // Create sink with template version 2 @@ -526,14 +548,14 @@ public void testInstantiateSinkDoesNotOverwriteNewerIndexTemplates( getTemplateRequest = new Request(HttpMethod.GET, "/" + templatePath + "/" + expectedIndexTemplateName + extraURI); getTemplateResponse = client.performRequest(getTemplateRequest); - MatcherAssert.assertThat(getTemplateResponse.getStatusLine().getStatusCode(), equalTo(SC_OK)); + assertThat(getTemplateResponse.getStatusLine().getStatusCode(), equalTo(SC_OK)); responseBody = EntityUtils.toString(getTemplateResponse.getEntity()); @SuppressWarnings("unchecked") final Integer secondResponseVersion = extractVersionFunction.apply(createContentParser(XContentType.JSON.xContent(), responseBody).map(), expectedIndexTemplateName); - MatcherAssert.assertThat(secondResponseVersion, equalTo(Integer.valueOf(2))); + assertThat(secondResponseVersion, equalTo(Integer.valueOf(2))); sink.shutdown(); // Create sink with template version 1 again @@ -543,7 +565,7 @@ public void testInstantiateSinkDoesNotOverwriteNewerIndexTemplates( getTemplateRequest = new Request(HttpMethod.GET, "/" + templatePath + "/" + expectedIndexTemplateName + extraURI); getTemplateResponse = client.performRequest(getTemplateRequest); - MatcherAssert.assertThat(getTemplateResponse.getStatusLine().getStatusCode(), equalTo(SC_OK)); + assertThat(getTemplateResponse.getStatusLine().getStatusCode(), equalTo(SC_OK)); responseBody = EntityUtils.toString(getTemplateResponse.getEntity()); @SuppressWarnings("unchecked") final Integer thirdResponseVersion = @@ -551,7 +573,7 @@ public void testInstantiateSinkDoesNotOverwriteNewerIndexTemplates( responseBody).map(), expectedIndexTemplateName); // Assert version 2 was not overwritten by version 1 - MatcherAssert.assertThat(thirdResponseVersion, equalTo(Integer.valueOf(2))); + assertThat(thirdResponseVersion, equalTo(Integer.valueOf(2))); sink.shutdown(); } @@ -568,7 +590,7 @@ public Stream provideArguments(ExtensionContext context) { ) ); - if(OpenSearchIntegrationHelper.getVersion().compareTo(DeclaredOpenSearchVersion.OPENDISTRO_1_9) >= 0) { + if (OpenSearchIntegrationHelper.getVersion().compareTo(DeclaredOpenSearchVersion.OPENDISTRO_1_9) >= 0) { arguments.add( arguments("index-template", "_index_template", TEST_INDEX_TEMPLATE_V1_FILE, TEST_INDEX_TEMPLATE_V2_FILE, @@ -581,6 +603,19 @@ public Stream provideArguments(ExtensionContext context) { } } + static class CreateSingleWithTemplatesArgumentsProvider implements ArgumentsProvider { + @Override + public Stream provideArguments(ExtensionContext context) { + final List arguments = new ArrayList<>(); + arguments.add(arguments("v1", TEST_TEMPLATE_V1_FILE)); + + if (OpenSearchIntegrationHelper.getVersion().compareTo(DeclaredOpenSearchVersion.OPENDISTRO_1_9) >= 0) { + arguments.add(arguments("index-template", TEST_INDEX_TEMPLATE_V1_FILE)); + } + return arguments.stream(); + } + } + @Test public void testOutputCustomIndex() throws IOException, InterruptedException { final String testIndexAlias = "test-alias"; @@ -594,15 +629,15 @@ public void testOutputCustomIndex() throws IOException, InterruptedException { final OpenSearchSink sink = createObjectUnderTest(pluginSetting, true); sink.output(testRecords); final List> retSources = getSearchResponseDocSources(testIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); - MatcherAssert.assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); + assertThat(retSources.size(), equalTo(1)); + assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); sink.shutdown(); // verify metrics final List bulkRequestLatencies = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(OpenSearchSink.BULKREQUEST_LATENCY).toString()); - MatcherAssert.assertThat(bulkRequestLatencies.size(), equalTo(3)); + assertThat(bulkRequestLatencies.size(), equalTo(3)); // COUNT Assert.assertEquals(1.0, bulkRequestLatencies.get(0).getValue(), 0); } @@ -621,15 +656,15 @@ public void testOpenSearchBulkActionsCreate() throws IOException, InterruptedExc final OpenSearchSink sink = createObjectUnderTest(pluginSetting, true); sink.output(testRecords); final List> retSources = getSearchResponseDocSources(testIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); - MatcherAssert.assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); + assertThat(retSources.size(), equalTo(1)); + assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); sink.shutdown(); // verify metrics final List bulkRequestLatencies = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(OpenSearchSink.BULKREQUEST_LATENCY).toString()); - MatcherAssert.assertThat(bulkRequestLatencies.size(), equalTo(3)); + assertThat(bulkRequestLatencies.size(), equalTo(3)); // COUNT Assert.assertEquals(1.0, bulkRequestLatencies.get(0).getValue(), 0); } @@ -644,7 +679,7 @@ public void testOpenSearchBulkActionsCreateWithExpression() throws IOException, final List> testRecords = Collections.singletonList(jsonStringToRecord(generateCustomRecordJson(testIdField, testId))); final PluginSetting pluginSetting = generatePluginSetting(null, testIndexAlias, testTemplateFile); pluginSetting.getSettings().put(IndexConfiguration.DOCUMENT_ID_FIELD, testIdField); - Event event = (Event)testRecords.get(0).getData(); + Event event = (Event) testRecords.get(0).getData(); event.getMetadata().setAttribute("action", "create"); final String actionFormatExpression = "${getMetadata(\"action\")}"; when(expressionEvaluator.isValidFormatExpressions(actionFormatExpression)).thenReturn(true); @@ -654,15 +689,15 @@ public void testOpenSearchBulkActionsCreateWithExpression() throws IOException, final OpenSearchSink sink = createObjectUnderTest(pluginSetting, true); sink.output(testRecords); final List> retSources = getSearchResponseDocSources(testIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); - MatcherAssert.assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); + assertThat(retSources.size(), equalTo(1)); + assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); sink.shutdown(); // verify metrics final List bulkRequestLatencies = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(OpenSearchSink.BULKREQUEST_LATENCY).toString()); - MatcherAssert.assertThat(bulkRequestLatencies.size(), equalTo(3)); + assertThat(bulkRequestLatencies.size(), equalTo(3)); // COUNT Assert.assertEquals(1.0, bulkRequestLatencies.get(0).getValue(), 0); } @@ -677,7 +712,7 @@ public void testOpenSearchBulkActionsCreateWithInvalidExpression() throws IOExce final List> testRecords = Collections.singletonList(jsonStringToRecord(generateCustomRecordJson(testIdField, testId))); final PluginSetting pluginSetting = generatePluginSetting(null, testIndexAlias, testTemplateFile); pluginSetting.getSettings().put(IndexConfiguration.DOCUMENT_ID_FIELD, testIdField); - Event event = (Event)testRecords.get(0).getData(); + Event event = (Event) testRecords.get(0).getData(); event.getMetadata().setAttribute("action", "unknown"); final String actionFormatExpression = "${getMetadata(\"action\")}"; when(expressionEvaluator.isValidFormatExpressions(actionFormatExpression)).thenReturn(true); @@ -687,8 +722,8 @@ public void testOpenSearchBulkActionsCreateWithInvalidExpression() throws IOExce final OpenSearchSink sink = createObjectUnderTest(pluginSetting, true); sink.output(testRecords); final List> retSources = getSearchResponseDocSources(testIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(0)); - MatcherAssert.assertThat(sink.getInvalidActionErrorsCount(), equalTo(1.0)); + assertThat(retSources.size(), equalTo(0)); + assertThat(sink.getInvalidActionErrorsCount(), equalTo(1.0)); sink.shutdown(); } @@ -711,15 +746,15 @@ public void testBulkActionCreateWithActions() throws IOException, InterruptedExc final OpenSearchSink sink = createObjectUnderTest(pluginSetting, true); sink.output(testRecords); final List> retSources = getSearchResponseDocSources(testIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); - MatcherAssert.assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); + assertThat(retSources.size(), equalTo(1)); + assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); sink.shutdown(); // verify metrics final List bulkRequestLatencies = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(OpenSearchSink.BULKREQUEST_LATENCY).toString()); - MatcherAssert.assertThat(bulkRequestLatencies.size(), equalTo(3)); + assertThat(bulkRequestLatencies.size(), equalTo(3)); // COUNT Assert.assertEquals(1.0, bulkRequestLatencies.get(0).getValue(), 0); } @@ -744,15 +779,15 @@ public void testBulkActionUpdateWithActions() throws IOException, InterruptedExc OpenSearchSink sink = createObjectUnderTest(pluginSetting, true); sink.output(testRecords); List> retSources = getSearchResponseDocSources(testIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); - MatcherAssert.assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); + assertThat(retSources.size(), equalTo(1)); + assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); sink.shutdown(); // verify metrics final List bulkRequestLatencies = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(OpenSearchSink.BULKREQUEST_LATENCY).toString()); - MatcherAssert.assertThat(bulkRequestLatencies.size(), equalTo(3)); + assertThat(bulkRequestLatencies.size(), equalTo(3)); // COUNT Assert.assertEquals(1.0, bulkRequestLatencies.get(0).getValue(), 0); testRecords = Collections.singletonList(jsonStringToRecord(generateCustomRecordJson2(testIdField, testId, "name", "value2"))); @@ -765,10 +800,10 @@ public void testBulkActionUpdateWithActions() throws IOException, InterruptedExc sink.output(testRecords); retSources = getSearchResponseDocSources(testIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); + assertThat(retSources.size(), equalTo(1)); Map source = retSources.get(0); - MatcherAssert.assertThat((String)source.get("name"), equalTo("value2")); - MatcherAssert.assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); + assertThat((String) source.get("name"), equalTo("value2")); + assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); sink.shutdown(); } @@ -792,15 +827,15 @@ public void testBulkActionUpsertWithActions() throws IOException, InterruptedExc OpenSearchSink sink = createObjectUnderTest(pluginSetting, true); sink.output(testRecords); List> retSources = getSearchResponseDocSources(testIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); - MatcherAssert.assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); + assertThat(retSources.size(), equalTo(1)); + assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); sink.shutdown(); // verify metrics final List bulkRequestLatencies = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(OpenSearchSink.BULKREQUEST_LATENCY).toString()); - MatcherAssert.assertThat(bulkRequestLatencies.size(), equalTo(3)); + assertThat(bulkRequestLatencies.size(), equalTo(3)); // COUNT Assert.assertEquals(1.0, bulkRequestLatencies.get(0).getValue(), 0); testRecords = Collections.singletonList(jsonStringToRecord(generateCustomRecordJson3(testIdField, testId, "name", "value3", "newKey", "newValue"))); @@ -813,11 +848,11 @@ public void testBulkActionUpsertWithActions() throws IOException, InterruptedExc sink.output(testRecords); retSources = getSearchResponseDocSources(testIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); + assertThat(retSources.size(), equalTo(1)); Map source = retSources.get(0); - MatcherAssert.assertThat((String)source.get("name"), equalTo("value3")); - MatcherAssert.assertThat((String)source.get("newKey"), equalTo("newValue")); - MatcherAssert.assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); + assertThat((String) source.get("name"), equalTo("value3")); + assertThat((String) source.get("newKey"), equalTo("newValue")); + assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); sink.shutdown(); } @@ -841,17 +876,17 @@ public void testBulkActionUpsertWithoutCreate() throws IOException, InterruptedE sink.output(testRecords); List> retSources = getSearchResponseDocSources(testIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); + assertThat(retSources.size(), equalTo(1)); Map source = retSources.get(0); - MatcherAssert.assertThat((String)source.get("name"), equalTo("value1")); - MatcherAssert.assertThat((String)source.get("newKey"), equalTo("newValue")); - MatcherAssert.assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); + assertThat((String) source.get("name"), equalTo("value1")); + assertThat((String) source.get("newKey"), equalTo("newValue")); + assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); sink.shutdown(); // verify metrics final List bulkRequestLatencies = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(OpenSearchSink.BULKREQUEST_LATENCY).toString()); - MatcherAssert.assertThat(bulkRequestLatencies.size(), equalTo(3)); + assertThat(bulkRequestLatencies.size(), equalTo(3)); // COUNT Assert.assertEquals(1.0, bulkRequestLatencies.get(0).getValue(), 0); } @@ -876,7 +911,7 @@ public void testBulkActionDeleteWithActions() throws IOException, InterruptedExc OpenSearchSink sink = createObjectUnderTest(pluginSetting, true); sink.output(testRecords); List> retSources = getSearchResponseDocSources(testIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(0)); + assertThat(retSources.size(), equalTo(0)); sink.shutdown(); } @@ -902,9 +937,9 @@ public void testEventOutputWithTags() throws IOException, InterruptedException { expectedContent.put("log", "foobar"); expectedContent.put(testTagsTargetKey, tagsList); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); - MatcherAssert.assertThat(retSources.containsAll(Arrays.asList(expectedContent)), equalTo(true)); - MatcherAssert.assertThat(getDocumentCount(expIndexAlias, "log", "foobar"), equalTo(Integer.valueOf(1))); + assertThat(retSources.size(), equalTo(1)); + assertThat(retSources.containsAll(Arrays.asList(expectedContent)), equalTo(true)); + assertThat(getDocumentCount(expIndexAlias, "log", "foobar"), equalTo(Integer.valueOf(1))); sink.shutdown(); } @@ -928,9 +963,9 @@ public void testEventOutput() throws IOException, InterruptedException { final Map expectedContent = new HashMap<>(); expectedContent.put("log", "foobar"); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); - MatcherAssert.assertThat(retSources.containsAll(Arrays.asList(expectedContent)), equalTo(true)); - MatcherAssert.assertThat(getDocumentCount(expIndexAlias, "log", "foobar"), equalTo(Integer.valueOf(1))); + assertThat(retSources.size(), equalTo(1)); + assertThat(retSources.containsAll(Arrays.asList(expectedContent)), equalTo(true)); + assertThat(getDocumentCount(expIndexAlias, "log", "foobar"), equalTo(Integer.valueOf(1))); sink.shutdown(); } @@ -954,7 +989,7 @@ public void testOpenSearchDocumentId(final String testDocumentIdField) throws IO final List docIds = getSearchResponseDocIds(testIndexAlias); for (String docId : docIds) { - MatcherAssert.assertThat(docId, equalTo(expectedId)); + assertThat(docId, equalTo(expectedId)); } sink.shutdown(); } @@ -979,7 +1014,7 @@ public void testOpenSearchRoutingField(final String testRoutingField) throws IOE final List routingFields = getSearchResponseRoutingFields(testIndexAlias); for (String routingField : routingFields) { - MatcherAssert.assertThat(routingField, equalTo(expectedRoutingField)); + assertThat(routingField, equalTo(expectedRoutingField)); } sink.shutdown(); } @@ -1006,8 +1041,8 @@ public void testOpenSearchDynamicIndex(final String testIndex) throws IOExceptio final OpenSearchSink sink = createObjectUnderTest(pluginSetting, true); sink.output(testRecords); final List> retSources = getSearchResponseDocSources(testIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); - MatcherAssert.assertThat(retSources, hasItem(expectedMap)); + assertThat(retSources.size(), equalTo(1)); + assertThat(retSources, hasItem(expectedMap)); sink.shutdown(); } @@ -1039,8 +1074,8 @@ public void testOpenSearchDynamicIndexWithDate(final String testIndex, final Str final OpenSearchSink sink = createObjectUnderTest(pluginSetting, true); sink.output(testRecords); final List> retSources = getSearchResponseDocSources(expectedIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); - MatcherAssert.assertThat(retSources, hasItem(expectedMap)); + assertThat(retSources.size(), equalTo(1)); + assertThat(retSources, hasItem(expectedMap)); sink.shutdown(); } @@ -1067,8 +1102,8 @@ public void testOpenSearchIndexWithDate(final String testDatePattern) throws IOE final OpenSearchSink sink = createObjectUnderTest(pluginSetting, true); sink.output(testRecords); final List> retSources = getSearchResponseDocSources(expectedIndexName); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); - MatcherAssert.assertThat(retSources, hasItem(expectedMap)); + assertThat(retSources.size(), equalTo(1)); + assertThat(retSources, hasItem(expectedMap)); sink.shutdown(); } @@ -1122,15 +1157,15 @@ public void testOutputManagementDisabled() throws IOException, InterruptedExcept sink.output(testRecords); final List> retSources = getSearchResponseDocSources(testIndexAlias); - MatcherAssert.assertThat(retSources.size(), equalTo(1)); - MatcherAssert.assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); + assertThat(retSources.size(), equalTo(1)); + assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1))); sink.shutdown(); // verify metrics final List bulkRequestLatencies = MetricsTestUtil.getMeasurementList( new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME) .add(OpenSearchSink.BULKREQUEST_LATENCY).toString()); - MatcherAssert.assertThat(bulkRequestLatencies.size(), equalTo(3)); + assertThat(bulkRequestLatencies.size(), equalTo(3)); // COUNT Assert.assertEquals(1.0, bulkRequestLatencies.get(0).getValue(), 0); } @@ -1313,6 +1348,20 @@ private Map getIndexMappings(final String index) throws IOExcept return mappings; } + private JsonNode getIndexSettings(final String index) throws IOException { + final String extraURI = DeclaredOpenSearchVersion.OPENDISTRO_0_10.compareTo( + OpenSearchIntegrationHelper.getVersion()) >= 0 ? INCLUDE_TYPE_NAME_FALSE_URI : ""; + final Request request = new Request(HttpMethod.GET, index + "/_settings" + extraURI); + final Response response = client.performRequest(request); + final String responseBody = EntityUtils.toString(response.getEntity()); + + Map responseMap = createContentParser(XContentType.JSON.xContent(), responseBody).map(); + + return new ObjectMapper().convertValue(responseMap, JsonNode.class) + .get(index) + .get("settings"); + } + private String getIndexPolicyId(final String index) throws IOException { // TODO: replace with new _opensearch API final Request request = new Request(HttpMethod.GET, "/_opendistro/_ism/explain/" + index); diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ComposableIndexTemplate.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ComposableIndexTemplate.java index f7a42d278b..d94fcfe8cd 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ComposableIndexTemplate.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ComposableIndexTemplate.java @@ -6,7 +6,9 @@ import java.util.Map; import java.util.Optional; -public class ComposableIndexTemplate implements IndexTemplate { +class ComposableIndexTemplate implements IndexTemplate { + static final String TEMPLATE_KEY = "template"; + static final String INDEX_SETTINGS_KEY = "settings"; private final Map indexTemplateMap; private String name; @@ -18,7 +20,6 @@ public ComposableIndexTemplate(final Map indexTemplateMap) { @Override public void setTemplateName(final String name) { this.name = name; - } @Override @@ -28,7 +29,11 @@ public void setIndexPatterns(final List indexPatterns) { @Override public void putCustomSetting(final String name, final Object value) { + Map template = (Map) indexTemplateMap.computeIfAbsent(TEMPLATE_KEY, key -> new HashMap<>()); + + Map settings = (Map) template.computeIfAbsent(INDEX_SETTINGS_KEY, key -> new HashMap<>()); + settings.put(name, value); } @Override diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ComposableTemplateAPIWrapper.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ComposableTemplateAPIWrapper.java index 1b78cb9da5..d2a72a6e75 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ComposableTemplateAPIWrapper.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ComposableTemplateAPIWrapper.java @@ -1,26 +1,21 @@ package org.opensearch.dataprepper.plugins.sink.opensearch.index; -import com.fasterxml.jackson.databind.ObjectMapper; -import jakarta.json.stream.JsonParser; -import org.opensearch.client.json.JsonpDeserializer; -import org.opensearch.client.json.JsonpMapper; -import org.opensearch.client.json.ObjectBuilderDeserializer; -import org.opensearch.client.json.ObjectDeserializer; import org.opensearch.client.opensearch.OpenSearchClient; +import org.opensearch.client.opensearch._types.ErrorResponse; import org.opensearch.client.opensearch.indices.ExistsIndexTemplateRequest; import org.opensearch.client.opensearch.indices.GetIndexTemplateRequest; import org.opensearch.client.opensearch.indices.GetIndexTemplateResponse; -import org.opensearch.client.opensearch.indices.PutIndexTemplateRequest; -import org.opensearch.client.opensearch.indices.put_index_template.IndexTemplateMapping; +import org.opensearch.client.opensearch.indices.PutIndexTemplateResponse; +import org.opensearch.client.transport.Endpoint; import org.opensearch.client.transport.endpoints.BooleanResponse; +import org.opensearch.client.transport.endpoints.SimpleEndpoint; -import java.io.ByteArrayInputStream; import java.io.IOException; -import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.Map; import java.util.Optional; public class ComposableTemplateAPIWrapper implements IndexTemplateAPIWrapper { - private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); private final OpenSearchClient openSearchClient; public ComposableTemplateAPIWrapper(final OpenSearchClient openSearchClient) { @@ -34,19 +29,12 @@ public void putTemplate(final IndexTemplate indexTemplate) throws IOException { } final ComposableIndexTemplate composableIndexTemplate = (ComposableIndexTemplate) indexTemplate; - final String indexTemplateString = OBJECT_MAPPER.writeValueAsString( - composableIndexTemplate.getIndexTemplateMap()); + Map indexTemplateMap = composableIndexTemplate.getIndexTemplateMap(); - final ByteArrayInputStream byteIn = new ByteArrayInputStream( - indexTemplateString.getBytes(StandardCharsets.UTF_8)); - final JsonpMapper mapper = openSearchClient._transport().jsonpMapper(); - final JsonParser parser = mapper.jsonProvider().createParser(byteIn); - - final PutIndexTemplateRequest putIndexTemplateRequest = PutIndexTemplateRequestDeserializer - .getJsonpDeserializer(composableIndexTemplate.getName()) - .deserialize(parser, mapper); - - openSearchClient.indices().putIndexTemplate(putIndexTemplateRequest); + openSearchClient._transport().performRequest( + indexTemplateMap, + createEndpoint(composableIndexTemplate), + openSearchClient._transportOptions()); } @Override @@ -66,24 +54,15 @@ public Optional getTemplate(final String indexTemplate return Optional.of(openSearchClient.indices().getIndexTemplate(getRequest)); } - private static class PutIndexTemplateRequestDeserializer { - private static void setupPutIndexTemplateRequestDeserializer(final ObjectDeserializer objectDeserializer) { + private Endpoint, PutIndexTemplateResponse, ErrorResponse> createEndpoint(final ComposableIndexTemplate composableIndexTemplate) { + final String path = "/_index_template/" + composableIndexTemplate.getName(); - objectDeserializer.add(PutIndexTemplateRequest.Builder::name, JsonpDeserializer.stringDeserializer(), "name"); - objectDeserializer.add(PutIndexTemplateRequest.Builder::indexPatterns, JsonpDeserializer.arrayDeserializer(JsonpDeserializer.stringDeserializer()), - "index_patterns"); - objectDeserializer.add(PutIndexTemplateRequest.Builder::version, JsonpDeserializer.longDeserializer(), "version"); - objectDeserializer.add(PutIndexTemplateRequest.Builder::priority, JsonpDeserializer.integerDeserializer(), "priority"); - objectDeserializer.add(PutIndexTemplateRequest.Builder::composedOf, JsonpDeserializer.arrayDeserializer(JsonpDeserializer.stringDeserializer()), - "composed_of"); - objectDeserializer.add(PutIndexTemplateRequest.Builder::template, IndexTemplateMapping._DESERIALIZER, "template"); - } - - static JsonpDeserializer getJsonpDeserializer(final String name) { - return ObjectBuilderDeserializer - .lazy( - () -> new PutIndexTemplateRequest.Builder().name(name), - PutIndexTemplateRequestDeserializer::setupPutIndexTemplateRequestDeserializer); - } + return new SimpleEndpoint<>( + request -> "PUT", + request -> path, + request -> Collections.emptyMap(), + SimpleEndpoint.emptyMap(), + true, + PutIndexTemplateResponse._DESERIALIZER); } } diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/LegacyIndexTemplate.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/LegacyIndexTemplate.java index 1d6e6be00e..4e54d4005c 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/LegacyIndexTemplate.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/LegacyIndexTemplate.java @@ -5,7 +5,7 @@ import java.util.Map; import java.util.Optional; -public class LegacyIndexTemplate implements IndexTemplate { +class LegacyIndexTemplate implements IndexTemplate { public static final String SETTINGS_KEY = "settings"; private final Map templateMap; diff --git a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ComposableIndexTemplateStrategyTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ComposableIndexTemplateStrategyTest.java index f9e0154007..7668d6fee2 100644 --- a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ComposableIndexTemplateStrategyTest.java +++ b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ComposableIndexTemplateStrategyTest.java @@ -26,9 +26,11 @@ import java.util.UUID; import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.instanceOf; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.notNullValue; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.hasKey; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; @@ -36,6 +38,8 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.plugins.sink.opensearch.index.ComposableIndexTemplate.INDEX_SETTINGS_KEY; +import static org.opensearch.dataprepper.plugins.sink.opensearch.index.ComposableIndexTemplate.TEMPLATE_KEY; @ExtendWith(MockitoExtension.class) class ComposableIndexTemplateStrategyTest { @@ -188,5 +192,90 @@ void getVersion_returns_version_from_root_map_when_provided_as_int() { assertThat(optionalVersion.isPresent(), equalTo(true)); assertThat(optionalVersion.get(), equalTo((long) version)); } + + @Test + void putCustomSetting_with_no_existing_template_adds_template_and_settings() { + final org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexTemplate indexTemplate = + createObjectUnderTest().createIndexTemplate(providedTemplateMap); + + String customKey = UUID.randomUUID().toString(); + String customValue = UUID.randomUUID().toString(); + + indexTemplate.putCustomSetting(customKey, customValue); + + Map indexTemplateMap = ((ComposableIndexTemplate) indexTemplate).getIndexTemplateMap(); + + assertThat(indexTemplateMap, hasKey(TEMPLATE_KEY)); + assertThat(indexTemplateMap.get(TEMPLATE_KEY), instanceOf(Map.class)); + Map templateMap = (Map) indexTemplateMap.get(TEMPLATE_KEY); + assertThat(templateMap, hasKey(INDEX_SETTINGS_KEY)); + assertThat(templateMap.get(INDEX_SETTINGS_KEY), instanceOf(Map.class)); + Map settingsMap = (Map) templateMap.get(INDEX_SETTINGS_KEY); + assertThat(settingsMap, hasKey(customKey)); + assertThat(settingsMap.get(customKey), equalTo(customValue)); + } + + @Test + void putCustomSetting_with_existing_template_adds_settings_to_that_template() { + String existingKey = UUID.randomUUID().toString(); + String existingValue = UUID.randomUUID().toString(); + Map template = new HashMap<>(); + template.put(existingKey, existingValue); + providedTemplateMap.put(TEMPLATE_KEY, template); + + final org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexTemplate indexTemplate = + createObjectUnderTest().createIndexTemplate(providedTemplateMap); + + String customKey = UUID.randomUUID().toString(); + String customValue = UUID.randomUUID().toString(); + + indexTemplate.putCustomSetting(customKey, customValue); + + Map indexTemplateMap = ((ComposableIndexTemplate) indexTemplate).getIndexTemplateMap(); + + assertThat(indexTemplateMap, hasKey(TEMPLATE_KEY)); + assertThat(indexTemplateMap.get(TEMPLATE_KEY), instanceOf(Map.class)); + Map templateMap = (Map) indexTemplateMap.get(TEMPLATE_KEY); + assertThat(templateMap, hasKey(INDEX_SETTINGS_KEY)); + assertThat(templateMap, hasKey(existingKey)); + assertThat(templateMap.get(existingKey), equalTo(existingValue)); + assertThat(templateMap.get(INDEX_SETTINGS_KEY), instanceOf(Map.class)); + Map settingsMap = (Map) templateMap.get(INDEX_SETTINGS_KEY); + assertThat(settingsMap, hasKey(customKey)); + assertThat(settingsMap.get(customKey), equalTo(customValue)); + } + + @Test + void putCustomSetting_with_existing_template_and_settings_puts_settings_to_that_settings() { + String existingKey = UUID.randomUUID().toString(); + String existingValue = UUID.randomUUID().toString(); + Map template = new HashMap<>(); + HashMap existingSettings = new HashMap<>(); + existingSettings.put(existingKey, existingValue); + template.put(INDEX_SETTINGS_KEY, existingSettings); + + providedTemplateMap.put(TEMPLATE_KEY, template); + + final org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexTemplate indexTemplate = + createObjectUnderTest().createIndexTemplate(providedTemplateMap); + + String customKey = UUID.randomUUID().toString(); + String customValue = UUID.randomUUID().toString(); + + indexTemplate.putCustomSetting(customKey, customValue); + + Map indexTemplateMap = ((ComposableIndexTemplate) indexTemplate).getIndexTemplateMap(); + + assertThat(indexTemplateMap, hasKey(TEMPLATE_KEY)); + assertThat(indexTemplateMap.get(TEMPLATE_KEY), instanceOf(Map.class)); + Map templateMap = (Map) indexTemplateMap.get(TEMPLATE_KEY); + assertThat(templateMap, hasKey(INDEX_SETTINGS_KEY)); + assertThat(templateMap.get(INDEX_SETTINGS_KEY), instanceOf(Map.class)); + Map settingsMap = (Map) templateMap.get(INDEX_SETTINGS_KEY); + assertThat(settingsMap, hasKey(customKey)); + assertThat(settingsMap.get(customKey), equalTo(customValue)); + assertThat(settingsMap, hasKey(existingKey)); + assertThat(settingsMap.get(existingKey), equalTo(existingValue)); + } } } \ No newline at end of file diff --git a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ComposableTemplateAPIWrapperTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ComposableTemplateAPIWrapperTest.java index d6040a017d..fe2829656d 100644 --- a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ComposableTemplateAPIWrapperTest.java +++ b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/ComposableTemplateAPIWrapperTest.java @@ -9,14 +9,17 @@ import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.client.opensearch.OpenSearchClient; +import org.opensearch.client.opensearch._types.ErrorResponse; import org.opensearch.client.opensearch.indices.ExistsIndexTemplateRequest; import org.opensearch.client.opensearch.indices.GetIndexTemplateRequest; import org.opensearch.client.opensearch.indices.GetIndexTemplateResponse; import org.opensearch.client.opensearch.indices.OpenSearchIndicesClient; -import org.opensearch.client.opensearch.indices.PutIndexTemplateRequest; +import org.opensearch.client.opensearch.indices.PutIndexTemplateResponse; +import org.opensearch.client.transport.Endpoint; +import org.opensearch.client.transport.JsonEndpoint; import org.opensearch.client.transport.OpenSearchTransport; +import org.opensearch.client.transport.TransportOptions; import org.opensearch.client.transport.endpoints.BooleanResponse; -import org.opensearch.dataprepper.plugins.sink.opensearch.bulk.PreSerializedJsonpMapper; import java.io.IOException; import java.util.Collections; @@ -28,11 +31,15 @@ import java.util.UUID; import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.instanceOf; +import static org.hamcrest.CoreMatchers.not; import static org.hamcrest.CoreMatchers.notNullValue; import static org.hamcrest.CoreMatchers.nullValue; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.hasKey; 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.lenient; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; @@ -47,11 +54,15 @@ class ComposableTemplateAPIWrapperTest { @Mock private OpenSearchTransport openSearchTransport; @Mock + private TransportOptions openSearchTransportOptions; + @Mock private GetIndexTemplateResponse getIndexTemplateResponse; @Mock private BooleanResponse booleanResponse; @Captor - private ArgumentCaptor putIndexTemplateRequestArgumentCaptor; + private ArgumentCaptor> putIndexTemplateRequestArgumentCaptor; + @Captor + private ArgumentCaptor, PutIndexTemplateResponse, ErrorResponse>> endpointArgumentCaptor; @Captor private ArgumentCaptor existsIndexTemplateRequestArgumentCaptor; @@ -76,7 +87,7 @@ void putTemplate_throws_if_template_is_not_ComposableIndexTemplate() { @Test void putTemplate_performs_putIndexTemplate_request() throws IOException { when(openSearchClient._transport()).thenReturn(openSearchTransport); - when(openSearchTransport.jsonpMapper()).thenReturn(new PreSerializedJsonpMapper()); + when(openSearchClient._transportOptions()).thenReturn(openSearchTransportOptions); final List indexPatterns = Collections.singletonList(UUID.randomUUID().toString()); final IndexTemplate indexTemplate = new ComposableIndexTemplate(new HashMap<>()); @@ -84,12 +95,25 @@ void putTemplate_performs_putIndexTemplate_request() throws IOException { indexTemplate.setIndexPatterns(indexPatterns); objectUnderTest.putTemplate(indexTemplate); - verify(openSearchIndicesClient).putIndexTemplate(putIndexTemplateRequestArgumentCaptor.capture()); + verify(openSearchTransport).performRequest( + putIndexTemplateRequestArgumentCaptor.capture(), + endpointArgumentCaptor.capture(), + eq(openSearchTransportOptions) + ); + + Map actualPutRequest = putIndexTemplateRequestArgumentCaptor.getValue(); - final PutIndexTemplateRequest actualPutRequest = putIndexTemplateRequestArgumentCaptor.getValue(); + assertThat(actualPutRequest.get("index_patterns"), equalTo(indexPatterns)); - assertThat(actualPutRequest.name(), equalTo(indexTemplateName)); - assertThat(actualPutRequest.indexPatterns(), equalTo(indexPatterns)); + Endpoint, PutIndexTemplateResponse, ErrorResponse> actualEndpoint = endpointArgumentCaptor.getValue(); + assertThat(actualEndpoint.method(null), equalTo("PUT")); + assertThat(actualEndpoint.requestUrl(null), equalTo("/_index_template/" + indexTemplateName)); + assertThat(actualEndpoint.queryParameters(null), equalTo(Collections.emptyMap())); + assertThat(actualEndpoint.headers(null), equalTo(Collections.emptyMap())); + assertThat(actualEndpoint.hasRequestBody(), equalTo(true)); + + assertThat(actualEndpoint, instanceOf(JsonEndpoint.class)); + assertThat(((JsonEndpoint)actualEndpoint).responseDeserializer(), equalTo(PutIndexTemplateResponse._DESERIALIZER)); } @Test @@ -139,18 +163,15 @@ void getExistingTemplate_should_return_template_if_template_exists() throws IOEx @Nested class IndexTemplateWithCreateTemplateTests { - private ArgumentCaptor putIndexTemplateRequestArgumentCaptor; private List indexPatterns; @BeforeEach void setUp() { - final OpenSearchTransport openSearchTransport = mock(OpenSearchTransport.class); when(openSearchClient._transport()).thenReturn(openSearchTransport); - when(openSearchTransport.jsonpMapper()).thenReturn(new PreSerializedJsonpMapper()); - - putIndexTemplateRequestArgumentCaptor = ArgumentCaptor.forClass(PutIndexTemplateRequest.class); indexPatterns = Collections.singletonList(UUID.randomUUID().toString()); + + when(openSearchClient._transportOptions()).thenReturn(openSearchTransportOptions); } @Test @@ -159,43 +180,64 @@ void putTemplate_with_setTemplateName_performs_putIndexTemplate_request() throws indexTemplate.setTemplateName(indexTemplateName); objectUnderTest.putTemplate(indexTemplate); - verify(openSearchIndicesClient).putIndexTemplate(putIndexTemplateRequestArgumentCaptor.capture()); - - final PutIndexTemplateRequest actualPutRequest = putIndexTemplateRequestArgumentCaptor.getValue(); - - assertThat(actualPutRequest.name(), equalTo(indexTemplateName)); - - assertThat(actualPutRequest.version(), nullValue()); - assertThat(actualPutRequest.indexPatterns(), notNullValue()); - assertThat(actualPutRequest.indexPatterns(), equalTo(Collections.emptyList())); - assertThat(actualPutRequest.template(), nullValue()); - assertThat(actualPutRequest.priority(), nullValue()); - assertThat(actualPutRequest.composedOf(), notNullValue()); - assertThat(actualPutRequest.composedOf(), equalTo(Collections.emptyList())); + verify(openSearchTransport).performRequest( + putIndexTemplateRequestArgumentCaptor.capture(), + endpointArgumentCaptor.capture(), + eq(openSearchTransportOptions) + ); + + final Map actualPutRequest = putIndexTemplateRequestArgumentCaptor.getValue(); + + final Endpoint, PutIndexTemplateResponse, ErrorResponse> actualEndpoint = endpointArgumentCaptor.getValue(); + assertThat(actualEndpoint.method(null), equalTo("PUT")); + assertThat(actualEndpoint.requestUrl(null), equalTo("/_index_template/" + indexTemplateName)); + assertThat(actualEndpoint.queryParameters(null), equalTo(Collections.emptyMap())); + assertThat(actualEndpoint.headers(null), equalTo(Collections.emptyMap())); + assertThat(actualEndpoint.hasRequestBody(), equalTo(true)); + + assertThat(actualEndpoint, instanceOf(JsonEndpoint.class)); + assertThat(((JsonEndpoint)actualEndpoint).responseDeserializer(), equalTo(PutIndexTemplateResponse._DESERIALIZER)); + + assertThat(actualPutRequest.get("version"), nullValue()); + assertThat(actualPutRequest.get("index_patterns"), nullValue()); + assertThat(actualPutRequest.get("template"), nullValue()); + assertThat(actualPutRequest.get("priority"), nullValue()); + assertThat(actualPutRequest.get("composed_of"), nullValue()); } @Test void putTemplate_with_setIndexPatterns_performs_putIndexTemplate_request() throws IOException { final List indexPatterns = Collections.singletonList(UUID.randomUUID().toString()); - final IndexConfiguration indexConfiguration = mock(IndexConfiguration.class); final IndexTemplate indexTemplate = new ComposableIndexTemplate(new HashMap<>()); indexTemplate.setTemplateName(indexTemplateName); indexTemplate.setIndexPatterns(indexPatterns); objectUnderTest.putTemplate(indexTemplate); - verify(openSearchIndicesClient).putIndexTemplate(putIndexTemplateRequestArgumentCaptor.capture()); + verify(openSearchTransport).performRequest( + putIndexTemplateRequestArgumentCaptor.capture(), + endpointArgumentCaptor.capture(), + eq(openSearchTransportOptions) + ); + + Map actualPutRequest = putIndexTemplateRequestArgumentCaptor.getValue(); + + assertThat(actualPutRequest.get("index_patterns"), equalTo(indexPatterns)); - final PutIndexTemplateRequest actualPutRequest = putIndexTemplateRequestArgumentCaptor.getValue(); + Endpoint, PutIndexTemplateResponse, ErrorResponse> actualEndpoint = endpointArgumentCaptor.getValue(); + assertThat(actualEndpoint.method(null), equalTo("PUT")); + assertThat(actualEndpoint.requestUrl(null), equalTo("/_index_template/" + indexTemplateName)); + assertThat(actualEndpoint.queryParameters(null), equalTo(Collections.emptyMap())); + assertThat(actualEndpoint.headers(null), equalTo(Collections.emptyMap())); + assertThat(actualEndpoint.hasRequestBody(), equalTo(true)); - assertThat(actualPutRequest.name(), equalTo(indexTemplateName)); - assertThat(actualPutRequest.indexPatterns(), equalTo(indexPatterns)); + assertThat(actualEndpoint, instanceOf(JsonEndpoint.class)); + assertThat(((JsonEndpoint)actualEndpoint).responseDeserializer(), equalTo(PutIndexTemplateResponse._DESERIALIZER)); - assertThat(actualPutRequest.version(), nullValue()); - assertThat(actualPutRequest.template(), nullValue()); - assertThat(actualPutRequest.priority(), nullValue()); - assertThat(actualPutRequest.composedOf(), notNullValue()); - assertThat(actualPutRequest.composedOf(), equalTo(Collections.emptyList())); + assertThat(actualPutRequest, not(hasKey("template"))); + assertThat(actualPutRequest, not(hasKey("priority"))); + assertThat(actualPutRequest, not(hasKey("composedOf"))); + assertThat(actualPutRequest, not(hasKey("template"))); } @Test @@ -205,7 +247,6 @@ void putTemplate_with_defined_template_values_performs_putIndexTemplate_request( final String numberOfShards = Integer.toString(random.nextInt(1000) + 100); final List composedOf = Collections.singletonList(UUID.randomUUID().toString()); - final IndexConfiguration indexConfiguration = mock(IndexConfiguration.class); final IndexTemplate indexTemplate = new ComposableIndexTemplate( Map.of("version", version, "priority", priority, @@ -220,21 +261,40 @@ void putTemplate_with_defined_template_values_performs_putIndexTemplate_request( indexTemplate.setIndexPatterns(indexPatterns); objectUnderTest.putTemplate(indexTemplate); - verify(openSearchIndicesClient).putIndexTemplate(putIndexTemplateRequestArgumentCaptor.capture()); - - final PutIndexTemplateRequest actualPutRequest = putIndexTemplateRequestArgumentCaptor.getValue(); - - assertThat(actualPutRequest.name(), equalTo(indexTemplateName)); - assertThat(actualPutRequest.indexPatterns(), equalTo(indexPatterns)); - assertThat(actualPutRequest.version(), equalTo(version)); - assertThat(actualPutRequest.priority(), equalTo(priority)); - assertThat(actualPutRequest.composedOf(), equalTo(composedOf)); - assertThat(actualPutRequest.template(), notNullValue()); - assertThat(actualPutRequest.template().mappings(), notNullValue()); - assertThat(actualPutRequest.template().mappings().dateDetection(), equalTo(true)); - assertThat(actualPutRequest.template().settings(), notNullValue()); - assertThat(actualPutRequest.template().settings().index(), notNullValue()); - assertThat(actualPutRequest.template().settings().index().numberOfShards(), equalTo(numberOfShards)); + verify(openSearchTransport).performRequest( + putIndexTemplateRequestArgumentCaptor.capture(), + endpointArgumentCaptor.capture(), + eq(openSearchTransportOptions) + ); + + final Map actualPutRequest = putIndexTemplateRequestArgumentCaptor.getValue(); + + assertThat(actualPutRequest.get("index_patterns"), equalTo(indexPatterns)); + + final Endpoint, PutIndexTemplateResponse, ErrorResponse> actualEndpoint = endpointArgumentCaptor.getValue(); + assertThat(actualEndpoint.method(null), equalTo("PUT")); + assertThat(actualEndpoint.requestUrl(null), equalTo("/_index_template/" + indexTemplateName)); + assertThat(actualEndpoint.queryParameters(null), equalTo(Collections.emptyMap())); + assertThat(actualEndpoint.headers(null), equalTo(Collections.emptyMap())); + assertThat(actualEndpoint.hasRequestBody(), equalTo(true)); + + assertThat(actualEndpoint, instanceOf(JsonEndpoint.class)); + assertThat(((JsonEndpoint)actualEndpoint).responseDeserializer(), equalTo(PutIndexTemplateResponse._DESERIALIZER)); + + assertThat(actualPutRequest, hasKey("template")); + assertThat(actualPutRequest.get("version"), equalTo(version)); + assertThat(actualPutRequest.get("priority"), equalTo(priority)); + assertThat(actualPutRequest.get("composed_of"), equalTo(composedOf)); + assertThat(actualPutRequest.get("template"), notNullValue()); + assertThat(actualPutRequest.get("template"), instanceOf(Map.class)); + Map actualTemplate = (Map) actualPutRequest.get("template"); + assertThat(actualTemplate.get("mappings"), instanceOf(Map.class)); + assertThat(((Map) actualTemplate.get("mappings")).get("date_detection"), equalTo(true)); + assertThat(actualTemplate.get("settings"), instanceOf(Map.class)); + Map actualSettings = (Map) actualTemplate.get("settings"); + assertThat(actualSettings.get("index"), instanceOf(Map.class)); + Map actualIndex = (Map) actualSettings.get("index"); + assertThat(actualIndex.get("number_of_shards"), equalTo(numberOfShards)); } } } \ No newline at end of file From 541a2a7798f97c5be563762853609f3453b4d35f Mon Sep 17 00:00:00 2001 From: kkondaka <41027584+kkondaka@users.noreply.github.com> Date: Wed, 8 Nov 2023 07:51:55 -0800 Subject: [PATCH 37/53] Change Latency Metric names (#3603) Signed-off-by: Krishna Kondaka Co-authored-by: Krishna Kondaka --- .../opensearch/dataprepper/model/sink/SinkLatencyMetrics.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetrics.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetrics.java index 31b584271f..3a39c75b96 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetrics.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/sink/SinkLatencyMetrics.java @@ -13,8 +13,8 @@ import java.time.Instant; public class SinkLatencyMetrics { - public static final String INTERNAL_LATENCY = "internalLatency"; - public static final String EXTERNAL_LATENCY = "externalLatency"; + public static final String INTERNAL_LATENCY = "PipelineLatency"; + public static final String EXTERNAL_LATENCY = "EndToEndLatency"; private final DistributionSummary internalLatencySummary; private final DistributionSummary externalLatencySummary; From 803671a3f2ae62ded315f66862ff5ff9b30c19ef Mon Sep 17 00:00:00 2001 From: David Venable Date: Wed, 8 Nov 2023 08:45:39 -0800 Subject: [PATCH 38/53] Some changes to the configurations for the new visibility timeout extension feature. Increase the maximum default to 2 hours, increase the maximum configurable extension to 24 hours, and rename to use the work "maximum" to remain consistent (e.g. maximum_messages). (#3604) Signed-off-by: David Venable --- .../plugins/source/s3/configuration/SqsOptions.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/configuration/SqsOptions.java b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/configuration/SqsOptions.java index c4acd3abfd..40aeb6ffd6 100644 --- a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/configuration/SqsOptions.java +++ b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/configuration/SqsOptions.java @@ -18,7 +18,7 @@ public class SqsOptions { private static final int DEFAULT_MAXIMUM_MESSAGES = 10; private static final Boolean DEFAULT_VISIBILITY_DUPLICATE_PROTECTION = false; private static final Duration DEFAULT_VISIBILITY_TIMEOUT_SECONDS = Duration.ofSeconds(30); - private static final Duration DEFAULT_MAX_VISIBILITY_TIMEOUT_EXTENSION = Duration.ofSeconds(1800); // 30 minutes + private static final Duration DEFAULT_MAX_VISIBILITY_TIMEOUT_EXTENSION = Duration.ofHours(2); private static final Duration DEFAULT_WAIT_TIME_SECONDS = Duration.ofSeconds(20); private static final Duration DEFAULT_POLL_DELAY_SECONDS = Duration.ofSeconds(0); @@ -39,9 +39,9 @@ public class SqsOptions { @JsonProperty("visibility_duplication_protection") private Boolean visibilityDuplicateProtection = DEFAULT_VISIBILITY_DUPLICATE_PROTECTION; - @JsonProperty("max_visibility_timeout_extension") + @JsonProperty("maximum_visibility_timeout_extension") @DurationMin(seconds = 30) - @DurationMax(seconds = 3600) + @DurationMax(hours = 24) private Duration maxVisibilityTimeoutExtension = DEFAULT_MAX_VISIBILITY_TIMEOUT_EXTENSION; @JsonProperty("wait_time") From c8bf0ffb8fd96a99ad7c18a1bee5a998baa1de94 Mon Sep 17 00:00:00 2001 From: Chase <62891993+engechas@users.noreply.github.com> Date: Wed, 8 Nov 2023 11:13:45 -0600 Subject: [PATCH 39/53] Check topic for emptiness during KafkaBuffer shutdown (#3545) * Add shutdown method to buffer API Signed-off-by: Chase Engelbrecht * Add unit tests Signed-off-by: Chase Engelbrecht * Only allow single thread to check emptiness Signed-off-by: Chase Engelbrecht * Fix logic and add 1 minute wait before requerying kafka Signed-off-by: Chase Engelbrecht * Add unit tests for thread safety logic Signed-off-by: Chase Engelbrecht * Refactor metadata related to emptiness into own class Signed-off-by: Chase Engelbrecht * Single instance per topic rather than per worker Signed-off-by: Chase Engelbrecht * Add unit tests for topic emptiness class Signed-off-by: Chase Engelbrecht * Resolve rebase conflicts Signed-off-by: Chase Engelbrecht --------- Signed-off-by: Chase Engelbrecht --- .../plugins/kafka/buffer/KafkaBuffer.java | 9 +- .../kafka/consumer/KafkaCustomConsumer.java | 44 +++- .../consumer/KafkaCustomConsumerFactory.java | 3 +- .../consumer/TopicEmptinessMetadata.java | 55 +++++ .../plugins/kafka/source/KafkaSource.java | 5 +- .../plugins/kafka/buffer/KafkaBufferTest.java | 111 +++++++++- .../consumer/KafkaCustomConsumerTest.java | 191 +++++++++++++++++- .../consumer/TopicEmptinessMetadataTest.java | 101 +++++++++ 8 files changed, 510 insertions(+), 9 deletions(-) create mode 100644 data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/TopicEmptinessMetadata.java create mode 100644 data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/TopicEmptinessMetadataTest.java 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 3a358c23fe..f8fc8337d2 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 @@ -46,6 +46,7 @@ public class KafkaBuffer> extends AbstractBuffer { static final String WRITE = "Write"; static final String READ = "Read"; private final KafkaCustomProducer producer; + private final List emptyCheckingConsumers; private final AbstractBuffer innerBuffer; private final ExecutorService executorService; private final Duration drainTimeout; @@ -69,6 +70,8 @@ public KafkaBuffer(final PluginSetting pluginSetting, final KafkaBufferConfig ka final PluginMetrics consumerMetrics = PluginMetrics.fromNames(metricPrefixName + READ, pluginSetting.getPipelineName()); final List consumers = kafkaCustomConsumerFactory.createConsumersForTopic(kafkaBufferConfig, kafkaBufferConfig.getTopic(), innerBuffer, consumerMetrics, acknowledgementSetManager, byteDecoder, shutdownInProgress, false); + emptyCheckingConsumers = kafkaCustomConsumerFactory.createConsumersForTopic(kafkaBufferConfig, kafkaBufferConfig.getTopic(), + innerBuffer, pluginMetrics, acknowledgementSetManager, byteDecoder, shutdownInProgress, false); this.executorService = Executors.newFixedThreadPool(consumers.size()); consumers.forEach(this.executorService::submit); @@ -124,8 +127,10 @@ public void doCheckpoint(CheckpointState checkpointState) { @Override public boolean isEmpty() { - // TODO: check Kafka topic is empty as well. - return innerBuffer.isEmpty(); + final boolean areTopicsEmpty = emptyCheckingConsumers.stream() + .allMatch(KafkaCustomConsumer::isTopicEmpty); + + return areTopicsEmpty && innerBuffer.isEmpty(); } @Override 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 d8d577e49a..73f0fe1278 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 @@ -16,6 +16,7 @@ import org.apache.kafka.clients.consumer.ConsumerRecords; import org.apache.kafka.clients.consumer.KafkaConsumer; import org.apache.kafka.clients.consumer.OffsetAndMetadata; +import org.apache.kafka.common.PartitionInfo; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.errors.AuthenticationException; import org.apache.kafka.common.errors.RecordDeserializationException; @@ -55,6 +56,7 @@ import java.util.Set; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; import static org.apache.commons.lang3.exception.ExceptionUtils.getRootCause; @@ -92,6 +94,7 @@ public class KafkaCustomConsumer implements Runnable, ConsumerRebalanceListener private long numRecordsCommitted = 0; private final LogRateLimiter errLogRateLimiter; private final ByteDecoder byteDecoder; + private final TopicEmptinessMetadata topicEmptinessMetadata; public KafkaCustomConsumer(final KafkaConsumer consumer, final AtomicBoolean shutdownInProgress, @@ -101,7 +104,8 @@ public KafkaCustomConsumer(final KafkaConsumer consumer, final String schemaType, final AcknowledgementSetManager acknowledgementSetManager, final ByteDecoder byteDecoder, - final KafkaTopicConsumerMetrics topicMetrics) { + final KafkaTopicConsumerMetrics topicMetrics, + final TopicEmptinessMetadata topicEmptinessMetadata) { this.topicName = topicConfig.getName(); this.topicConfig = topicConfig; this.shutdownInProgress = shutdownInProgress; @@ -125,6 +129,7 @@ public KafkaCustomConsumer(final KafkaConsumer consumer, this.lastCommitTime = System.currentTimeMillis(); this.numberOfAcksPending = new AtomicInteger(0); this.errLogRateLimiter = new LogRateLimiter(2, System.currentTimeMillis()); + this.topicEmptinessMetadata = topicEmptinessMetadata; } KafkaTopicConsumerMetrics getTopicMetrics() { @@ -528,4 +533,41 @@ final String getTopicPartitionOffset(final Map offsetMap, final Long offset = offsetMap.get(topicPartition); return Objects.isNull(offset) ? "-" : offset.toString(); } + + public synchronized boolean isTopicEmpty() { + final long currentThreadId = Thread.currentThread().getId(); + if (Objects.isNull(topicEmptinessMetadata.getTopicEmptyCheckingOwnerThreadId())) { + topicEmptinessMetadata.setTopicEmptyCheckingOwnerThreadId(currentThreadId); + } + + if (currentThreadId != topicEmptinessMetadata.getTopicEmptyCheckingOwnerThreadId() || + topicEmptinessMetadata.isWithinCheckInterval(System.currentTimeMillis())) { + return topicEmptinessMetadata.isTopicEmpty(); + } + + final List partitions = consumer.partitionsFor(topicName); + final List topicPartitions = partitions.stream() + .map(partitionInfo -> new TopicPartition(topicName, partitionInfo.partition())) + .collect(Collectors.toList()); + + final Map committedOffsets = consumer.committed(new HashSet<>(topicPartitions)); + final Map endOffsets = consumer.endOffsets(topicPartitions); + + for (TopicPartition topicPartition : topicPartitions) { + final OffsetAndMetadata offsetAndMetadata = committedOffsets.get(topicPartition); + final Long endOffset = endOffsets.get(topicPartition); + topicEmptinessMetadata.updateTopicEmptinessStatus(topicPartition, true); + + // If there is data in the partition + if (endOffset != 0L) { + // If there is no committed offset for the partition or the committed offset is behind the end offset + if (Objects.isNull(offsetAndMetadata) || offsetAndMetadata.offset() < endOffset) { + topicEmptinessMetadata.updateTopicEmptinessStatus(topicPartition, false); + } + } + } + + topicEmptinessMetadata.setLastIsEmptyCheckTime(System.currentTimeMillis()); + return topicEmptinessMetadata.isTopicEmpty(); + } } 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 5954e91c16..89132ea9b3 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 @@ -87,6 +87,7 @@ public List createConsumersForTopic(final KafkaConsumerConf try { final int numWorkers = topic.getWorkers(); + final TopicEmptinessMetadata topicEmptinessMetadata = new TopicEmptinessMetadata(); IntStream.range(0, numWorkers).forEach(index -> { KafkaDataConfig dataConfig = new KafkaDataConfigAdapter(keyFactory, topic); Deserializer keyDeserializer = (Deserializer) serializationFactory.getDeserializer(PlaintextKafkaDataConfig.plaintextDataConfig(dataConfig)); @@ -100,7 +101,7 @@ public List createConsumersForTopic(final KafkaConsumerConf final KafkaConsumer kafkaConsumer = new KafkaConsumer<>(consumerProperties, keyDeserializer, valueDeserializer); consumers.add(new KafkaCustomConsumer(kafkaConsumer, shutdownInProgress, buffer, kafkaConsumerConfig, topic, - schemaType, acknowledgementSetManager, byteDecoder, topicMetrics)); + schemaType, acknowledgementSetManager, byteDecoder, topicMetrics, topicEmptinessMetadata)); }); } catch (Exception e) { diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/TopicEmptinessMetadata.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/TopicEmptinessMetadata.java new file mode 100644 index 0000000000..bec9a17c91 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/TopicEmptinessMetadata.java @@ -0,0 +1,55 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.kafka.consumer; + +import org.apache.kafka.common.TopicPartition; + +import java.util.concurrent.ConcurrentHashMap; + +public class TopicEmptinessMetadata { + static final long IS_EMPTY_CHECK_INTERVAL_MS = 60000L; + + private long lastIsEmptyCheckTime; + private Long topicEmptyCheckingOwnerThreadId; + private ConcurrentHashMap topicPartitionToIsEmpty; + + public TopicEmptinessMetadata() { + this.lastIsEmptyCheckTime = 0; + this.topicEmptyCheckingOwnerThreadId = null; + this.topicPartitionToIsEmpty = new ConcurrentHashMap<>(); + } + + public void setLastIsEmptyCheckTime(final long timestamp) { + this.lastIsEmptyCheckTime = timestamp; + } + + public void setTopicEmptyCheckingOwnerThreadId(final Long threadId) { + this.topicEmptyCheckingOwnerThreadId = threadId; + } + + public void updateTopicEmptinessStatus(final TopicPartition topicPartition, final Boolean isEmpty) { + topicPartitionToIsEmpty.put(topicPartition, isEmpty); + } + + public long getLastIsEmptyCheckTime() { + return this.lastIsEmptyCheckTime; + } + + public Long getTopicEmptyCheckingOwnerThreadId() { + return this.topicEmptyCheckingOwnerThreadId; + } + + public ConcurrentHashMap getTopicPartitionToIsEmpty() { + return this.topicPartitionToIsEmpty; + } + + public boolean isTopicEmpty() { + return topicPartitionToIsEmpty.values().stream().allMatch(isEmpty -> isEmpty); + } + + public boolean isWithinCheckInterval(final long epochTimestamp) { + return epochTimestamp < lastIsEmptyCheckTime + IS_EMPTY_CHECK_INTERVAL_MS; + } +} 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 5d625439cf..aa52a1fd1a 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 @@ -37,6 +37,7 @@ import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaRegistryType; import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.opensearch.dataprepper.plugins.kafka.consumer.KafkaCustomConsumer; +import org.opensearch.dataprepper.plugins.kafka.consumer.TopicEmptinessMetadata; import org.opensearch.dataprepper.plugins.kafka.extension.KafkaClusterConfigSupplier; import org.opensearch.dataprepper.plugins.kafka.util.ClientDNSLookupType; import org.opensearch.dataprepper.plugins.kafka.util.KafkaSecurityConfigurer; @@ -116,6 +117,7 @@ public void start(Buffer> buffer) { int numWorkers = topic.getWorkers(); executorService = Executors.newFixedThreadPool(numWorkers); allTopicExecutorServices.add(executorService); + final TopicEmptinessMetadata topicEmptinessMetadata = new TopicEmptinessMetadata(); IntStream.range(0, numWorkers).forEach(index -> { while (true) { @@ -138,7 +140,8 @@ public void start(Buffer> buffer) { } } - consumer = new KafkaCustomConsumer(kafkaConsumer, shutdownInProgress, buffer, sourceConfig, topic, schemaType, acknowledgementSetManager, null, topicMetrics); + consumer = new KafkaCustomConsumer(kafkaConsumer, shutdownInProgress, buffer, sourceConfig, topic, schemaType, + acknowledgementSetManager, null, topicMetrics, topicEmptinessMetadata); allTopicConsumers.add(consumer); executorService.submit(consumer); 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 59c7641b76..ddcdb77edd 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 @@ -28,6 +28,8 @@ import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionType; import org.opensearch.dataprepper.plugins.kafka.configuration.PlainTextAuthConfig; +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.producer.ProducerWorker; @@ -36,6 +38,7 @@ import java.time.Duration; import java.util.Arrays; import java.util.Collections; +import java.util.List; import java.util.Objects; import java.util.Random; import java.util.UUID; @@ -59,7 +62,9 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mockConstruction; import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; import static org.opensearch.dataprepper.plugins.kafka.buffer.KafkaBuffer.EXECUTOR_SERVICE_SHUTDOWN_TIMEOUT; @@ -107,6 +112,12 @@ class KafkaBufferTest { @Mock KafkaCustomProducer producer; + @Mock + private KafkaCustomConsumerFactory consumerFactory; + + @Mock + private KafkaCustomConsumer consumer; + @Mock BlockingBuffer> blockingBuffer; @@ -114,7 +125,10 @@ class KafkaBufferTest { private AwsCredentialsSupplier awsCredentialsSupplier; public KafkaBuffer> createObjectUnderTest() { + return createObjectUnderTest(List.of(consumer)); + } + public KafkaBuffer> createObjectUnderTest(final List consumers) { try ( final MockedStatic executorsMockedStatic = mockStatic(Executors.class); final MockedConstruction producerFactoryMock = @@ -122,6 +136,11 @@ public KafkaBuffer> createObjectUnderTest() { producerFactory = mock; when(producerFactory.createProducer(any() ,any(), any(), isNull(), isNull(), any(), anyBoolean())).thenReturn(producer); }); + final MockedConstruction consumerFactoryMock = + mockConstruction(KafkaCustomConsumerFactory.class, (mock, context) -> { + consumerFactory = mock; + when(consumerFactory.createConsumersForTopic(any(), any(), any(), any(), any(), any(), any(), anyBoolean())).thenReturn(consumers); + }); final MockedConstruction blockingBufferMock = mockConstruction(BlockingBuffer.class, (mock, context) -> { blockingBuffer = mock; @@ -210,12 +229,100 @@ void test_kafkaBuffer_doWriteAll() throws Exception { } @Test - void test_kafkaBuffer_isEmpty() { + void test_kafkaBuffer_isEmpty_True() { + kafkaBuffer = createObjectUnderTest(); + assertTrue(Objects.nonNull(kafkaBuffer)); + when(blockingBuffer.isEmpty()).thenReturn(true); + when(consumer.isTopicEmpty()).thenReturn(true); + + final boolean result = kafkaBuffer.isEmpty(); + assertThat(result, equalTo(true)); + + verify(blockingBuffer).isEmpty(); + verify(consumer).isTopicEmpty(); + } + + @Test + void test_kafkaBuffer_isEmpty_BufferNotEmpty() { + kafkaBuffer = createObjectUnderTest(); + assertTrue(Objects.nonNull(kafkaBuffer)); + when(blockingBuffer.isEmpty()).thenReturn(false); + when(consumer.isTopicEmpty()).thenReturn(true); + + final boolean result = kafkaBuffer.isEmpty(); + assertThat(result, equalTo(false)); + + verify(blockingBuffer).isEmpty(); + verify(consumer).isTopicEmpty(); + } + + @Test + void test_kafkaBuffer_isEmpty_TopicNotEmpty() { kafkaBuffer = createObjectUnderTest(); assertTrue(Objects.nonNull(kafkaBuffer)); + when(blockingBuffer.isEmpty()).thenReturn(true); + when(consumer.isTopicEmpty()).thenReturn(false); + + final boolean result = kafkaBuffer.isEmpty(); + assertThat(result, equalTo(false)); + + verifyNoInteractions(blockingBuffer); + verify(consumer).isTopicEmpty(); + } + + @Test + void test_kafkaBuffer_isEmpty_MultipleTopics_AllNotEmpty() { + kafkaBuffer = createObjectUnderTest(List.of(consumer, consumer)); + assertTrue(Objects.nonNull(kafkaBuffer)); + when(blockingBuffer.isEmpty()).thenReturn(true); + when(consumer.isTopicEmpty()).thenReturn(false).thenReturn(false); + + final boolean result = kafkaBuffer.isEmpty(); + assertThat(result, equalTo(false)); + + verifyNoInteractions(blockingBuffer); + verify(consumer).isTopicEmpty(); + } + + @Test + void test_kafkaBuffer_isEmpty_MultipleTopics_SomeNotEmpty() { + kafkaBuffer = createObjectUnderTest(List.of(consumer, consumer)); + assertTrue(Objects.nonNull(kafkaBuffer)); + when(blockingBuffer.isEmpty()).thenReturn(true); + when(consumer.isTopicEmpty()).thenReturn(true).thenReturn(false); + + final boolean result = kafkaBuffer.isEmpty(); + assertThat(result, equalTo(false)); + + verifyNoInteractions(blockingBuffer); + verify(consumer, times(2)).isTopicEmpty(); + } + + @Test + void test_kafkaBuffer_isEmpty_MultipleTopics_AllEmpty() { + kafkaBuffer = createObjectUnderTest(List.of(consumer, consumer)); + assertTrue(Objects.nonNull(kafkaBuffer)); + when(blockingBuffer.isEmpty()).thenReturn(true); + when(consumer.isTopicEmpty()).thenReturn(true).thenReturn(true); + + final boolean result = kafkaBuffer.isEmpty(); + assertThat(result, equalTo(true)); + + verify(blockingBuffer).isEmpty(); + verify(consumer, times(2)).isTopicEmpty(); + } + + @Test + void test_kafkaBuffer_isEmpty_ZeroTopics() { + kafkaBuffer = createObjectUnderTest(Collections.emptyList()); + assertTrue(Objects.nonNull(kafkaBuffer)); + when(blockingBuffer.isEmpty()).thenReturn(true); + + final boolean result = kafkaBuffer.isEmpty(); + assertThat(result, equalTo(true)); - kafkaBuffer.isEmpty(); verify(blockingBuffer).isEmpty(); + verifyNoInteractions(consumer); } @Test 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 d707cbf82d..8c7675795c 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 @@ -14,6 +14,7 @@ import org.apache.kafka.clients.consumer.ConsumerRecords; import org.apache.kafka.clients.consumer.KafkaConsumer; import org.apache.kafka.clients.consumer.OffsetAndMetadata; +import org.apache.kafka.common.PartitionInfo; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.errors.RecordDeserializationException; import org.junit.jupiter.api.Assertions; @@ -42,25 +43,38 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; +import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.concurrent.ScheduledExecutorService; +import java.util.Random; import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.stream.Collectors; +import java.util.stream.IntStream; import static org.awaitility.Awaitility.await; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; 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.anyCollection; +import static org.mockito.ArgumentMatchers.anySet; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; @ExtendWith(MockitoExtension.class) @MockitoSettings(strictness = Strictness.LENIENT) public class KafkaCustomConsumerTest { + private static final String TOPIC_NAME = "topic1"; + private static final Random RANDOM = new Random(); @Mock private KafkaConsumer kafkaConsumer; @@ -80,6 +94,10 @@ public class KafkaCustomConsumerTest { @Mock private KafkaTopicConsumerMetrics topicMetrics; + @Mock + private PartitionInfo partitionInfo; + @Mock + private OffsetAndMetadata offsetAndMetadata; private KafkaCustomConsumer consumer; @@ -105,6 +123,7 @@ public class KafkaCustomConsumerTest { private Duration delayTime; private double posCount; private double negCount; + private TopicEmptinessMetadata topicEmptinessMetadata; @BeforeEach public void setUp() { @@ -140,12 +159,14 @@ public void setUp() { sourceConfig = mock(KafkaConsumerConfig.class); buffer = getBuffer(); shutdownInProgress = new AtomicBoolean(false); - when(topicConfig.getName()).thenReturn("topic1"); + when(topicConfig.getName()).thenReturn(TOPIC_NAME); } public KafkaCustomConsumer createObjectUnderTest(String schemaType, boolean acknowledgementsEnabled) { + topicEmptinessMetadata = new TopicEmptinessMetadata(); when(sourceConfig.getAcknowledgementsEnabled()).thenReturn(acknowledgementsEnabled); - return new KafkaCustomConsumer(kafkaConsumer, shutdownInProgress, buffer, sourceConfig, topicConfig, schemaType, acknowledgementSetManager, null, topicMetrics); + return new KafkaCustomConsumer(kafkaConsumer, shutdownInProgress, buffer, sourceConfig, topicConfig, schemaType, + acknowledgementSetManager, null, topicMetrics, topicEmptinessMetadata); } private BlockingBuffer> getBuffer() { @@ -454,6 +475,172 @@ public void testAwsGlueErrorWithAcknowledgements() throws Exception { }); } + @Test + public void isTopicEmpty_OnePartition_IsEmpty() { + final Long offset = RANDOM.nextLong(); + final List topicPartitions = buildTopicPartitions(1); + + consumer = createObjectUnderTest("json", true); + + when(kafkaConsumer.partitionsFor(TOPIC_NAME)).thenReturn(List.of(partitionInfo)); + when(partitionInfo.partition()).thenReturn(0); + when(kafkaConsumer.committed(anySet())).thenReturn(getTopicPartitionToMap(topicPartitions, offsetAndMetadata)); + when(kafkaConsumer.endOffsets(anyCollection())).thenReturn(getTopicPartitionToMap(topicPartitions, offset)); + when(offsetAndMetadata.offset()).thenReturn(offset); + + assertThat(consumer.isTopicEmpty(), equalTo(true)); + + verify(kafkaConsumer).partitionsFor(TOPIC_NAME); + verify(kafkaConsumer).committed(new HashSet<>(topicPartitions)); + verify(kafkaConsumer).endOffsets(topicPartitions); + verify(partitionInfo).partition(); + verify(offsetAndMetadata).offset(); + } + + @Test + public void isTopicEmpty_OnePartition_PartitionNeverHadData() { + final Long offset = 0L; + final List topicPartitions = buildTopicPartitions(1); + + consumer = createObjectUnderTest("json", true); + + when(kafkaConsumer.partitionsFor(TOPIC_NAME)).thenReturn(List.of(partitionInfo)); + when(partitionInfo.partition()).thenReturn(0); + when(kafkaConsumer.committed(anySet())).thenReturn(getTopicPartitionToMap(topicPartitions, offsetAndMetadata)); + when(kafkaConsumer.endOffsets(anyCollection())).thenReturn(getTopicPartitionToMap(topicPartitions, offset)); + when(offsetAndMetadata.offset()).thenReturn(offset - 1); + + assertThat(consumer.isTopicEmpty(), equalTo(true)); + + verify(kafkaConsumer).partitionsFor(TOPIC_NAME); + verify(kafkaConsumer).committed(new HashSet<>(topicPartitions)); + verify(kafkaConsumer).endOffsets(topicPartitions); + verify(partitionInfo).partition(); + } + + @Test + public void isTopicEmpty_OnePartition_IsNotEmpty() { + final Long offset = RANDOM.nextLong(); + final List topicPartitions = buildTopicPartitions(1); + + consumer = createObjectUnderTest("json", true); + + when(kafkaConsumer.partitionsFor(TOPIC_NAME)).thenReturn(List.of(partitionInfo)); + when(partitionInfo.partition()).thenReturn(0); + when(kafkaConsumer.committed(anySet())).thenReturn(getTopicPartitionToMap(topicPartitions, offsetAndMetadata)); + when(kafkaConsumer.endOffsets(anyCollection())).thenReturn(getTopicPartitionToMap(topicPartitions, offset)); + when(offsetAndMetadata.offset()).thenReturn(offset - 1); + + assertThat(consumer.isTopicEmpty(), equalTo(false)); + + verify(kafkaConsumer).partitionsFor(TOPIC_NAME); + verify(kafkaConsumer).committed(new HashSet<>(topicPartitions)); + verify(kafkaConsumer).endOffsets(topicPartitions); + verify(partitionInfo).partition(); + verify(offsetAndMetadata).offset(); + } + + @Test + public void isTopicEmpty_OnePartition_NoCommittedPartition() { + final Long offset = RANDOM.nextLong(); + final List topicPartitions = buildTopicPartitions(1); + + consumer = createObjectUnderTest("json", true); + + when(kafkaConsumer.partitionsFor(TOPIC_NAME)).thenReturn(List.of(partitionInfo)); + when(partitionInfo.partition()).thenReturn(0); + when(kafkaConsumer.committed(anySet())).thenReturn(Collections.emptyMap()); + when(kafkaConsumer.endOffsets(anyCollection())).thenReturn(getTopicPartitionToMap(topicPartitions, offset)); + + assertThat(consumer.isTopicEmpty(), equalTo(false)); + + verify(kafkaConsumer).partitionsFor(TOPIC_NAME); + verify(kafkaConsumer).committed(new HashSet<>(topicPartitions)); + verify(kafkaConsumer).endOffsets(topicPartitions); + verify(partitionInfo).partition(); + } + + @Test + public void isTopicEmpty_MultiplePartitions_AllEmpty() { + final Long offset1 = RANDOM.nextLong(); + final Long offset2 = RANDOM.nextLong(); + final List topicPartitions = buildTopicPartitions(2); + + consumer = createObjectUnderTest("json", true); + + when(kafkaConsumer.partitionsFor(TOPIC_NAME)).thenReturn(List.of(partitionInfo, partitionInfo)); + when(partitionInfo.partition()).thenReturn(0).thenReturn(1); + when(kafkaConsumer.committed(anySet())).thenReturn(getTopicPartitionToMap(topicPartitions, offsetAndMetadata)); + final Map endOffsets = getTopicPartitionToMap(topicPartitions, offset1); + endOffsets.put(topicPartitions.get(1), offset2); + when(kafkaConsumer.endOffsets(anyCollection())).thenReturn(endOffsets); + when(offsetAndMetadata.offset()).thenReturn(offset1).thenReturn(offset2); + + assertThat(consumer.isTopicEmpty(), equalTo(true)); + + verify(kafkaConsumer).partitionsFor(TOPIC_NAME); + verify(kafkaConsumer).committed(new HashSet<>(topicPartitions)); + verify(kafkaConsumer).endOffsets(topicPartitions); + verify(partitionInfo, times(2)).partition(); + verify(offsetAndMetadata, times(2)).offset(); + } + + @Test + public void isTopicEmpty_MultiplePartitions_OneNotEmpty() { + final Long offset1 = RANDOM.nextLong(); + final Long offset2 = RANDOM.nextLong(); + final List topicPartitions = buildTopicPartitions(2); + + consumer = createObjectUnderTest("json", true); + + when(kafkaConsumer.partitionsFor(TOPIC_NAME)).thenReturn(List.of(partitionInfo, partitionInfo)); + when(partitionInfo.partition()).thenReturn(0).thenReturn(1); + when(kafkaConsumer.committed(anySet())).thenReturn(getTopicPartitionToMap(topicPartitions, offsetAndMetadata)); + final Map endOffsets = getTopicPartitionToMap(topicPartitions, offset1); + endOffsets.put(topicPartitions.get(1), offset2); + when(kafkaConsumer.endOffsets(anyCollection())).thenReturn(endOffsets); + when(offsetAndMetadata.offset()).thenReturn(offset1).thenReturn(offset2 - 1); + + assertThat(consumer.isTopicEmpty(), equalTo(false)); + + verify(kafkaConsumer).partitionsFor(TOPIC_NAME); + verify(kafkaConsumer).committed(new HashSet<>(topicPartitions)); + verify(kafkaConsumer).endOffsets(topicPartitions); + verify(partitionInfo, times(2)).partition(); + verify(offsetAndMetadata, times(2)).offset(); + } + + @Test + public void isTopicEmpty_NonCheckerThread_ShortCircuits() { + consumer = createObjectUnderTest("json", true); + + topicEmptinessMetadata.setTopicEmptyCheckingOwnerThreadId(Thread.currentThread().getId() - 1); + assertThat(consumer.isTopicEmpty(), equalTo(true)); + + verifyNoInteractions(kafkaConsumer); + } + + @Test + public void isTopicEmpty_CheckedWithinDelay_ShortCircuits() { + consumer = createObjectUnderTest("json", true); + + topicEmptinessMetadata.setLastIsEmptyCheckTime(System.currentTimeMillis()); + assertThat(consumer.isTopicEmpty(), equalTo(true)); + + verifyNoInteractions(kafkaConsumer); + } + + private List buildTopicPartitions(final int partitionCount) { + return IntStream.range(0, partitionCount) + .mapToObj(i -> new TopicPartition(TOPIC_NAME, i)) + .collect(Collectors.toList()); + } + + private Map getTopicPartitionToMap(final List topicPartitions, final T value) { + return topicPartitions.stream() + .collect(Collectors.toMap(i -> i, i -> value)); + } + private ConsumerRecords createPlainTextRecords(String topic, final long startOffset) { Map> records = new HashMap<>(); ConsumerRecord record1 = new ConsumerRecord<>(topic, testPartition, startOffset, testKey1, testValue1); diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/TopicEmptinessMetadataTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/TopicEmptinessMetadataTest.java new file mode 100644 index 0000000000..25f1950936 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/TopicEmptinessMetadataTest.java @@ -0,0 +1,101 @@ +package org.opensearch.dataprepper.plugins.kafka.consumer; + +import org.apache.kafka.common.TopicPartition; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.opensearch.dataprepper.plugins.kafka.consumer.TopicEmptinessMetadata.IS_EMPTY_CHECK_INTERVAL_MS; + +public class TopicEmptinessMetadataTest { + @Mock + private TopicPartition topicPartition; + @Mock + private TopicPartition topicPartition2; + + private TopicEmptinessMetadata topicEmptinessMetadata; + + @BeforeEach + void setup() { + MockitoAnnotations.openMocks(this); + this.topicEmptinessMetadata = new TopicEmptinessMetadata(); + } + + @Test + void updateTopicEmptinessStatus_AddEntry() { + topicEmptinessMetadata.updateTopicEmptinessStatus(topicPartition, false); + assertThat(topicEmptinessMetadata.getTopicPartitionToIsEmpty().containsKey(topicPartition), equalTo(true)); + assertThat(topicEmptinessMetadata.getTopicPartitionToIsEmpty().get(topicPartition), equalTo(false)); + } + + @Test + void updateTopicEmptinessStatus_UpdateEntry() { + topicEmptinessMetadata.updateTopicEmptinessStatus(topicPartition, false); + assertThat(topicEmptinessMetadata.getTopicPartitionToIsEmpty().containsKey(topicPartition), equalTo(true)); + assertThat(topicEmptinessMetadata.getTopicPartitionToIsEmpty().get(topicPartition), equalTo(false)); + + topicEmptinessMetadata.updateTopicEmptinessStatus(topicPartition, true); + assertThat(topicEmptinessMetadata.getTopicPartitionToIsEmpty().containsKey(topicPartition), equalTo(true)); + assertThat(topicEmptinessMetadata.getTopicPartitionToIsEmpty().get(topicPartition), equalTo(true)); + } + + @Test + void isTopicEmpty_NoItems() { + assertThat(topicEmptinessMetadata.isTopicEmpty(), equalTo(true)); + } + + @Test + void isTopicEmpty_OnePartition_IsNotEmpty() { + topicEmptinessMetadata.updateTopicEmptinessStatus(topicPartition, false); + assertThat(topicEmptinessMetadata.isTopicEmpty(), equalTo(false)); + } + + @Test + void isTopicEmpty_OnePartition_IsEmpty() { + topicEmptinessMetadata.updateTopicEmptinessStatus(topicPartition, true); + assertThat(topicEmptinessMetadata.isTopicEmpty(), equalTo(true)); + } + + @Test + void isTopicEmpty_MultiplePartitions_OneNotEmpty() { + topicEmptinessMetadata.updateTopicEmptinessStatus(topicPartition, true); + topicEmptinessMetadata.updateTopicEmptinessStatus(topicPartition2, false); + assertThat(topicEmptinessMetadata.isTopicEmpty(), equalTo(false)); + } + + @Test + void isCheckDurationExceeded_NoPreviousChecks() { + assertThat(topicEmptinessMetadata.isWithinCheckInterval(System.currentTimeMillis()), equalTo(false)); + } + + @Test + void isCheckDurationExceeded_CurrentTimeBeforeLastCheck() { + final long time = System.currentTimeMillis(); + topicEmptinessMetadata.setLastIsEmptyCheckTime(time); + assertThat(topicEmptinessMetadata.isWithinCheckInterval(time - 1), equalTo(true)); + } + + @Test + void isCheckDurationExceeded_CurrentTimeAfterLastCheck_BeforeInterval() { + final long time = System.currentTimeMillis(); + topicEmptinessMetadata.setLastIsEmptyCheckTime(time); + assertThat(topicEmptinessMetadata.isWithinCheckInterval((time + IS_EMPTY_CHECK_INTERVAL_MS) - 1), equalTo(true)); + } + + @Test + void isCheckDurationExceeded_CurrentTimeAfterLastCheck_AtInterval() { + final long time = System.currentTimeMillis(); + topicEmptinessMetadata.setLastIsEmptyCheckTime(time); + assertThat(topicEmptinessMetadata.isWithinCheckInterval(time + IS_EMPTY_CHECK_INTERVAL_MS), equalTo(false)); + } + + @Test + void isCheckDurationExceeded_CurrentTimeAfterLastCheck_AfterInterval() { + final long time = System.currentTimeMillis(); + topicEmptinessMetadata.setLastIsEmptyCheckTime(time); + assertThat(topicEmptinessMetadata.isWithinCheckInterval(time + IS_EMPTY_CHECK_INTERVAL_MS + 1), equalTo(false)); + } +} From 476227a948b63ad0901db7596c7976c9ede75532 Mon Sep 17 00:00:00 2001 From: Asif Sohail Mohammed Date: Wed, 8 Nov 2023 11:40:42 -0600 Subject: [PATCH 40/53] Updated visibility protection timeout (#3608) Signed-off-by: Asif Sohail Mohammed --- .../dataprepper/plugins/source/s3/SqsWorkerIT.java | 4 ++-- .../dataprepper/plugins/source/s3/SqsWorker.java | 2 +- .../plugins/source/s3/configuration/SqsOptions.java | 10 +++++----- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerIT.java b/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerIT.java index dc09cd989c..e2ede9ef8a 100644 --- a/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerIT.java +++ b/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerIT.java @@ -344,7 +344,7 @@ void processSqsMessages_with_acks_and_progress_check_callbacks(final int numberO }).when(s3Service).addS3Object(any(S3ObjectReference.class), any(AcknowledgementSet.class)); ScheduledExecutorService executor = Executors.newScheduledThreadPool(2); when(sqsOptions.getVisibilityTimeout()).thenReturn(Duration.ofSeconds(6)); - when(sqsOptions.getMaxVisibilityTimeoutExtension()).thenReturn(Duration.ofSeconds(60)); + when(sqsOptions.getVisibilityDuplicateProtectionTimeout()).thenReturn(Duration.ofSeconds(60)); when(sqsOptions.getVisibilityDuplicateProtection()).thenReturn(true); acknowledgementSetManager = new DefaultAcknowledgementSetManager(executor); final SqsWorker objectUnderTest = createObjectUnderTest(); @@ -419,7 +419,7 @@ void processSqsMessages_with_acks_and_progress_check_callbacks_expires(final int }).when(s3Service).addS3Object(any(S3ObjectReference.class), any(AcknowledgementSet.class)); ScheduledExecutorService executor = Executors.newScheduledThreadPool(2); when(sqsOptions.getVisibilityTimeout()).thenReturn(Duration.ofSeconds(6)); - when(sqsOptions.getMaxVisibilityTimeoutExtension()).thenReturn(Duration.ofSeconds(60)); + when(sqsOptions.getVisibilityDuplicateProtectionTimeout()).thenReturn(Duration.ofSeconds(60)); when(sqsOptions.getVisibilityDuplicateProtection()).thenReturn(true); acknowledgementSetManager = new DefaultAcknowledgementSetManager(executor); final SqsWorker objectUnderTest = createObjectUnderTest(); diff --git a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorker.java b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorker.java index a0f7a0bb16..a390c48d6f 100644 --- a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorker.java +++ b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorker.java @@ -235,7 +235,7 @@ && isEventBridgeEventTypeCreated(parsedMessage)) { List waitingForAcknowledgements = new ArrayList<>(); AcknowledgementSet acknowledgementSet = null; final int visibilityTimeout = (int)sqsOptions.getVisibilityTimeout().getSeconds(); - final int maxVisibilityTimeout = (int)sqsOptions.getMaxVisibilityTimeoutExtension().getSeconds(); + final int maxVisibilityTimeout = (int)sqsOptions.getVisibilityDuplicateProtectionTimeout().getSeconds(); final int progressCheckInterval = visibilityTimeout/2 - 1; if (endToEndAcknowledgementsEnabled) { int expiryTimeout = visibilityTimeout - 2; diff --git a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/configuration/SqsOptions.java b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/configuration/SqsOptions.java index 40aeb6ffd6..1242a6525b 100644 --- a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/configuration/SqsOptions.java +++ b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/configuration/SqsOptions.java @@ -18,7 +18,7 @@ public class SqsOptions { private static final int DEFAULT_MAXIMUM_MESSAGES = 10; private static final Boolean DEFAULT_VISIBILITY_DUPLICATE_PROTECTION = false; private static final Duration DEFAULT_VISIBILITY_TIMEOUT_SECONDS = Duration.ofSeconds(30); - private static final Duration DEFAULT_MAX_VISIBILITY_TIMEOUT_EXTENSION = Duration.ofHours(2); + private static final Duration DEFAULT_VISIBILITY_DUPLICATE_PROTECTION_TIMEOUT = Duration.ofHours(2); private static final Duration DEFAULT_WAIT_TIME_SECONDS = Duration.ofSeconds(20); private static final Duration DEFAULT_POLL_DELAY_SECONDS = Duration.ofSeconds(0); @@ -39,10 +39,10 @@ public class SqsOptions { @JsonProperty("visibility_duplication_protection") private Boolean visibilityDuplicateProtection = DEFAULT_VISIBILITY_DUPLICATE_PROTECTION; - @JsonProperty("maximum_visibility_timeout_extension") + @JsonProperty("visibility_duplicate_protection_timeout") @DurationMin(seconds = 30) @DurationMax(hours = 24) - private Duration maxVisibilityTimeoutExtension = DEFAULT_MAX_VISIBILITY_TIMEOUT_EXTENSION; + private Duration visibilityDuplicateProtectionTimeout = DEFAULT_VISIBILITY_DUPLICATE_PROTECTION_TIMEOUT; @JsonProperty("wait_time") @DurationMin(seconds = 0) @@ -65,8 +65,8 @@ public Duration getVisibilityTimeout() { return visibilityTimeout; } - public Duration getMaxVisibilityTimeoutExtension() { - return maxVisibilityTimeoutExtension; + public Duration getVisibilityDuplicateProtectionTimeout() { + return visibilityDuplicateProtectionTimeout; } public Boolean getVisibilityDuplicateProtection() { From b80b565583c998e43a809a12e43dca9f9a20f6f0 Mon Sep 17 00:00:00 2001 From: Adi Suresh Date: Wed, 8 Nov 2023 12:07:27 -0600 Subject: [PATCH 41/53] Move opensearch source into same module as opensearch sink. (#3607) This change is required to share code between the source and sink plugins. Signed-off-by: Adi Suresh --- .../opensearch-source/README.md | 195 ----------------- .../opensearch-source/build.gradle | 29 --- data-prepper-plugins/opensearch/README.md | 196 ++++++++++++++++++ data-prepper-plugins/opensearch/build.gradle | 10 +- .../source/opensearch/ClientRefresher.java | 0 .../OpenSearchIndexProgressState.java | 0 .../source/opensearch/OpenSearchService.java | 0 .../source/opensearch/OpenSearchSource.java | 0 .../OpenSearchSourceConfiguration.java | 0 .../AwsAuthenticationConfiguration.java | 0 .../ConnectionConfiguration.java | 0 .../IndexParametersConfiguration.java | 0 .../configuration/OpenSearchIndex.java | 0 .../SchedulingParameterConfiguration.java | 0 .../configuration/SearchConfiguration.java | 0 .../OpenSearchSourcePluginMetrics.java | 0 .../worker/NoSearchContextWorker.java | 0 ...nSearchIndexPartitionCreationSupplier.java | 0 .../source/opensearch/worker/PitWorker.java | 0 .../opensearch/worker/ScrollWorker.java | 0 .../opensearch/worker/SearchWorker.java | 0 .../opensearch/worker/WorkerCommonUtils.java | 0 .../worker/client/ClusterClientFactory.java | 0 .../worker/client/ElasticsearchAccessor.java | 0 .../worker/client/OpenSearchAccessor.java | 0 .../client/OpenSearchClientFactory.java | 0 .../worker/client/SearchAccessor.java | 0 .../worker/client/SearchAccessorStrategy.java | 0 .../worker/client/X509TrustAllManager.java | 0 .../exceptions/IndexNotFoundException.java | 0 .../SearchContextLimitException.java | 0 .../model/CreatePointInTimeRequest.java | 0 .../model/CreatePointInTimeResponse.java | 0 .../client/model/CreateScrollRequest.java | 0 .../client/model/CreateScrollResponse.java | 0 .../model/DeletePointInTimeRequest.java | 0 .../client/model/DeleteScrollRequest.java | 0 .../client/model/MetadataKeyAttributes.java | 0 .../model/NoSearchContextSearchRequest.java | 0 .../client/model/SearchContextType.java | 0 .../model/SearchPointInTimeRequest.java | 0 .../client/model/SearchScrollRequest.java | 0 .../client/model/SearchScrollResponse.java | 0 .../model/SearchWithSearchAfterResults.java | 0 .../worker/client/model/SortingOptions.java | 0 .../opensearch/ClientRefresherTest.java | 0 .../opensearch/OpenSearchServiceTest.java | 0 .../OpenSearchSourceConfigurationTest.java | 0 .../opensearch/OpenSearchSourceTest.java | 0 .../AwsAuthenticationConfigurationTest.java | 0 .../ConnectionConfigurationTest.java | 0 .../IndexParametersConfigurationTest.java | 0 .../configuration/OpenSearchIndexTest.java | 0 .../SchedulingParameterConfigurationTest.java | 0 .../SearchConfigurationTest.java | 0 .../worker/NoSearchContextWorkerTest.java | 0 .../opensearch/worker/PitWorkerTest.java | 0 .../opensearch/worker/ScrollWorkerTest.java | 0 .../worker/WorkerCommonUtilsTest.java | 0 .../client/ElasticsearchAccessorTest.java | 0 .../worker/client/OpenSearchAccessorTest.java | 0 .../client/OpenSearchClientFactoryTest.java | 0 ...rchIndexPartitionCreationSupplierTest.java | 0 .../client/SearchAccessStrategyTest.java | 0 settings.gradle | 1 - 65 files changed, 205 insertions(+), 226 deletions(-) delete mode 100644 data-prepper-plugins/opensearch-source/README.md delete mode 100644 data-prepper-plugins/opensearch-source/build.gradle rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/ClientRefresher.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchIndexProgressState.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchService.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSource.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSourceConfiguration.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/AwsAuthenticationConfiguration.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/ConnectionConfiguration.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/IndexParametersConfiguration.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/OpenSearchIndex.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SchedulingParameterConfiguration.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SearchConfiguration.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/metrics/OpenSearchSourcePluginMetrics.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorker.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/OpenSearchIndexPartitionCreationSupplier.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorker.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorker.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/SearchWorker.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtils.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/ClusterClientFactory.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/ElasticsearchAccessor.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchAccessor.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchClientFactory.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/SearchAccessor.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/SearchAccessorStrategy.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/X509TrustAllManager.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/exceptions/IndexNotFoundException.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/exceptions/SearchContextLimitException.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreatePointInTimeRequest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreatePointInTimeResponse.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreateScrollRequest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreateScrollResponse.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/DeletePointInTimeRequest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/DeleteScrollRequest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/MetadataKeyAttributes.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/NoSearchContextSearchRequest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchContextType.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchPointInTimeRequest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchScrollRequest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchScrollResponse.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchWithSearchAfterResults.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SortingOptions.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/ClientRefresherTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchServiceTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSourceConfigurationTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSourceTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/AwsAuthenticationConfigurationTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/ConnectionConfigurationTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/IndexParametersConfigurationTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/OpenSearchIndexTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SchedulingParameterConfigurationTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SearchConfigurationTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorkerTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorkerTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorkerTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtilsTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/ElasticsearchAccessorTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchAccessorTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchClientFactoryTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchIndexPartitionCreationSupplierTest.java (100%) rename data-prepper-plugins/{opensearch-source => opensearch}/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/SearchAccessStrategyTest.java (100%) diff --git a/data-prepper-plugins/opensearch-source/README.md b/data-prepper-plugins/opensearch-source/README.md deleted file mode 100644 index 9424bf5093..0000000000 --- a/data-prepper-plugins/opensearch-source/README.md +++ /dev/null @@ -1,195 +0,0 @@ -# OpenSearch Source - -This is the Date Prepper OpenSearch source plugin that processes indices for either OpenSearch, Elasticsearch, -or Amazon OpenSearch Service clusters. It is ideal for migrating index data from a cluster. - -Note: Only fully tested versions will be listed below. It is likely many more versions are supported already, but it is untested. - -The OpenSearch source is compatible with the following OpenSearch versions: -* 2.5 - -And is compatible with the following Elasticsearch versions: -* 7.10 - -# Usages - -### Minimum required config with username and password - -```yaml -opensearch-source-pipeline: - source: - opensearch: - connection: - insecure: true - hosts: [ "https://localhost:9200" ] - username: "username" - password: "password" -``` - -### Full config example - -```yaml -opensearch-source-pipeline: - source: - opensearch: - indices: - include: - - index_name_regex: "test-index-.*" - exclude: - - index_name_regex: "test-index-[1-9].*" - scheduling: - rate: "PT1H" - start_time: "2023-06-02T22:01:30.00Z" - job_count: 2 - search_options: - search_context_type: "none" - batch_size: 1000 - connection: - insecure: false - cert: "/path/to/cert.crt" - socket_timeout: "100ms" - connection_timeout: "100ms" - hosts: [ "https://localhost:9200" ] - username: "username" - password: "password" -``` - -### Amazon OpenSearch Service - -The OpenSearch source can also be configured for an Amazon OpenSearch Service domain. - -```yaml -opensearch-source-pipeline: - source: - opensearch: - connection: - insecure: true - hosts: [ "https://search-my-domain-soopywaovobopgs8ywurr3utsu.us-east-1.es.amazonaws.com" ] - aws: - region: "us-east-1" - sts_role_arn: "arn:aws:iam::123456789012:role/my-domain-role" -``` - -### Using Metadata - -When the OpenSearch source constructs Data Prepper Events from documents in the cluster, the -document index is stored in the `EventMetadata` with an `opensearch-index` key, and the document_id is -stored in the `EventMetadata` with a `opensearch-document_id` key. This allows conditional routing based on the index or document_id, -among other things. For example, one could send to an OpenSearch sink and use the same index and document_id from the source cluster in -the destination cluster. A full config example for this use case is below - -```yaml -opensearch-source-pipeline: - source: - opensearch: - connection: - insecure: true - hosts: [ "https://source-cluster:9200" ] - username: "username" - password: "password" - sink: - - opensearch: - hosts: [ "https://sink-cluster:9200" ] - username: "username" - password: "password" - document_id_field: "getMetadata(\"opensearch-document_id\")" - index: "${getMetadata(\"opensearch-index\"}" -``` - -## Configuration - -- `hosts` (Required) : A list of IP addresses of OpenSearch or Elasticsearch nodes. - - -- `username` (Optional) : A String of username used in the internal users of OpenSearch cluster. Default is null. - - -- `password` (Optional) : A String of password used in the internal users of OpenSearch cluster. Default is null. - - -- `disable_authentication` (Optional) : A boolean that can disable authentication if the cluster supports it. Defaults to false. - - -- `aws` (Optional) : AWS configurations. See [AWS Configuration](#aws_configuration) for details. SigV4 is enabled by default when this option is used. - - -- `search_options` (Optional) : See [Search Configuration](#search_configuration) for details - - -- `indices` (Optional): See [Indices Configurations](#indices_configuration) for filtering options. - - -- `scheduling` (Optional): See [Scheduling Configuration](#scheduling_configuration) for details - - -- `connection` (Optional): See [Connection Configuration](#connection_configuration) - -### AWS Configuration - -* `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 STS role to assume for requests to AWS. 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_header_overrides` (Optional): A map of header overrides to make when assuming the IAM role for the source plugin. - -### Search Configuration - -* `search_context_type` (Optional) : A direct override for which type of search context should be used to search documents. - Options include `point_in_time`, `scroll`, or `none` (just search after). - By default, the OpenSearch source will attempt to use `point_in_time` on a cluster by auto-detecting that the cluster version and distribution -supports Point in Time. If the cluster does not support `point_in_time`, then `scroll` is the default behavior. - - -* `batch_size` (Optional) : The amount of documents to read in at once while searching. -This size is passed to the search requests for all search context types (`none` (search_after), `point_in_time`, or `scroll`). -Defaults to 1,000. - -### Scheduling Configuration - -Schedule the start time and amount of times an index should be processed. For example, -a `rate` of `PT1H` and a `job_count` of 3 would result in each index getting processed 3 times, starting at `start_time` -and then every hour after the first time the index is processed. - -* `rate` (Optional) : A String that indicates the rate to process an index based on the `job_count`. -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 `job_count` is greater than 1. - - - -* `job_count` (Optional) : An Integer that specifies how many times each index should be processed. Defaults to 1. - - - -* `start_time` (Optional) : A String in the format of a timestamp that is compatible with Java Instant (i.e. `2023-06-02T22:01:30.00Z`). -Processing will be delayed until this timestamp is reached. The default start time is to start immediately. - -### Connection Configuration - -* `insecure` (Optional): A boolean flag to turn off SSL certificate verification. If set to true, CA certificate verification will be turned off and insecure HTTP requests will be sent. Default to false. - - -* `cert` (Optional) : CA certificate that is pem encoded. Accepts both .pem or .crt. This enables the client to trust the CA that has signed the certificate that the OpenSearch cluster is using. Default is null. - - -* `socket_timeout` (Optional) : A String that indicates the timeout duration for waiting for data. Supports ISO_8601 notation Strings ("PT20.345S", "PT15M", etc.) as well as simple notation Strings for seconds ("60s") and milliseconds ("1500ms"). If this timeout value not set, the underlying Apache HttpClient would rely on operating system settings for managing socket timeouts. - - -* `connection_timeout` (Optional) : A String that indicates the timeout duration used when requesting a connection from the connection manager. Supports ISO_8601 notation Strings ("PT20.345S", "PT15M", etc.) as well as simple notation Strings for seconds ("60s") and milliseconds ("1500ms"). If this timeout value is either negative or not set, the underlying Apache HttpClient would rely on operating system settings for managing connection timeouts. - -### Indices Configuration - -Can be used to filter which indices should be processed. -An index will be processed if its name matches one of the `index_name_regex` -patterns in the `include` list, and does not match any of the pattern in the `exclude` list. -The default behavior is to process all indices. - -* `include` (Optional) : A List of [Index Configuration](#index_configuration) that defines which indices should be processed. Defaults to an empty list. - - -* `exclude` (Optional) : A List of [Index Configuration](#index_configuration) that defines which indices should not be processed. - -#### Index Configuration - -* `index_name_regex`: A regex pattern to represent the index names for filtering diff --git a/data-prepper-plugins/opensearch-source/build.gradle b/data-prepper-plugins/opensearch-source/build.gradle deleted file mode 100644 index c080eff561..0000000000 --- a/data-prepper-plugins/opensearch-source/build.gradle +++ /dev/null @@ -1,29 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -dependencies { - implementation project(path: ':data-prepper-api') - implementation project(':data-prepper-plugins:buffer-common') - implementation project(':data-prepper-plugins:aws-plugin-api') - implementation 'software.amazon.awssdk:apache-client' - implementation 'software.amazon.awssdk:netty-nio-client' - implementation 'io.micrometer:micrometer-core' - implementation 'com.fasterxml.jackson.core:jackson-databind' - implementation 'com.fasterxml.jackson.datatype:jackson-datatype-jsr310:2.15.2' - implementation 'software.amazon.awssdk:s3' - implementation 'software.amazon.awssdk:sts' - testImplementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml' - implementation 'org.opensearch.client:opensearch-java:2.5.0' - implementation 'org.opensearch.client:opensearch-rest-client:2.7.0' - implementation 'co.elastic.clients:elasticsearch-java:7.17.0' - implementation libs.commons.lang3 - implementation('org.apache.maven:maven-artifact:3.0.3') { - exclude group: 'org.codehaus.plexus' - } - testImplementation testLibs.mockito.inline -} - -test { - useJUnitPlatform() -} \ No newline at end of file diff --git a/data-prepper-plugins/opensearch/README.md b/data-prepper-plugins/opensearch/README.md index 5bff437fa0..6ea7f3972c 100644 --- a/data-prepper-plugins/opensearch/README.md +++ b/data-prepper-plugins/opensearch/README.md @@ -357,3 +357,199 @@ This plugin is compatible with Java 8. 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) + +# OpenSearch Source + +This is the Date Prepper OpenSearch source plugin that processes indices for either OpenSearch, Elasticsearch, +or Amazon OpenSearch Service clusters. It is ideal for migrating index data from a cluster. + +Note: Only fully tested versions will be listed below. It is likely many more versions are supported already, but it is untested. + +The OpenSearch source is compatible with the following OpenSearch versions: +* 2.5 + +And is compatible with the following Elasticsearch versions: +* 7.10 + +# Usages + +### Minimum required config with username and password + +```yaml +opensearch-source-pipeline: + source: + opensearch: + connection: + insecure: true + hosts: [ "https://localhost:9200" ] + username: "username" + password: "password" +``` + +### Full config example + +```yaml +opensearch-source-pipeline: + source: + opensearch: + indices: + include: + - index_name_regex: "test-index-.*" + exclude: + - index_name_regex: "test-index-[1-9].*" + scheduling: + rate: "PT1H" + start_time: "2023-06-02T22:01:30.00Z" + job_count: 2 + search_options: + search_context_type: "none" + batch_size: 1000 + connection: + insecure: false + cert: "/path/to/cert.crt" + socket_timeout: "100ms" + connection_timeout: "100ms" + hosts: [ "https://localhost:9200" ] + username: "username" + password: "password" +``` + +### Amazon OpenSearch Service + +The OpenSearch source can also be configured for an Amazon OpenSearch Service domain. + +```yaml +opensearch-source-pipeline: + source: + opensearch: + connection: + insecure: true + hosts: [ "https://search-my-domain-soopywaovobopgs8ywurr3utsu.us-east-1.es.amazonaws.com" ] + aws: + region: "us-east-1" + sts_role_arn: "arn:aws:iam::123456789012:role/my-domain-role" +``` + +### Using Metadata + +When the OpenSearch source constructs Data Prepper Events from documents in the cluster, the +document index is stored in the `EventMetadata` with an `opensearch-index` key, and the document_id is +stored in the `EventMetadata` with a `opensearch-document_id` key. This allows conditional routing based on the index or document_id, +among other things. For example, one could send to an OpenSearch sink and use the same index and document_id from the source cluster in +the destination cluster. A full config example for this use case is below + +```yaml +opensearch-source-pipeline: + source: + opensearch: + connection: + insecure: true + hosts: [ "https://source-cluster:9200" ] + username: "username" + password: "password" + sink: + - opensearch: + hosts: [ "https://sink-cluster:9200" ] + username: "username" + password: "password" + document_id_field: "getMetadata(\"opensearch-document_id\")" + index: "${getMetadata(\"opensearch-index\"}" +``` + +## Configuration + +- `hosts` (Required) : A list of IP addresses of OpenSearch or Elasticsearch nodes. + + +- `username` (Optional) : A String of username used in the internal users of OpenSearch cluster. Default is null. + + +- `password` (Optional) : A String of password used in the internal users of OpenSearch cluster. Default is null. + + +- `disable_authentication` (Optional) : A boolean that can disable authentication if the cluster supports it. Defaults to false. + + +- `aws` (Optional) : AWS configurations. See [AWS Configuration](#aws_configuration) for details. SigV4 is enabled by default when this option is used. + + +- `search_options` (Optional) : See [Search Configuration](#search_configuration) for details + + +- `indices` (Optional): See [Indices Configurations](#indices_configuration) for filtering options. + + +- `scheduling` (Optional): See [Scheduling Configuration](#scheduling_configuration) for details + + +- `connection` (Optional): See [Connection Configuration](#connection_configuration) + +### AWS Configuration + +* `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 STS role to assume for requests to AWS. 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_header_overrides` (Optional): A map of header overrides to make when assuming the IAM role for the source plugin. + +### Search Configuration + +* `search_context_type` (Optional) : A direct override for which type of search context should be used to search documents. + Options include `point_in_time`, `scroll`, or `none` (just search after). + By default, the OpenSearch source will attempt to use `point_in_time` on a cluster by auto-detecting that the cluster version and distribution + supports Point in Time. If the cluster does not support `point_in_time`, then `scroll` is the default behavior. + + +* `batch_size` (Optional) : The amount of documents to read in at once while searching. + This size is passed to the search requests for all search context types (`none` (search_after), `point_in_time`, or `scroll`). + Defaults to 1,000. + +### Scheduling Configuration + +Schedule the start time and amount of times an index should be processed. For example, +a `rate` of `PT1H` and a `job_count` of 3 would result in each index getting processed 3 times, starting at `start_time` +and then every hour after the first time the index is processed. + +* `rate` (Optional) : A String that indicates the rate to process an index based on the `job_count`. + 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 `job_count` is greater than 1. + + + +* `job_count` (Optional) : An Integer that specifies how many times each index should be processed. Defaults to 1. + + + +* `start_time` (Optional) : A String in the format of a timestamp that is compatible with Java Instant (i.e. `2023-06-02T22:01:30.00Z`). + Processing will be delayed until this timestamp is reached. The default start time is to start immediately. + +### Connection Configuration + +* `insecure` (Optional): A boolean flag to turn off SSL certificate verification. If set to true, CA certificate verification will be turned off and insecure HTTP requests will be sent. Default to false. + + +* `cert` (Optional) : CA certificate that is pem encoded. Accepts both .pem or .crt. This enables the client to trust the CA that has signed the certificate that the OpenSearch cluster is using. Default is null. + + +* `socket_timeout` (Optional) : A String that indicates the timeout duration for waiting for data. Supports ISO_8601 notation Strings ("PT20.345S", "PT15M", etc.) as well as simple notation Strings for seconds ("60s") and milliseconds ("1500ms"). If this timeout value not set, the underlying Apache HttpClient would rely on operating system settings for managing socket timeouts. + + +* `connection_timeout` (Optional) : A String that indicates the timeout duration used when requesting a connection from the connection manager. Supports ISO_8601 notation Strings ("PT20.345S", "PT15M", etc.) as well as simple notation Strings for seconds ("60s") and milliseconds ("1500ms"). If this timeout value is either negative or not set, the underlying Apache HttpClient would rely on operating system settings for managing connection timeouts. + +### Indices Configuration + +Can be used to filter which indices should be processed. +An index will be processed if its name matches one of the `index_name_regex` +patterns in the `include` list, and does not match any of the pattern in the `exclude` list. +The default behavior is to process all indices. + +* `include` (Optional) : A List of [Index Configuration](#index_configuration) that defines which indices should be processed. Defaults to an empty list. + + +* `exclude` (Optional) : A List of [Index Configuration](#index_configuration) that defines which indices should not be processed. + +#### Index Configuration + +* `index_name_regex`: A regex pattern to represent the index names for filtering diff --git a/data-prepper-plugins/opensearch/build.gradle b/data-prepper-plugins/opensearch/build.gradle index 17c632a359..2609f148b0 100644 --- a/data-prepper-plugins/opensearch/build.gradle +++ b/data-prepper-plugins/opensearch/build.gradle @@ -7,14 +7,16 @@ dependencies { implementation project(':data-prepper-api') implementation libs.armeria.core testImplementation project(':data-prepper-api').sourceSets.test.output + implementation project(':data-prepper-plugins:aws-plugin-api') + 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:aws-plugin-api') implementation libs.opensearch.client implementation libs.opensearch.rhlc implementation libs.opensearch.java implementation 'com.fasterxml.jackson.core:jackson-databind' implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml' + implementation 'com.fasterxml.jackson.datatype:jackson-datatype-jsr310:2.15.2' implementation 'javax.ws.rs:javax.ws.rs-api:2.1.1' implementation 'software.amazon.awssdk:auth' implementation 'software.amazon.awssdk:http-client-spi' @@ -31,11 +33,17 @@ dependencies { implementation libs.commons.lang3 implementation 'com.github.ben-manes.caffeine:caffeine:3.1.8' implementation 'software.amazon.awssdk:apache-client' + implementation 'software.amazon.awssdk:netty-nio-client' + implementation 'co.elastic.clients:elasticsearch-java:7.17.0' + implementation('org.apache.maven:maven-artifact:3.0.3') { + exclude group: 'org.codehaus.plexus' + } testImplementation testLibs.junit.vintage testImplementation libs.commons.io testImplementation 'net.bytebuddy:byte-buddy:1.14.8' testImplementation 'net.bytebuddy:byte-buddy-agent:1.14.9' testImplementation testLibs.slf4j.simple + testImplementation testLibs.mockito.inline } sourceSets { diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/ClientRefresher.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/ClientRefresher.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/ClientRefresher.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/ClientRefresher.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchIndexProgressState.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchIndexProgressState.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchIndexProgressState.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchIndexProgressState.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchService.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchService.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchService.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchService.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSource.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSource.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSource.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSource.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSourceConfiguration.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSourceConfiguration.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSourceConfiguration.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSourceConfiguration.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/AwsAuthenticationConfiguration.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/AwsAuthenticationConfiguration.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/AwsAuthenticationConfiguration.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/AwsAuthenticationConfiguration.java diff --git a/data-prepper-plugins/opensearch-source/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 similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/ConnectionConfiguration.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/ConnectionConfiguration.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/IndexParametersConfiguration.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/IndexParametersConfiguration.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/IndexParametersConfiguration.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/IndexParametersConfiguration.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/OpenSearchIndex.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/OpenSearchIndex.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/OpenSearchIndex.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/OpenSearchIndex.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SchedulingParameterConfiguration.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SchedulingParameterConfiguration.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SchedulingParameterConfiguration.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SchedulingParameterConfiguration.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SearchConfiguration.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SearchConfiguration.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SearchConfiguration.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SearchConfiguration.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/metrics/OpenSearchSourcePluginMetrics.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/metrics/OpenSearchSourcePluginMetrics.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/metrics/OpenSearchSourcePluginMetrics.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/metrics/OpenSearchSourcePluginMetrics.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorker.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorker.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorker.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorker.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/OpenSearchIndexPartitionCreationSupplier.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/OpenSearchIndexPartitionCreationSupplier.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/OpenSearchIndexPartitionCreationSupplier.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/OpenSearchIndexPartitionCreationSupplier.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorker.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorker.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorker.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorker.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorker.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorker.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorker.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorker.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/SearchWorker.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/SearchWorker.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/SearchWorker.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/SearchWorker.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtils.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtils.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtils.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtils.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/ClusterClientFactory.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/ClusterClientFactory.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/ClusterClientFactory.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/ClusterClientFactory.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/ElasticsearchAccessor.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/ElasticsearchAccessor.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/ElasticsearchAccessor.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/ElasticsearchAccessor.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchAccessor.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchAccessor.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchAccessor.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchAccessor.java diff --git a/data-prepper-plugins/opensearch-source/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 similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchClientFactory.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchClientFactory.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/SearchAccessor.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/SearchAccessor.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/SearchAccessor.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/SearchAccessor.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/SearchAccessorStrategy.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/SearchAccessorStrategy.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/SearchAccessorStrategy.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/SearchAccessorStrategy.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/X509TrustAllManager.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/X509TrustAllManager.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/X509TrustAllManager.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/X509TrustAllManager.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/exceptions/IndexNotFoundException.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/exceptions/IndexNotFoundException.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/exceptions/IndexNotFoundException.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/exceptions/IndexNotFoundException.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/exceptions/SearchContextLimitException.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/exceptions/SearchContextLimitException.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/exceptions/SearchContextLimitException.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/exceptions/SearchContextLimitException.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreatePointInTimeRequest.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreatePointInTimeRequest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreatePointInTimeRequest.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreatePointInTimeRequest.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreatePointInTimeResponse.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreatePointInTimeResponse.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreatePointInTimeResponse.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreatePointInTimeResponse.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreateScrollRequest.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreateScrollRequest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreateScrollRequest.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreateScrollRequest.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreateScrollResponse.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreateScrollResponse.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreateScrollResponse.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/CreateScrollResponse.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/DeletePointInTimeRequest.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/DeletePointInTimeRequest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/DeletePointInTimeRequest.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/DeletePointInTimeRequest.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/DeleteScrollRequest.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/DeleteScrollRequest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/DeleteScrollRequest.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/DeleteScrollRequest.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/MetadataKeyAttributes.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/MetadataKeyAttributes.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/MetadataKeyAttributes.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/MetadataKeyAttributes.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/NoSearchContextSearchRequest.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/NoSearchContextSearchRequest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/NoSearchContextSearchRequest.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/NoSearchContextSearchRequest.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchContextType.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchContextType.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchContextType.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchContextType.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchPointInTimeRequest.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchPointInTimeRequest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchPointInTimeRequest.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchPointInTimeRequest.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchScrollRequest.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchScrollRequest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchScrollRequest.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchScrollRequest.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchScrollResponse.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchScrollResponse.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchScrollResponse.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchScrollResponse.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchWithSearchAfterResults.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchWithSearchAfterResults.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchWithSearchAfterResults.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SearchWithSearchAfterResults.java diff --git a/data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SortingOptions.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SortingOptions.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SortingOptions.java rename to data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/model/SortingOptions.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/ClientRefresherTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/ClientRefresherTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/ClientRefresherTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/ClientRefresherTest.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchServiceTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchServiceTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchServiceTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchServiceTest.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSourceConfigurationTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSourceConfigurationTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSourceConfigurationTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSourceConfigurationTest.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSourceTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSourceTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSourceTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/OpenSearchSourceTest.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/AwsAuthenticationConfigurationTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/AwsAuthenticationConfigurationTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/AwsAuthenticationConfigurationTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/AwsAuthenticationConfigurationTest.java diff --git a/data-prepper-plugins/opensearch-source/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 similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/ConnectionConfigurationTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/ConnectionConfigurationTest.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/IndexParametersConfigurationTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/IndexParametersConfigurationTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/IndexParametersConfigurationTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/IndexParametersConfigurationTest.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/OpenSearchIndexTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/OpenSearchIndexTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/OpenSearchIndexTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/OpenSearchIndexTest.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SchedulingParameterConfigurationTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SchedulingParameterConfigurationTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SchedulingParameterConfigurationTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SchedulingParameterConfigurationTest.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SearchConfigurationTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SearchConfigurationTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SearchConfigurationTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/configuration/SearchConfigurationTest.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorkerTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorkerTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorkerTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/NoSearchContextWorkerTest.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorkerTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorkerTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorkerTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorkerTest.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorkerTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorkerTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorkerTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/ScrollWorkerTest.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtilsTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtilsTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtilsTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtilsTest.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/ElasticsearchAccessorTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/ElasticsearchAccessorTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/ElasticsearchAccessorTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/ElasticsearchAccessorTest.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchAccessorTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchAccessorTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchAccessorTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchAccessorTest.java diff --git a/data-prepper-plugins/opensearch-source/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 similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchClientFactoryTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchClientFactoryTest.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchIndexPartitionCreationSupplierTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchIndexPartitionCreationSupplierTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchIndexPartitionCreationSupplierTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/OpenSearchIndexPartitionCreationSupplierTest.java diff --git a/data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/SearchAccessStrategyTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/SearchAccessStrategyTest.java similarity index 100% rename from data-prepper-plugins/opensearch-source/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/SearchAccessStrategyTest.java rename to data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/client/SearchAccessStrategyTest.java diff --git a/settings.gradle b/settings.gradle index 5ce7405bbc..b674d1689d 100644 --- a/settings.gradle +++ b/settings.gradle @@ -132,7 +132,6 @@ 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:opensearch-source' include 'data-prepper-plugins:user-agent-processor' include 'data-prepper-plugins:in-memory-source-coordination-store' include 'data-prepper-plugins:aws-plugin-api' From 3e558c04f62d60a111a11c8afbe6cf842a801db2 Mon Sep 17 00:00:00 2001 From: David Venable Date: Wed, 8 Nov 2023 11:21:20 -0800 Subject: [PATCH 42/53] Pause reading from Kafka in the Kafka buffer when the circuit breaker is open (#3595) Use the CircuitBreaker in the Kafka buffer to stop reading data from the Kafka topic and putting it into the in-memory buffer. Moves the CircuitBreaker class into data-prepper-api. Adds a DelegatingBuffer class to data-prepper-api. Resolves #3578. Signed-off-by: David Venable --- .../model}/breaker/CircuitBreaker.java | 6 +- .../model/buffer/DelegatingBuffer.java | 85 ++++++++ .../model/buffer/DelegatingBufferTest.java | 202 ++++++++++++++++++ .../breaker/CircuitBreakerAppConfig.java | 7 + .../breaker/CircuitBreakerManager.java | 2 + .../breaker/HeapCircuitBreaker.java | 1 + .../breaker/InnerCircuitBreaker.java | 2 + .../parser/CircuitBreakingBuffer.java | 49 +---- .../ComponentPluginArgumentsContext.java | 9 + .../plugin/DefaultPluginFactory.java | 10 +- .../dataprepper/breaker/CircuitBreakerIT.java | 1 + .../breaker/CircuitBreakerManagerTest.java | 1 + .../parser/CircuitBreakingBufferTest.java | 2 +- .../parser/PipelineTransformerTests.java | 2 +- .../plugin/DefaultPluginFactoryTest.java | 5 +- .../plugins/kafka/buffer/KafkaBufferIT.java | 8 +- .../plugins/kafka/buffer/KafkaBuffer.java | 21 +- .../kafka/consumer/KafkaCustomConsumer.java | 10 +- .../consumer/KafkaCustomConsumerFactory.java | 8 +- .../kafka/consumer/PauseConsumePredicate.java | 44 ++++ .../plugins/kafka/source/KafkaSource.java | 3 +- .../plugins/kafka/buffer/KafkaBufferTest.java | 14 +- .../consumer/KafkaCustomConsumerTest.java | 5 +- .../consumer/PauseConsumePredicateTest.java | 54 +++++ 24 files changed, 477 insertions(+), 74 deletions(-) rename {data-prepper-core/src/main/java/org/opensearch/dataprepper => data-prepper-api/src/main/java/org/opensearch/dataprepper/model}/breaker/CircuitBreaker.java (81%) create mode 100644 data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/DelegatingBuffer.java create mode 100644 data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/DelegatingBufferTest.java create mode 100644 data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/PauseConsumePredicate.java create mode 100644 data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/PauseConsumePredicateTest.java diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreaker.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/breaker/CircuitBreaker.java similarity index 81% rename from data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreaker.java rename to data-prepper-api/src/main/java/org/opensearch/dataprepper/model/breaker/CircuitBreaker.java index 9774648b59..25fa34c56a 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreaker.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/breaker/CircuitBreaker.java @@ -3,12 +3,12 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.dataprepper.breaker; +package org.opensearch.dataprepper.model.breaker; /** * Represents a circuit breaker in Data Prepper. * - * @since 2.1 + * @since 2.6 */ public interface CircuitBreaker { /** @@ -16,7 +16,7 @@ public interface CircuitBreaker { * been tripped. * * @return true if open; false if closed. - * @since 2.1 + * @since 2.6 */ boolean isOpen(); } diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/DelegatingBuffer.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/DelegatingBuffer.java new file mode 100644 index 0000000000..adc6b76b9d --- /dev/null +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/DelegatingBuffer.java @@ -0,0 +1,85 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.model.buffer; + +import org.opensearch.dataprepper.model.CheckpointState; +import org.opensearch.dataprepper.model.record.Record; + +import java.time.Duration; +import java.util.Collection; +import java.util.Map; +import java.util.Objects; +import java.util.concurrent.TimeoutException; + +/** + * An implementation of {@link Buffer} which delegates all calls to a delgate + * (or inner) buffer. + *

+ * This class exists to help with writing decorators of the {@link Buffer} interface. + * + * @param The type of data in the buffer + * + * @since 2.6 + */ +public abstract class DelegatingBuffer> implements Buffer { + private final Buffer delegateBuffer; + + /** + * Constructor for subclasses to use. + * + * @param delegateBuffer The delegate (or inner) buffer. + * + * @since 2.6 + */ + protected DelegatingBuffer(final Buffer delegateBuffer) { + this.delegateBuffer = Objects.requireNonNull(delegateBuffer); + } + + @Override + public void write(final T record, final int timeoutInMillis) throws TimeoutException { + delegateBuffer.write(record, timeoutInMillis); + } + + @Override + public void writeAll(final Collection records, final int timeoutInMillis) throws Exception { + delegateBuffer.writeAll(records, timeoutInMillis); + } + + @Override + public void writeBytes(final byte[] bytes, final String key, final int timeoutInMillis) throws Exception { + delegateBuffer.writeBytes(bytes, key, timeoutInMillis); + } + + @Override + public Map.Entry, CheckpointState> read(final int timeoutInMillis) { + return delegateBuffer.read(timeoutInMillis); + } + + @Override + public void checkpoint(final CheckpointState checkpointState) { + delegateBuffer.checkpoint(checkpointState); + } + + @Override + public boolean isEmpty() { + return delegateBuffer.isEmpty(); + } + + @Override + public boolean isByteBuffer() { + return delegateBuffer.isByteBuffer(); + } + + @Override + public Duration getDrainTimeout() { + return delegateBuffer.getDrainTimeout(); + } + + @Override + public void shutdown() { + delegateBuffer.shutdown(); + } +} diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/DelegatingBufferTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/DelegatingBufferTest.java new file mode 100644 index 0000000000..6934104593 --- /dev/null +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/DelegatingBufferTest.java @@ -0,0 +1,202 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.model.buffer; + +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.ValueSource; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.model.CheckpointState; +import org.opensearch.dataprepper.model.record.Record; + +import java.time.Duration; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.UUID; +import java.util.concurrent.TimeoutException; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.sameInstance; +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.anyInt; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +class DelegatingBufferTest { + @Mock + private Buffer> innerBuffer; + + @Mock + private Record record; + + private Collection> records; + + private int timeoutInMillis; + private Random random; + + @BeforeEach + void setUp() { + random = new Random(); + timeoutInMillis = random.nextInt(10_000) + 100; + + records = List.of(record, mock(Record.class)); + } + + private static class TestDelegatingBuffer extends DelegatingBuffer> { + TestDelegatingBuffer(final Buffer> delegateBuffer) { + super(delegateBuffer); + } + } + + private DelegatingBuffer> createObjectUnderTest() { + return new TestDelegatingBuffer(innerBuffer); + } + + @Test + void constructor_throws_if_delegate_is_null() { + innerBuffer = null; + + assertThrows(NullPointerException.class, this::createObjectUnderTest); + } + + @Test + void write_calls_inner_write() throws TimeoutException { + createObjectUnderTest().write(record, timeoutInMillis); + + verify(innerBuffer).write(record, timeoutInMillis); + } + + @ParameterizedTest + @ValueSource(classes = {RuntimeException.class, TimeoutException.class}) + void write_throws_exceptions_from_inner_write(final Class exceptionType) throws TimeoutException { + final Throwable exception = mock(exceptionType); + doThrow(exception).when(innerBuffer).write(any(), anyInt()); + + final DelegatingBuffer> objectUnderTest = createObjectUnderTest(); + final Exception actualException = assertThrows(Exception.class, () -> objectUnderTest.write(record, timeoutInMillis)); + + assertThat(actualException, sameInstance(exception)); + } + + @Test + void writeAll_calls_inner_writeAll() throws Exception { + createObjectUnderTest().writeAll(records, timeoutInMillis); + + verify(innerBuffer).writeAll(records, timeoutInMillis); + } + + @ParameterizedTest + @ValueSource(classes = {Exception.class, RuntimeException.class, TimeoutException.class}) + void writeAll_throws_exceptions_from_inner_writeAll(final Class exceptionType) throws Exception { + final Throwable exception = mock(exceptionType); + doThrow(exception).when(innerBuffer).writeAll(any(), anyInt()); + + final DelegatingBuffer> objectUnderTest = createObjectUnderTest(); + final Exception actualException = assertThrows(Exception.class, () -> objectUnderTest.writeAll(records, timeoutInMillis)); + + assertThat(actualException, sameInstance(exception)); + } + + @Test + void writeBytes_calls_inner_writeBytes() throws Exception { + final byte[] bytesToWrite = new byte[64]; + random.nextBytes(bytesToWrite); + final String key = UUID.randomUUID().toString(); + createObjectUnderTest().writeBytes(bytesToWrite, key, timeoutInMillis); + + verify(innerBuffer).writeBytes(bytesToWrite, key, timeoutInMillis); + } + + @ParameterizedTest + @ValueSource(classes = {Exception.class, RuntimeException.class, TimeoutException.class}) + void writeBytes_throws_exceptions_from_inner_writeBytes(final Class exceptionType) throws Exception { + final Throwable exception = mock(exceptionType); + doThrow(exception).when(innerBuffer).writeBytes(any(), any(), anyInt()); + + final byte[] bytesToWrite = new byte[64]; + random.nextBytes(bytesToWrite); + final String key = UUID.randomUUID().toString(); + + final DelegatingBuffer> objectUnderTest = createObjectUnderTest(); + final Exception actualException = assertThrows(Exception.class, () -> objectUnderTest.writeBytes(bytesToWrite, key, timeoutInMillis)); + + assertThat(actualException, sameInstance(exception)); + } + + @Test + void read_returns_inner_read() { + final Map.Entry>, CheckpointState> readResult = mock(Map.Entry.class); + when(innerBuffer.read(timeoutInMillis)).thenReturn(readResult); + + assertThat(createObjectUnderTest().read(timeoutInMillis), + equalTo(readResult)); + } + + @Test + void read_throws_exceptions_from_inner_read() { + final RuntimeException exception = mock(RuntimeException.class); + + when(innerBuffer.read(timeoutInMillis)).thenThrow(exception); + + final DelegatingBuffer> objectUnderTest = createObjectUnderTest(); + final Exception actualException = assertThrows(Exception.class, () -> objectUnderTest.read(timeoutInMillis)); + + assertThat(actualException, sameInstance(exception)); + } + + @Test + void checkpoint_calls_inner_checkpoint() { + final CheckpointState checkpointState = mock(CheckpointState.class); + + createObjectUnderTest().checkpoint(checkpointState); + + verify(innerBuffer).checkpoint(checkpointState); + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void isEmpty_returns_inner_isEmpty(final boolean isEmpty) { + when(innerBuffer.isEmpty()).thenReturn(isEmpty); + + assertThat(createObjectUnderTest().isEmpty(), + equalTo(isEmpty)); + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void isByteBuffer_returns_inner_isByteBuffer(final boolean isByteBuffer) { + when(innerBuffer.isByteBuffer()).thenReturn(isByteBuffer); + + assertThat(createObjectUnderTest().isByteBuffer(), + equalTo(isByteBuffer)); + } + + @Test + void getDrainTimeout_returns_inner_getDrainTimeout() { + final Duration drainTimeout = Duration.ofSeconds(random.nextInt(10_000) + 100); + when(innerBuffer.getDrainTimeout()).thenReturn(drainTimeout); + + assertThat(createObjectUnderTest().getDrainTimeout(), + equalTo(drainTimeout)); + } + + @Test + void shutdown_calls_inner_shutdown() { + createObjectUnderTest().shutdown(); + + verify(innerBuffer).shutdown(); + } +} \ No newline at end of file diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreakerAppConfig.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreakerAppConfig.java index 19fb683dcd..c8d6393682 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreakerAppConfig.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreakerAppConfig.java @@ -5,12 +5,14 @@ package org.opensearch.dataprepper.breaker; +import org.opensearch.dataprepper.model.breaker.CircuitBreaker; import org.opensearch.dataprepper.parser.model.CircuitBreakerConfig; import org.opensearch.dataprepper.parser.model.DataPrepperConfiguration; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import java.util.List; +import java.util.Optional; /** * The application config for circuit breakers. Used for wiring beans @@ -34,4 +36,9 @@ InnerCircuitBreaker heapCircuitBreaker(final DataPrepperConfiguration dataPreppe return null; } } + + @Bean + public Optional circuitBreaker(final CircuitBreakerManager circuitBreakerManager) { + return circuitBreakerManager.getGlobalCircuitBreaker(); + } } diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreakerManager.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreakerManager.java index 23360ab04c..405b9e3bdc 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreakerManager.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/CircuitBreakerManager.java @@ -5,6 +5,8 @@ package org.opensearch.dataprepper.breaker; +import org.opensearch.dataprepper.model.breaker.CircuitBreaker; + import java.util.List; import java.util.Optional; diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/HeapCircuitBreaker.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/HeapCircuitBreaker.java index dc34bd5cd0..9993b9e4a8 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/HeapCircuitBreaker.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/HeapCircuitBreaker.java @@ -6,6 +6,7 @@ package org.opensearch.dataprepper.breaker; import io.micrometer.core.instrument.Metrics; +import org.opensearch.dataprepper.model.breaker.CircuitBreaker; import org.opensearch.dataprepper.parser.model.HeapCircuitBreakerConfig; import org.slf4j.Logger; import org.slf4j.LoggerFactory; diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/InnerCircuitBreaker.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/InnerCircuitBreaker.java index 50fb2450cb..ebaa81fa7d 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/InnerCircuitBreaker.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/breaker/InnerCircuitBreaker.java @@ -5,6 +5,8 @@ package org.opensearch.dataprepper.breaker; +import org.opensearch.dataprepper.model.breaker.CircuitBreaker; + /** * Interface to signal that this {@link CircuitBreaker} to prevent * access beyond the {@link CircuitBreakerManager}. diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/CircuitBreakingBuffer.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/CircuitBreakingBuffer.java index 2242f2db65..d61f183f05 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/CircuitBreakingBuffer.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/CircuitBreakingBuffer.java @@ -5,14 +5,12 @@ package org.opensearch.dataprepper.parser; -import org.opensearch.dataprepper.breaker.CircuitBreaker; -import org.opensearch.dataprepper.model.CheckpointState; +import org.opensearch.dataprepper.model.breaker.CircuitBreaker; import org.opensearch.dataprepper.model.buffer.Buffer; +import org.opensearch.dataprepper.model.buffer.DelegatingBuffer; import org.opensearch.dataprepper.model.record.Record; -import java.time.Duration; import java.util.Collection; -import java.util.Map; import java.util.concurrent.TimeoutException; import static java.util.Objects.requireNonNull; @@ -24,8 +22,7 @@ * @param The type of record. * @since 2.1 */ -class CircuitBreakingBuffer> implements Buffer { - private final Buffer buffer; +class CircuitBreakingBuffer> extends DelegatingBuffer implements Buffer { private final CircuitBreaker circuitBreaker; /** @@ -35,7 +32,7 @@ class CircuitBreakingBuffer> implements Buffer { * @param circuitBreaker The circuit breaker to check */ public CircuitBreakingBuffer(final Buffer buffer, final CircuitBreaker circuitBreaker) { - this.buffer = requireNonNull(buffer); + super(buffer); this.circuitBreaker = requireNonNull(circuitBreaker); } @@ -43,55 +40,25 @@ public CircuitBreakingBuffer(final Buffer buffer, final CircuitBreaker circui public void write(final T record, final int timeoutInMillis) throws TimeoutException { checkBreaker(); - buffer.write(record, timeoutInMillis); + super.write(record, timeoutInMillis); } @Override public void writeAll(final Collection records, final int timeoutInMillis) throws Exception { checkBreaker(); - buffer.writeAll(records, timeoutInMillis); + super.writeAll(records, timeoutInMillis); } @Override - public void writeBytes(byte[] bytes, String key, int timeoutInMillis) throws Exception { + public void writeBytes(final byte[] bytes, final String key, final int timeoutInMillis) throws Exception { checkBreaker(); - buffer.writeBytes(bytes, key, timeoutInMillis); + super.writeBytes(bytes, key, timeoutInMillis); } private void checkBreaker() throws TimeoutException { if(circuitBreaker.isOpen()) throw new TimeoutException("Circuit breaker is open. Unable to write to buffer."); } - - @Override - public Map.Entry, CheckpointState> read(final int timeoutInMillis) { - return buffer.read(timeoutInMillis); - } - - @Override - public void checkpoint(final CheckpointState checkpointState) { - buffer.checkpoint(checkpointState); - } - - @Override - public boolean isEmpty() { - return buffer.isEmpty(); - } - - @Override - public boolean isByteBuffer() { - return buffer.isByteBuffer(); - } - - @Override - public Duration getDrainTimeout() { - return buffer.getDrainTimeout(); - } - - @Override - public void shutdown() { - buffer.shutdown(); - } } diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugin/ComponentPluginArgumentsContext.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugin/ComponentPluginArgumentsContext.java index bb688abb8f..6be9244c08 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugin/ComponentPluginArgumentsContext.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugin/ComponentPluginArgumentsContext.java @@ -5,6 +5,7 @@ package org.opensearch.dataprepper.plugin; +import org.opensearch.dataprepper.model.breaker.CircuitBreaker; import org.opensearch.dataprepper.model.plugin.PluginConfigObservable; import org.opensearch.dataprepper.model.sink.SinkContext; import org.opensearch.dataprepper.metrics.PluginMetrics; @@ -74,6 +75,8 @@ private ComponentPluginArgumentsContext(final Builder builder) { if (builder.sinkContext != null) { typedArgumentsSuppliers.put(SinkContext.class, () -> builder.sinkContext); } + + typedArgumentsSuppliers.put(CircuitBreaker.class, () -> builder.circuitBreaker); } @Override @@ -135,6 +138,7 @@ static class Builder { private AcknowledgementSetManager acknowledgementSetManager; private PluginConfigObservable pluginConfigObservable; private SinkContext sinkContext; + private CircuitBreaker circuitBreaker; Builder withPluginConfiguration(final Object pluginConfiguration) { this.pluginConfiguration = pluginConfiguration; @@ -181,6 +185,11 @@ Builder withPluginConfigurationObservable(final PluginConfigObservable pluginCon return this; } + Builder withCircuitBreaker(final CircuitBreaker circuitBreaker) { + this.circuitBreaker = circuitBreaker; + return this; + } + ComponentPluginArgumentsContext build() { return new ComponentPluginArgumentsContext(this); } diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugin/DefaultPluginFactory.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugin/DefaultPluginFactory.java index bd53b9678a..8e6f070526 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugin/DefaultPluginFactory.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugin/DefaultPluginFactory.java @@ -5,6 +5,7 @@ package org.opensearch.dataprepper.plugin; +import org.opensearch.dataprepper.model.breaker.CircuitBreaker; import org.opensearch.dataprepper.model.plugin.PluginConfigObservable; import org.opensearch.dataprepper.model.sink.SinkContext; import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; @@ -15,6 +16,7 @@ import org.opensearch.dataprepper.acknowledgements.DefaultAcknowledgementSetManager; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.springframework.beans.factory.annotation.Autowired; import org.springframework.context.annotation.DependsOn; import javax.inject.Inject; @@ -43,6 +45,7 @@ public class DefaultPluginFactory implements PluginFactory { private final DefaultEventFactory eventFactory; private final DefaultAcknowledgementSetManager acknowledgementSetManager; private final PluginConfigurationObservableFactory pluginConfigurationObservableFactory; + private final CircuitBreaker circuitBreaker; @Inject DefaultPluginFactory( @@ -52,8 +55,10 @@ public class DefaultPluginFactory implements PluginFactory { final PluginBeanFactoryProvider pluginBeanFactoryProvider, final DefaultEventFactory eventFactory, final DefaultAcknowledgementSetManager acknowledgementSetManager, - final PluginConfigurationObservableFactory pluginConfigurationObservableFactory - ) { + final PluginConfigurationObservableFactory pluginConfigurationObservableFactory, + @Autowired(required = false) final CircuitBreaker circuitBreaker + ) { + this.circuitBreaker = circuitBreaker; Objects.requireNonNull(pluginProviderLoader); Objects.requireNonNull(pluginConfigurationObservableFactory); this.pluginCreator = Objects.requireNonNull(pluginCreator); @@ -131,6 +136,7 @@ private ComponentPluginArgumentsContext getConstructionContext(final PluginS .withAcknowledgementSetManager(acknowledgementSetManager) .withPluginConfigurationObservable(pluginConfigObservable) .withSinkContext(sinkContext) + .withCircuitBreaker(circuitBreaker) .build(); } diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/breaker/CircuitBreakerIT.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/breaker/CircuitBreakerIT.java index ff14013db3..44c32b4f4d 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/breaker/CircuitBreakerIT.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/breaker/CircuitBreakerIT.java @@ -13,6 +13,7 @@ import org.junit.jupiter.params.provider.CsvSource; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.model.breaker.CircuitBreaker; import org.opensearch.dataprepper.model.types.ByteCount; import org.opensearch.dataprepper.parser.model.CircuitBreakerConfig; import org.opensearch.dataprepper.parser.model.DataPrepperConfiguration; diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/breaker/CircuitBreakerManagerTest.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/breaker/CircuitBreakerManagerTest.java index a57a83caa1..e92bf39e61 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/breaker/CircuitBreakerManagerTest.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/breaker/CircuitBreakerManagerTest.java @@ -13,6 +13,7 @@ import org.junit.jupiter.params.provider.ValueSource; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.dataprepper.model.breaker.CircuitBreaker; import java.util.Collections; import java.util.List; diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/CircuitBreakingBufferTest.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/CircuitBreakingBufferTest.java index 346d4fb8bf..faf8c25120 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/CircuitBreakingBufferTest.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/CircuitBreakingBufferTest.java @@ -14,7 +14,7 @@ import org.junit.jupiter.params.provider.ValueSource; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.dataprepper.breaker.CircuitBreaker; +import org.opensearch.dataprepper.model.breaker.CircuitBreaker; import org.opensearch.dataprepper.model.CheckpointState; import org.opensearch.dataprepper.model.buffer.Buffer; import org.opensearch.dataprepper.model.record.Record; diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/PipelineTransformerTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/PipelineTransformerTests.java index ce44b298fd..8182b4ab37 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/PipelineTransformerTests.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/PipelineTransformerTests.java @@ -16,7 +16,7 @@ import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.dataprepper.TestDataProvider; -import org.opensearch.dataprepper.breaker.CircuitBreaker; +import org.opensearch.dataprepper.model.breaker.CircuitBreaker; import org.opensearch.dataprepper.breaker.CircuitBreakerManager; import org.opensearch.dataprepper.model.buffer.Buffer; import org.opensearch.dataprepper.model.configuration.PipelinesDataFlowModel; diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugin/DefaultPluginFactoryTest.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugin/DefaultPluginFactoryTest.java index a43b9c17c4..1c920b6fc3 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugin/DefaultPluginFactoryTest.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugin/DefaultPluginFactoryTest.java @@ -7,6 +7,7 @@ import org.mockito.ArgumentCaptor; import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; +import org.opensearch.dataprepper.model.breaker.CircuitBreaker; import org.opensearch.dataprepper.model.configuration.PipelineDescription; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.plugin.NoPluginFoundException; @@ -57,6 +58,7 @@ class DefaultPluginFactoryTest { private BeanFactory beanFactory; private String pipelineName; private DefaultAcknowledgementSetManager acknowledgementSetManager; + private CircuitBreaker circuitBreaker; private DefaultEventFactory eventFactory; private PluginConfigurationObservableFactory pluginConfigurationObservableFactory; private PluginConfigObservable pluginConfigObservable; @@ -68,6 +70,7 @@ void setUp() { pluginConfigurationConverter = mock(PluginConfigurationConverter.class); acknowledgementSetManager = mock(DefaultAcknowledgementSetManager.class); + circuitBreaker = mock(CircuitBreaker.class); eventFactory = mock(DefaultEventFactory.class); pluginProviders = new ArrayList<>(); given(pluginProviderLoader.getPluginProviders()).willReturn(pluginProviders); @@ -94,7 +97,7 @@ void setUp() { private DefaultPluginFactory createObjectUnderTest() { return new DefaultPluginFactory( pluginProviderLoader, pluginCreator, pluginConfigurationConverter, beanFactoryProvider, eventFactory, - acknowledgementSetManager, pluginConfigurationObservableFactory); + acknowledgementSetManager, pluginConfigurationObservableFactory, circuitBreaker); } @Test 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 9f2a782ddf..48e1ce6907 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 @@ -93,13 +93,13 @@ void setUp() { when(kafkaBufferConfig.getEncryptionConfig()).thenReturn(encryptionConfig); } - private KafkaBuffer> createObjectUnderTest() { - return new KafkaBuffer<>(pluginSetting, kafkaBufferConfig, pluginFactory, acknowledgementSetManager, pluginMetrics, null, null); + private KafkaBuffer createObjectUnderTest() { + return new KafkaBuffer(pluginSetting, kafkaBufferConfig, pluginFactory, acknowledgementSetManager, pluginMetrics, null, null, null); } @Test void write_and_read() throws TimeoutException { - KafkaBuffer> objectUnderTest = createObjectUnderTest(); + KafkaBuffer objectUnderTest = createObjectUnderTest(); Record record = createRecord(); objectUnderTest.write(record, 1_000); @@ -123,7 +123,7 @@ void write_and_read() throws TimeoutException { void write_and_read_encrypted() throws TimeoutException, NoSuchAlgorithmException { when(topicConfig.getEncryptionKey()).thenReturn(createAesKey()); - KafkaBuffer> objectUnderTest = createObjectUnderTest(); + KafkaBuffer objectUnderTest = createObjectUnderTest(); Record record = createRecord(); objectUnderTest.write(record, 1_000); 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 f8fc8337d2..296571bf1c 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 @@ -8,6 +8,7 @@ import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.CheckpointState; +import org.opensearch.dataprepper.model.breaker.CircuitBreaker; import org.opensearch.dataprepper.model.codec.ByteDecoder; import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; @@ -15,6 +16,7 @@ import org.opensearch.dataprepper.model.buffer.AbstractBuffer; import org.opensearch.dataprepper.model.buffer.Buffer; 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; @@ -37,7 +39,7 @@ import java.util.concurrent.atomic.AtomicBoolean; @DataPrepperPlugin(name = "kafka", pluginType = Buffer.class, pluginConfigurationType = KafkaBufferConfig.class) -public class KafkaBuffer> extends AbstractBuffer { +public class KafkaBuffer extends AbstractBuffer> { private static final Logger LOG = LoggerFactory.getLogger(KafkaBuffer.class); static final long EXECUTOR_SERVICE_SHUTDOWN_TIMEOUT = 30L; @@ -47,7 +49,7 @@ public class KafkaBuffer> extends AbstractBuffer { static final String READ = "Read"; private final KafkaCustomProducer producer; private final List emptyCheckingConsumers; - private final AbstractBuffer innerBuffer; + private final AbstractBuffer> innerBuffer; private final ExecutorService executorService; private final Duration drainTimeout; private AtomicBoolean shutdownInProgress; @@ -56,7 +58,8 @@ public class KafkaBuffer> extends AbstractBuffer { @DataPrepperPluginConstructor public KafkaBuffer(final PluginSetting pluginSetting, final KafkaBufferConfig kafkaBufferConfig, final PluginFactory pluginFactory, final AcknowledgementSetManager acknowledgementSetManager, final PluginMetrics pluginMetrics, - final ByteDecoder byteDecoder, final AwsCredentialsSupplier awsCredentialsSupplier) { + final ByteDecoder byteDecoder, final AwsCredentialsSupplier awsCredentialsSupplier, + final CircuitBreaker circuitBreaker) { super(pluginSetting); SerializationFactory serializationFactory = new SerializationFactory(); final KafkaCustomProducerFactory kafkaCustomProducerFactory = new KafkaCustomProducerFactory(serializationFactory, awsCredentialsSupplier); @@ -69,9 +72,9 @@ public KafkaBuffer(final PluginSetting pluginSetting, final KafkaBufferConfig ka this.shutdownInProgress = new AtomicBoolean(false); final PluginMetrics consumerMetrics = PluginMetrics.fromNames(metricPrefixName + READ, pluginSetting.getPipelineName()); final List consumers = kafkaCustomConsumerFactory.createConsumersForTopic(kafkaBufferConfig, kafkaBufferConfig.getTopic(), - innerBuffer, consumerMetrics, acknowledgementSetManager, byteDecoder, shutdownInProgress, false); + innerBuffer, consumerMetrics, acknowledgementSetManager, byteDecoder, shutdownInProgress, false, circuitBreaker); emptyCheckingConsumers = kafkaCustomConsumerFactory.createConsumersForTopic(kafkaBufferConfig, kafkaBufferConfig.getTopic(), - innerBuffer, pluginMetrics, acknowledgementSetManager, byteDecoder, shutdownInProgress, false); + innerBuffer, pluginMetrics, acknowledgementSetManager, byteDecoder, shutdownInProgress, false, null); this.executorService = Executors.newFixedThreadPool(consumers.size()); consumers.forEach(this.executorService::submit); @@ -89,7 +92,7 @@ public void writeBytes(final byte[] bytes, final String key, int timeoutInMillis } @Override - public void doWrite(T record, int timeoutInMillis) throws TimeoutException { + public void doWrite(Record record, int timeoutInMillis) throws TimeoutException { try { producer.produceRecords(record); } catch (final Exception e) { @@ -104,14 +107,14 @@ public boolean isByteBuffer() { } @Override - public void doWriteAll(Collection records, int timeoutInMillis) throws Exception { - for ( T record: records ) { + public void doWriteAll(Collection> records, int timeoutInMillis) throws Exception { + for ( Record record: records ) { doWrite(record, timeoutInMillis); } } @Override - public Map.Entry, CheckpointState> doRead(int timeoutInMillis) { + public Map.Entry>, CheckpointState> doRead(int timeoutInMillis) { return innerBuffer.doRead(timeoutInMillis); } 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 73f0fe1278..d9ab131a68 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 @@ -89,6 +89,7 @@ public class KafkaCustomConsumer implements Runnable, ConsumerRebalanceListener private final boolean acknowledgementsEnabled; private final Duration acknowledgementsTimeout; private final KafkaTopicConsumerMetrics topicMetrics; + private final PauseConsumePredicate pauseConsumePredicate; private long metricsUpdatedTime; private final AtomicInteger numberOfAcksPending; private long numRecordsCommitted = 0; @@ -105,7 +106,8 @@ public KafkaCustomConsumer(final KafkaConsumer consumer, final AcknowledgementSetManager acknowledgementSetManager, final ByteDecoder byteDecoder, final KafkaTopicConsumerMetrics topicMetrics, - final TopicEmptinessMetadata topicEmptinessMetadata) { + final TopicEmptinessMetadata topicEmptinessMetadata, + final PauseConsumePredicate pauseConsumePredicate) { this.topicName = topicConfig.getName(); this.topicConfig = topicConfig; this.shutdownInProgress = shutdownInProgress; @@ -113,6 +115,7 @@ public KafkaCustomConsumer(final KafkaConsumer consumer, this.buffer = buffer; this.byteDecoder = byteDecoder; this.topicMetrics = topicMetrics; + this.pauseConsumePredicate = pauseConsumePredicate; this.topicMetrics.register(consumer); this.offsetsToCommit = new HashMap<>(); this.ownedPartitionsEpoch = new HashMap<>(); @@ -171,7 +174,7 @@ private AcknowledgementSet createAcknowledgementSet(Map void consumeRecords() throws Exception { + void consumeRecords() throws Exception { try { ConsumerRecords records = consumer.poll(Duration.ofMillis(topicConfig.getThreadWaitingTime().toMillis()/2)); @@ -331,7 +334,8 @@ public void run() { boolean retryingAfterException = false; while (!shutdownInProgress.get()) { try { - if (retryingAfterException) { + if (retryingAfterException || pauseConsumePredicate.pauseConsuming()) { + LOG.debug("Pause consuming from Kafka topic."); Thread.sleep(10000); } synchronized(this) { 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 89132ea9b3..a50702ff3a 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 @@ -20,6 +20,7 @@ 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.breaker.CircuitBreaker; import org.opensearch.dataprepper.model.buffer.Buffer; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.record.Record; @@ -73,7 +74,8 @@ public List createConsumersForTopic(final KafkaConsumerConf final AcknowledgementSetManager acknowledgementSetManager, final ByteDecoder byteDecoder, final AtomicBoolean shutdownInProgress, - final boolean topicNameInMetrics) { + final boolean topicNameInMetrics, + final CircuitBreaker circuitBreaker) { Properties authProperties = new Properties(); KafkaSecurityConfigurer.setAuthProperties(authProperties, kafkaConsumerConfig, LOG); KafkaTopicConsumerMetrics topicMetrics = new KafkaTopicConsumerMetrics(topic.getName(), pluginMetrics, topicNameInMetrics); @@ -85,6 +87,8 @@ public List createConsumersForTopic(final KafkaConsumerConf AwsContext awsContext = new AwsContext(kafkaConsumerConfig, awsCredentialsSupplier); KeyFactory keyFactory = new KeyFactory(awsContext); + PauseConsumePredicate pauseConsumePredicate = PauseConsumePredicate.circuitBreakingPredicate(circuitBreaker); + try { final int numWorkers = topic.getWorkers(); final TopicEmptinessMetadata topicEmptinessMetadata = new TopicEmptinessMetadata(); @@ -101,7 +105,7 @@ public List createConsumersForTopic(final KafkaConsumerConf final KafkaConsumer kafkaConsumer = new KafkaConsumer<>(consumerProperties, keyDeserializer, valueDeserializer); consumers.add(new KafkaCustomConsumer(kafkaConsumer, shutdownInProgress, buffer, kafkaConsumerConfig, topic, - schemaType, acknowledgementSetManager, byteDecoder, topicMetrics, topicEmptinessMetadata)); + schemaType, acknowledgementSetManager, byteDecoder, topicMetrics, topicEmptinessMetadata, pauseConsumePredicate)); }); } catch (Exception 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 new file mode 100644 index 0000000000..0d5260d92a --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/PauseConsumePredicate.java @@ -0,0 +1,44 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.consumer; + +import org.opensearch.dataprepper.model.breaker.CircuitBreaker; + +/** + * Represents if the {@link KafkaCustomConsumer} should pause consuming + * from a Kafka topic due to an external reason. + */ +@FunctionalInterface +public interface PauseConsumePredicate { + /** + * Returns whether to pause consumption of a Kafka topic. + * + * @return True if the consumer should pause. False if there is it does not need to pause. + */ + boolean pauseConsuming(); + + /** + * Returns a {@link PauseConsumePredicate} from a {@link CircuitBreaker}. This value may + * be null, in which case, it will not pause. + * + * @param circuitBreaker The {@link CircuitBreaker} or null + * @return a predicate based on the circuit breaker. + */ + static PauseConsumePredicate circuitBreakingPredicate(final CircuitBreaker circuitBreaker) { + if(circuitBreaker == null) + return noPause(); + return circuitBreaker::isOpen; + } + + /** + * Returns a {@link PauseConsumePredicate} that never pauses. + * + * @return a predicate that does not pause + */ + static PauseConsumePredicate noPause() { + return () -> false; + } +} 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 aa52a1fd1a..49fec5646a 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 @@ -37,6 +37,7 @@ import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaRegistryType; import org.opensearch.dataprepper.plugins.kafka.configuration.TopicConfig; import org.opensearch.dataprepper.plugins.kafka.consumer.KafkaCustomConsumer; +import org.opensearch.dataprepper.plugins.kafka.consumer.PauseConsumePredicate; import org.opensearch.dataprepper.plugins.kafka.consumer.TopicEmptinessMetadata; import org.opensearch.dataprepper.plugins.kafka.extension.KafkaClusterConfigSupplier; import org.opensearch.dataprepper.plugins.kafka.util.ClientDNSLookupType; @@ -141,7 +142,7 @@ public void start(Buffer> buffer) { } consumer = new KafkaCustomConsumer(kafkaConsumer, shutdownInProgress, buffer, sourceConfig, topic, schemaType, - acknowledgementSetManager, null, topicMetrics, topicEmptinessMetadata); + acknowledgementSetManager, null, topicMetrics, topicEmptinessMetadata, PauseConsumePredicate.noPause()); allTopicConsumers.add(consumer); executorService.submit(consumer); 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 ddcdb77edd..70665dcd72 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 @@ -18,6 +18,7 @@ import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.CheckpointState; import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; +import org.opensearch.dataprepper.model.breaker.CircuitBreaker; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.event.JacksonEvent; @@ -74,7 +75,7 @@ class KafkaBufferTest { private static final String TEST_GROUP_ID = "testGroupId"; - private KafkaBuffer> kafkaBuffer; + private KafkaBuffer kafkaBuffer; ExecutorService executorService; @Mock private KafkaBufferConfig bufferConfig; @@ -124,11 +125,14 @@ class KafkaBufferTest { @Mock private AwsCredentialsSupplier awsCredentialsSupplier; - public KafkaBuffer> createObjectUnderTest() { + @Mock + private CircuitBreaker circuitBreaker; + + public KafkaBuffer createObjectUnderTest() { return createObjectUnderTest(List.of(consumer)); } - public KafkaBuffer> createObjectUnderTest(final List consumers) { + public KafkaBuffer createObjectUnderTest(final List consumers) { try ( final MockedStatic executorsMockedStatic = mockStatic(Executors.class); final MockedConstruction producerFactoryMock = @@ -139,7 +143,7 @@ public KafkaBuffer> createObjectUnderTest(final List consumerFactoryMock = mockConstruction(KafkaCustomConsumerFactory.class, (mock, context) -> { consumerFactory = mock; - when(consumerFactory.createConsumersForTopic(any(), any(), any(), any(), any(), any(), any(), anyBoolean())).thenReturn(consumers); + when(consumerFactory.createConsumersForTopic(any(), any(), any(), any(), any(), any(), any(), anyBoolean(), any())).thenReturn(consumers); }); final MockedConstruction blockingBufferMock = mockConstruction(BlockingBuffer.class, (mock, context) -> { @@ -147,7 +151,7 @@ public KafkaBuffer> createObjectUnderTest(final List Executors.newFixedThreadPool(anyInt())).thenReturn(executorService); - return new KafkaBuffer>(pluginSetting, bufferConfig, pluginFactory, acknowledgementSetManager, pluginMetrics, null, awsCredentialsSupplier); + return new KafkaBuffer(pluginSetting, bufferConfig, pluginFactory, acknowledgementSetManager, pluginMetrics, null, awsCredentialsSupplier, circuitBreaker); } } 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 8c7675795c..7943c07419 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 @@ -99,6 +99,9 @@ public class KafkaCustomConsumerTest { @Mock private OffsetAndMetadata offsetAndMetadata; + @Mock + private PauseConsumePredicate pauseConsumePredicate; + private KafkaCustomConsumer consumer; private ConsumerRecords consumerRecords; @@ -166,7 +169,7 @@ public KafkaCustomConsumer createObjectUnderTest(String schemaType, boolean ackn topicEmptinessMetadata = new TopicEmptinessMetadata(); when(sourceConfig.getAcknowledgementsEnabled()).thenReturn(acknowledgementsEnabled); return new KafkaCustomConsumer(kafkaConsumer, shutdownInProgress, buffer, sourceConfig, topicConfig, schemaType, - acknowledgementSetManager, null, topicMetrics, topicEmptinessMetadata); + acknowledgementSetManager, null, topicMetrics, topicEmptinessMetadata, pauseConsumePredicate); } private BlockingBuffer> getBuffer() { 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 new file mode 100644 index 0000000000..0d2743b8ab --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/PauseConsumePredicateTest.java @@ -0,0 +1,54 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.kafka.consumer; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.opensearch.dataprepper.model.breaker.CircuitBreaker; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.notNullValue; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +class PauseConsumePredicateTest { + @Test + void noPause_returns_predicate_with_pauseConsuming_returning_false() { + final PauseConsumePredicate pauseConsumePredicate = PauseConsumePredicate.noPause(); + + assertThat(pauseConsumePredicate, notNullValue()); + + assertThat(pauseConsumePredicate.pauseConsuming(), equalTo(false)); + } + + @Test + void circuitBreakingPredicate_with_a_null_circuit_breaker_returns_predicate_with_pauseConsuming_returning_false() { + final PauseConsumePredicate pauseConsumePredicate = PauseConsumePredicate.circuitBreakingPredicate(null); + + assertThat(pauseConsumePredicate, notNullValue()); + + assertThat(pauseConsumePredicate.pauseConsuming(), equalTo(false)); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void circuitBreakingPredicate_with_a_circuit_breaker_returns_predicate_with_pauseConsuming_returning_value_of_isOpen(final boolean isOpen) { + final CircuitBreaker circuitBreaker = mock(CircuitBreaker.class); + + final PauseConsumePredicate pauseConsumePredicate = PauseConsumePredicate.circuitBreakingPredicate(circuitBreaker); + + verifyNoInteractions(circuitBreaker); + + assertThat(pauseConsumePredicate, notNullValue()); + + when(circuitBreaker.isOpen()).thenReturn(isOpen); + + assertThat(pauseConsumePredicate.pauseConsuming(), equalTo(isOpen)); + } +} \ No newline at end of file From f08c16c38728f68696572ae6ace8177be18320b1 Mon Sep 17 00:00:00 2001 From: Taylor Gray Date: Wed, 8 Nov 2023 14:08:42 -0600 Subject: [PATCH 43/53] Fix issue where we start from checkpoint for PIT with acks to instead start from beginning (#3610) Signed-off-by: Taylor Gray --- .../plugins/source/opensearch/worker/PitWorker.java | 3 ++- .../plugins/source/opensearch/worker/WorkerCommonUtils.java | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorker.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorker.java index 9050b6fd87..bc9f531d95 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorker.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/PitWorker.java @@ -156,7 +156,8 @@ private void processIndex(final SourcePartition op LOG.info("Starting processing for index: '{}'", indexName); Optional openSearchIndexProgressStateOptional = openSearchIndexPartition.getPartitionState(); - if (openSearchIndexProgressStateOptional.isEmpty()) { + // We can't checkpoint acks yet so need to restart from the beginning of index when acks are enabled for now + if (openSearchSourceConfiguration.isAcknowledgmentsEnabled() || openSearchIndexProgressStateOptional.isEmpty()) { openSearchIndexProgressStateOptional = Optional.of(initializeProgressState()); } diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtils.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtils.java index 7c5a132017..5e490f1ff1 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtils.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/source/opensearch/worker/WorkerCommonUtils.java @@ -28,7 +28,7 @@ public class WorkerCommonUtils { static final Duration BACKOFF_ON_EXCEPTION = Duration.ofSeconds(60); static final long DEFAULT_CHECKPOINT_INTERVAL_MILLS = 5 * 60_000; - static final Duration ACKNOWLEDGEMENT_SET_TIMEOUT = Duration.ofHours(1); + static final Duration ACKNOWLEDGEMENT_SET_TIMEOUT = Duration.ofMinutes(20); static final Duration STARTING_BACKOFF = Duration.ofMillis(500); static final Duration MAX_BACKOFF = Duration.ofSeconds(60); static final int BACKOFF_RATE = 2; From 621de7ba8c0536b10a7e104de67601b83a4fb32a Mon Sep 17 00:00:00 2001 From: Dinu John <86094133+dinujoh@users.noreply.github.com> Date: Wed, 8 Nov 2023 14:46:35 -0600 Subject: [PATCH 44/53] Add Kafka Producer metrics for send record failures (#3611) Signed-off-by: Dinu John <86094133+dinujoh@users.noreply.github.com> --- .../kafka/producer/KafkaCustomProducer.java | 16 ++- .../producer/KafkaCustomProducerFactory.java | 5 +- .../kafka/util/KafkaTopicProducerMetrics.java | 22 +++- .../producer/KafkaCustomProducerTest.java | 123 ++++++++++++++++-- 4 files changed, 144 insertions(+), 22 deletions(-) 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 ab745b9ebf..c661bd5889 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 @@ -78,7 +78,8 @@ public KafkaCustomProducer(final KafkaProducer producer, final DLQSink dlqSink, final ExpressionEvaluator expressionEvaluator, final String tagTargetKey, - final KafkaTopicProducerMetrics topicMetrics + final KafkaTopicProducerMetrics topicMetrics, + final SchemaService schemaService ) { this.producer = producer; this.kafkaProducerConfig = kafkaProducerConfig; @@ -88,7 +89,7 @@ public KafkaCustomProducer(final KafkaProducer producer, this.tagTargetKey = tagTargetKey; this.topicName = ObjectUtils.isEmpty(kafkaProducerConfig.getTopic()) ? null : kafkaProducerConfig.getTopic().getName(); this.serdeFormat = ObjectUtils.isEmpty(kafkaProducerConfig.getSerdeFormat()) ? null : kafkaProducerConfig.getSerdeFormat(); - this.schemaService = new SchemaService.SchemaServiceBuilder().getFetchSchemaService(topicName, kafkaProducerConfig.getSchemaConfig()).build(); + this.schemaService = schemaService; this.topicMetrics = topicMetrics; this.topicMetrics.register(this.producer); } @@ -102,8 +103,8 @@ public void produceRawData(final byte[] bytes, final String key) { send(topicName, key, bytes).get(); topicMetrics.update(producer); } catch (Exception e) { - // TODO: Metrics - LOG.error("Error occurred while publishing {}", e.getMessage()); + topicMetrics.getNumberOfRawDataSendErrors().increment(); + LOG.error("Error occurred while publishing raw data {}", e.getMessage()); } } @@ -121,8 +122,8 @@ public void produceRecords(final Record record) { } topicMetrics.update(producer); } catch (Exception e) { - // TODO: Metrics - LOG.error("Error occurred while publishing {}", e.getMessage()); + LOG.error("Error occurred while publishing record {}", e.getMessage()); + topicMetrics.getNumberOfRecordSendErrors().increment(); releaseEventHandles(false); } @@ -183,7 +184,8 @@ public boolean validateSchema(final String jsonData, final String schemaJson) th private Callback callBack(final Object dataForDlq) { return (metadata, exception) -> { if (null != exception) { - LOG.error("Error occured while publishing " + exception.getMessage()); + LOG.error("Error occurred while publishing {}", exception.getMessage()); + topicMetrics.getNumberOfRecordProcessingErrors().increment(); releaseEventHandles(false); dlqSink.perform(dataForDlq, exception); } else { 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 fa378bfc8c..c4cebed653 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 @@ -5,6 +5,7 @@ package org.opensearch.dataprepper.plugins.kafka.producer; +import org.apache.commons.lang3.ObjectUtils; import org.apache.kafka.clients.producer.KafkaProducer; import org.apache.kafka.common.serialization.Serializer; import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; @@ -63,9 +64,11 @@ public KafkaCustomProducer createProducer(final KafkaProducerConfig kafkaProduce final KafkaProducer producer = new KafkaProducer<>(properties, keyDeserializer, valueSerializer); final DLQSink dlqSink = new DLQSink(pluginFactory, kafkaProducerConfig, pluginSetting); final KafkaTopicProducerMetrics topicMetrics = new KafkaTopicProducerMetrics(topic.getName(), pluginMetrics, topicNameInMetrics); + 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); + expressionEvaluator, Objects.nonNull(sinkContext) ? sinkContext.getTagsTargetKey() : null, topicMetrics, schemaService); } private void prepareTopicAndSchema(final KafkaProducerConfig kafkaProducerConfig) { checkTopicCreationCriteriaAndCreateTopic(kafkaProducerConfig); diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicProducerMetrics.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicProducerMetrics.java index 525b85d0d5..0b2b28f60d 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicProducerMetrics.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicProducerMetrics.java @@ -28,13 +28,18 @@ public class KafkaTopicProducerMetrics { static final String BYTE_SEND_TOTAL = "byteSendTotal"; static final String RECORD_SEND_RATE_MAP_VALUE = "recordSendRate"; static final String RECORD_SEND_TOTAL_MAP_VALUE = "recordSendTotal"; - + static final String NUMBER_OF_RAW_DATA_SEND_ERRORS = "numberOfRawDataSendErrors"; + static final String NUMBER_OF_RECORD_SEND_ERRORS = "numberOfRecordSendErrors"; + static final String NUMBER_OF_RECORD_PROCESSING_ERRORS = "numberOfRecordProcessingErrors"; private final String topicName; private Map metricsNameMap; private Map> metricValues; private final PluginMetrics pluginMetrics; private final Counter numberOfRecordsSent; private final Counter numberOfBytesSent; + private final Counter numberOfRawDataSendErrors; + private final Counter numberOfRecordSendErrors; + private final Counter numberOfRecordProcessingErrors; public KafkaTopicProducerMetrics(final String topicName, final PluginMetrics pluginMetrics, final boolean topicNameInMetrics) { @@ -44,6 +49,9 @@ public KafkaTopicProducerMetrics(final String topicName, final PluginMetrics plu initializeMetricNamesMap(topicNameInMetrics); this.numberOfRecordsSent = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_RECORDS_SENT, topicNameInMetrics)); this.numberOfBytesSent = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_BYTES_SENT, topicNameInMetrics)); + this.numberOfRawDataSendErrors = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_RAW_DATA_SEND_ERRORS, topicNameInMetrics)); + this.numberOfRecordSendErrors = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_RECORD_SEND_ERRORS, topicNameInMetrics)); + this.numberOfRecordProcessingErrors = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_RECORD_PROCESSING_ERRORS, topicNameInMetrics)); } private void initializeMetricNamesMap(final boolean topicNameInMetrics) { @@ -85,6 +93,18 @@ Counter getNumberOfBytesSent() { return numberOfBytesSent; } + public Counter getNumberOfRawDataSendErrors() { + return numberOfRawDataSendErrors; + } + + public Counter getNumberOfRecordSendErrors() { + return numberOfRecordSendErrors; + } + + public Counter getNumberOfRecordProcessingErrors() { + return numberOfRecordProcessingErrors; + } + private String getTopicMetricName(final String metricName, final boolean topicNameInMetrics) { if (topicNameInMetrics) { return "topic." + topicName + "." + metricName; diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerTest.java index 2a03e5d9db..54a955d3ac 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducerTest.java @@ -7,12 +7,13 @@ import com.github.fge.jsonschema.core.exceptions.ProcessingException; import io.confluent.kafka.schemaregistry.client.SchemaMetadata; -import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; +import io.micrometer.core.instrument.Counter; import org.apache.avro.Schema; import org.apache.avro.generic.GenericRecord; import org.apache.kafka.clients.producer.Callback; import org.apache.kafka.clients.producer.KafkaProducer; import org.apache.kafka.clients.producer.ProducerRecord; +import org.apache.kafka.common.KafkaException; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -29,6 +30,7 @@ import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.TopicProducerConfig; import org.opensearch.dataprepper.plugins.kafka.configuration.SchemaConfig; +import org.opensearch.dataprepper.plugins.kafka.service.SchemaService; import org.opensearch.dataprepper.plugins.kafka.sink.DLQSink; import org.opensearch.dataprepper.plugins.kafka.util.KafkaTopicProducerMetrics; @@ -36,7 +38,6 @@ import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Method; import java.util.UUID; -import java.util.concurrent.ExecutionException; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -44,6 +45,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; @@ -59,6 +61,15 @@ public class KafkaCustomProducerTest { @Mock private KafkaTopicProducerMetrics kafkaTopicProducerMetrics; + @Mock + private SchemaService schemaService; + @Mock + private Counter numberOfRawDataSendErrors; + @Mock + private Counter numberOfRecordSendErrors; + @Mock + private Counter numberOfRecordProcessingError; + private Record record; KafkaCustomProducer sinkProducer; @@ -83,10 +94,52 @@ record = new Record<>(event); } @Test - public void producePlainTextRecordsTest() throws ExecutionException, InterruptedException { + public void produceRawDataTest() { + when(kafkaSinkConfig.getSerdeFormat()).thenReturn("plaintext"); + KafkaProducer kafkaProducer = mock(KafkaProducer.class); + producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), + null, kafkaTopicProducerMetrics, schemaService); + when(kafkaTopicProducerMetrics.getNumberOfRawDataSendErrors()).thenReturn(numberOfRawDataSendErrors); + sinkProducer = spy(producer); + final String key = UUID.randomUUID().toString(); + final byte[] byteData = record.getData().toJsonString().getBytes(); + sinkProducer.produceRawData(byteData, key); + verify(sinkProducer).produceRawData(record.getData().toJsonString().getBytes(), key); + final ArgumentCaptor recordArgumentCaptor = ArgumentCaptor.forClass(ProducerRecord.class); + verify(kafkaProducer).send(recordArgumentCaptor.capture(), any(Callback.class)); + assertEquals(recordArgumentCaptor.getValue().topic(), kafkaSinkConfig.getTopic().getName()); + assertEquals(recordArgumentCaptor.getValue().value(), byteData); + assertEquals(recordArgumentCaptor.getValue().key(), key); + verifyNoInteractions(numberOfRecordSendErrors); + } + + @Test + public void produceRawData_sendError() { + when(kafkaSinkConfig.getSerdeFormat()).thenReturn("plaintext"); + KafkaProducer kafkaProducer = mock(KafkaProducer.class); + producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), + null, kafkaTopicProducerMetrics, schemaService); + when(kafkaTopicProducerMetrics.getNumberOfRawDataSendErrors()).thenReturn(numberOfRawDataSendErrors); + when(kafkaProducer.send(any(ProducerRecord.class), any(Callback.class))).thenThrow(new KafkaException()); + sinkProducer = spy(producer); + final String key = UUID.randomUUID().toString(); + final byte[] byteData = record.getData().toJsonString().getBytes(); + sinkProducer.produceRawData(byteData, key); + verify(sinkProducer).produceRawData(record.getData().toJsonString().getBytes(), key); + final ArgumentCaptor recordArgumentCaptor = ArgumentCaptor.forClass(ProducerRecord.class); + verify(kafkaProducer).send(recordArgumentCaptor.capture(), any(Callback.class)); + assertEquals(recordArgumentCaptor.getValue().topic(), kafkaSinkConfig.getTopic().getName()); + assertEquals(recordArgumentCaptor.getValue().value(), byteData); + assertEquals(recordArgumentCaptor.getValue().key(), key); + verify(numberOfRawDataSendErrors).increment(); + } + + @Test + public void producePlainTextRecords() { when(kafkaSinkConfig.getSerdeFormat()).thenReturn("plaintext"); KafkaProducer kafkaProducer = mock(KafkaProducer.class); - producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), null, kafkaTopicProducerMetrics); + producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), + null, kafkaTopicProducerMetrics, schemaService); sinkProducer = spy(producer); sinkProducer.produceRecords(record); verify(sinkProducer).produceRecords(record); @@ -94,14 +147,53 @@ public void producePlainTextRecordsTest() throws ExecutionException, Interrupted verify(kafkaProducer).send(recordArgumentCaptor.capture(), any(Callback.class)); assertEquals(recordArgumentCaptor.getValue().topic(), kafkaSinkConfig.getTopic().getName()); assertEquals(recordArgumentCaptor.getValue().value(), record.getData().toJsonString()); + verifyNoInteractions(numberOfRecordSendErrors); + } + @Test + public void producePlainTextRecords_sendError() { + when(kafkaSinkConfig.getSerdeFormat()).thenReturn("plaintext"); + KafkaProducer kafkaProducer = mock(KafkaProducer.class); + producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), + null, kafkaTopicProducerMetrics, schemaService); + when(kafkaTopicProducerMetrics.getNumberOfRecordSendErrors()).thenReturn(numberOfRecordSendErrors); + when(kafkaProducer.send(any(ProducerRecord.class), any(Callback.class))).thenThrow(new KafkaException()); + sinkProducer = spy(producer); + sinkProducer.produceRecords(record); + verify(sinkProducer).produceRecords(record); + final ArgumentCaptor recordArgumentCaptor = ArgumentCaptor.forClass(ProducerRecord.class); + verify(kafkaProducer).send(recordArgumentCaptor.capture(), any(Callback.class)); + assertEquals(recordArgumentCaptor.getValue().topic(), kafkaSinkConfig.getTopic().getName()); + assertEquals(recordArgumentCaptor.getValue().value(), record.getData().toJsonString()); + verify(numberOfRecordSendErrors).increment(); } @Test - public void produceJsonRecordsTest() throws RestClientException, IOException { + public void producePlainTextRecords_callbackException() { + when(kafkaSinkConfig.getSerdeFormat()).thenReturn("plaintext"); + KafkaProducer kafkaProducer = mock(KafkaProducer.class); + producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), + null, kafkaTopicProducerMetrics, schemaService); + when(kafkaTopicProducerMetrics.getNumberOfRecordProcessingErrors()).thenReturn(numberOfRecordProcessingError); + sinkProducer = spy(producer); + sinkProducer.produceRecords(record); + verify(sinkProducer).produceRecords(record); + final ArgumentCaptor recordArgumentCaptor = ArgumentCaptor.forClass(ProducerRecord.class); + final ArgumentCaptor callbackArgumentCaptor = ArgumentCaptor.forClass(Callback.class); + verify(kafkaProducer).send(recordArgumentCaptor.capture(), callbackArgumentCaptor.capture()); + assertEquals(recordArgumentCaptor.getValue().topic(), kafkaSinkConfig.getTopic().getName()); + assertEquals(recordArgumentCaptor.getValue().value(), record.getData().toJsonString()); + final Callback actualCallback = callbackArgumentCaptor.getValue(); + actualCallback.onCompletion(null, new RuntimeException()); + verify(numberOfRecordProcessingError).increment(); + } + + @Test + public void produceJsonRecordsTest() { when(kafkaSinkConfig.getSerdeFormat()).thenReturn("JSON"); KafkaProducer kafkaProducer = mock(KafkaProducer.class); - producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), null, kafkaTopicProducerMetrics); + producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), + null, kafkaTopicProducerMetrics, schemaService); SchemaMetadata schemaMetadata = mock(SchemaMetadata.class); String jsonSchema = "{\n" + " \"$schema\": \"http://json-schema.org/draft-07/schema#\",\n" + @@ -121,31 +213,35 @@ public void produceJsonRecordsTest() throws RestClientException, IOException { verify(kafkaProducer).send(recordArgumentCaptor.capture(), any(Callback.class)); assertEquals(recordArgumentCaptor.getValue().topic(), kafkaSinkConfig.getTopic().getName()); assertEquals(recordArgumentCaptor.getValue().value(), record.getData().getJsonNode()); + verifyNoInteractions(numberOfRecordSendErrors); } @Test - public void produceAvroRecordsTest() throws Exception { + public void produceAvroRecordsTest() { when(kafkaSinkConfig.getSerdeFormat()).thenReturn("AVRO"); KafkaProducer kafkaProducer = mock(KafkaProducer.class); - producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), null, kafkaTopicProducerMetrics); - SchemaMetadata schemaMetadata = mock(SchemaMetadata.class); + producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), + null, kafkaTopicProducerMetrics, schemaService); + when(kafkaTopicProducerMetrics.getNumberOfRecordSendErrors()).thenReturn(numberOfRecordSendErrors); String avroSchema = "{\"type\":\"record\",\"name\":\"MyMessage\",\"fields\":[{\"name\":\"message\",\"type\":\"string\"}]}"; - when(schemaMetadata.getSchema()).thenReturn(avroSchema); + when(schemaService.getSchema(kafkaSinkConfig.getTopic().getName())).thenReturn(new Schema.Parser().parse(avroSchema)); sinkProducer = spy(producer); sinkProducer.produceRecords(record); verify(sinkProducer).produceRecords(record); - + verifyNoInteractions(numberOfRecordSendErrors); } @Test public void testGetGenericRecord() throws NoSuchMethodException, InvocationTargetException, IllegalAccessException { KafkaProducer kafkaProducer = mock(KafkaProducer.class); - producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), null, kafkaTopicProducerMetrics); + producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), + null, kafkaTopicProducerMetrics, schemaService); final Schema schema = createMockSchema(); Method privateMethod = KafkaCustomProducer.class.getDeclaredMethod("getGenericRecord", Event.class, Schema.class); privateMethod.setAccessible(true); GenericRecord result = (GenericRecord) privateMethod.invoke(producer, event, schema); Assertions.assertNotNull(result); + verifyNoInteractions(numberOfRecordSendErrors); } private Schema createMockSchema() { @@ -158,7 +254,8 @@ private Schema createMockSchema() { public void validateSchema() throws IOException, ProcessingException { when(kafkaSinkConfig.getSerdeFormat()).thenReturn("avro"); KafkaProducer kafkaProducer = mock(KafkaProducer.class); - producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), null, kafkaTopicProducerMetrics); + producer = new KafkaCustomProducer(kafkaProducer, kafkaSinkConfig, dlqSink, mock(ExpressionEvaluator.class), + null, kafkaTopicProducerMetrics, schemaService); String jsonSchema = "{\"type\": \"object\",\"properties\": {\"Year\": {\"type\": \"string\"},\"Age\": {\"type\": \"string\"},\"Ethnic\": {\"type\":\"string\",\"default\": null}}}"; String jsonSchema2 = "{\"type\": \"object\",\"properties\": {\"Year\": {\"type\": \"string\"},\"Age\": {\"type\": \"string\"},\"Ethnic\": {\"type\":\"string\",\"default\": null}}}"; assertTrue(producer.validateSchema(jsonSchema, jsonSchema2)); From d3acfa72009a78202b44445b1f73792813dc7635 Mon Sep 17 00:00:00 2001 From: kkondaka <41027584+kkondaka@users.noreply.github.com> Date: Thu, 9 Nov 2023 10:09:20 -0800 Subject: [PATCH 45/53] Fix race condition in DefaultEventHandle (#3618) Signed-off-by: Krishna Kondaka Co-authored-by: Krishna Kondaka --- .../dataprepper/model/event/DefaultEventHandle.java | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java index cb7c12da84..743309bf75 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventHandle.java @@ -56,8 +56,10 @@ public Instant getExternalOriginationTime() { @Override public void release(boolean result) { - for (final BiConsumer consumer: releaseConsumers) { - consumer.accept(this, result); + synchronized (releaseConsumers) { + for (final BiConsumer consumer: releaseConsumers) { + consumer.accept(this, result); + } } AcknowledgementSet acknowledgementSet = getAcknowledgementSet(); if (acknowledgementSet != null) { @@ -67,6 +69,8 @@ public void release(boolean result) { @Override public void onRelease(BiConsumer releaseConsumer) { - releaseConsumers.add(releaseConsumer); + synchronized (releaseConsumers) { + releaseConsumers.add(releaseConsumer); + } } } From 85aba4d9fdb1c1ba75779f658a0cd24bfe570033 Mon Sep 17 00:00:00 2001 From: Andriy Redko Date: Thu, 9 Nov 2023 18:21:03 -0500 Subject: [PATCH 46/53] Add Java 11/17/21 matrix for build, test and performance test checks (#3622) Signed-off-by: Andriy Redko --- .github/workflows/gradle.yml | 2 +- .github/workflows/performance-test-compile.yml | 2 +- build.gradle | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/gradle.yml b/.github/workflows/gradle.yml index a083132203..1f375461d9 100644 --- a/.github/workflows/gradle.yml +++ b/.github/workflows/gradle.yml @@ -12,7 +12,7 @@ jobs: build: strategy: matrix: - java: [11, 17] + java: [11, 17, 21] fail-fast: false runs-on: ubuntu-latest diff --git a/.github/workflows/performance-test-compile.yml b/.github/workflows/performance-test-compile.yml index 6c8ea1b9e4..f7c7840761 100644 --- a/.github/workflows/performance-test-compile.yml +++ b/.github/workflows/performance-test-compile.yml @@ -19,7 +19,7 @@ jobs: build: strategy: matrix: - java: [11, 17] + java: [11, 17, 21] runs-on: ubuntu-latest diff --git a/build.gradle b/build.gradle index db2552e774..3d720a6f34 100644 --- a/build.gradle +++ b/build.gradle @@ -26,7 +26,7 @@ buildscript { } plugins { - id 'com.diffplug.spotless' version '6.11.0' + id 'com.diffplug.spotless' version '6.22.0' id 'io.spring.dependency-management' version '1.1.0' } From bc504fd1a7842a8832bda5e70f40c847944fd668 Mon Sep 17 00:00:00 2001 From: kkondaka <41027584+kkondaka@users.noreply.github.com> Date: Thu, 9 Nov 2023 15:32:41 -0800 Subject: [PATCH 47/53] Add support for OTEL metrics source to use Kafka buffer (#3539) * Add support for OTEL metrics source to use Kafka buffer Signed-off-by: Krishna Kondaka * Added tests and fixed test failures Signed-off-by: Krishna Kondaka --------- Signed-off-by: Krishna Kondaka Co-authored-by: Krishna Kondaka --- .../metric/JacksonExponentialHistogram.java | 4 +- .../model/metric/JacksonHistogram.java | 2 +- .../model/metric/JacksonMetric.java | 2 +- .../otelmetrics/OTelMetricsRawProcessor.java | 286 ++++-------------- .../OtelMetricsRawProcessorConfig.java | 4 +- .../otelmetrics/MetricsPluginSumTest.java | 15 +- .../otelmetrics/MetricsPluginSummaryTest.java | 8 +- .../otel-metrics-source/build.gradle | 1 + .../otelmetrics/OTelMetricsGrpcService.java | 6 +- .../source/otelmetrics/OTelMetricsSource.java | 9 + .../plugins/otel/codec/OTelMetricDecoder.java | 39 +++ .../plugins/otel/codec/OTelProtoCodec.java | 276 +++++++++++++++++ .../otel/codec/OTelProtoCodecTest.java | 85 +++++- .../test/resources/test-gauge-metrics.json | 44 +++ .../resources/test-histogram-metrics.json | 50 +++ .../src/test/resources/test-sum-metrics.json | 45 +++ 16 files changed, 626 insertions(+), 250 deletions(-) create mode 100644 data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OTelMetricDecoder.java create mode 100644 data-prepper-plugins/otel-proto-common/src/test/resources/test-gauge-metrics.json create mode 100644 data-prepper-plugins/otel-proto-common/src/test/resources/test-histogram-metrics.json create mode 100644 data-prepper-plugins/otel-proto-common/src/test/resources/test-sum-metrics.json 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 b52c850e46..b865ce0eb5 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 @@ -27,8 +27,8 @@ public class JacksonExponentialHistogram extends JacksonMetric implements Expone private static final String SCALE_KEY = "scale"; private static final String AGGREGATION_TEMPORALITY_KEY = "aggregationTemporality"; private static final String ZERO_COUNT_KEY = "zeroCount"; - private static final String POSITIVE_BUCKETS_KEY = "positiveBuckets"; - private static final String NEGATIVE_BUCKETS_KEY = "negativeBuckets"; + public static final String POSITIVE_BUCKETS_KEY = "positiveBuckets"; + public static final String NEGATIVE_BUCKETS_KEY = "negativeBuckets"; private static final String NEGATIVE_KEY = "negative"; private static final String POSITIVE_KEY = "positive"; private static final String NEGATIVE_OFFSET_KEY = "negativeOffset"; 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 0209f7012d..f9e066875d 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 @@ -29,7 +29,7 @@ public class JacksonHistogram extends JacksonMetric implements Histogram { private static final String AGGREGATION_TEMPORALITY_KEY = "aggregationTemporality"; private static final String BUCKET_COUNTS_KEY = "bucketCounts"; private static final String EXPLICIT_BOUNDS_COUNT_KEY = "explicitBoundsCount"; - private static final String BUCKETS_KEY = "buckets"; + public static final String BUCKETS_KEY = "buckets"; private static final String BUCKET_COUNTS_LIST_KEY = "bucketCountsList"; private static final String EXPLICIT_BOUNDS_KEY = "explicitBounds"; 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 0ab81ed7e0..8d8ebf0f87 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 @@ -28,7 +28,7 @@ public abstract class JacksonMetric extends JacksonEvent implements Metric { protected static final String SERVICE_NAME_KEY = "serviceName"; protected static final String KIND_KEY = "kind"; protected static final String UNIT_KEY = "unit"; - protected static final String ATTRIBUTES_KEY = "attributes"; + public static final String ATTRIBUTES_KEY = "attributes"; protected static final String SCHEMA_URL_KEY = "schemaUrl"; protected static final String EXEMPLARS_KEY = "exemplars"; protected static final String FLAGS_KEY = "flags"; 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 e180e48ac3..679eef3224 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 @@ -6,35 +6,32 @@ package org.opensearch.dataprepper.plugins.processor.otelmetrics; import io.opentelemetry.proto.collector.metrics.v1.ExportMetricsServiceRequest; -import io.opentelemetry.proto.metrics.v1.InstrumentationLibraryMetrics; -import io.opentelemetry.proto.metrics.v1.ResourceMetrics; -import io.opentelemetry.proto.metrics.v1.ScopeMetrics; import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; import org.opensearch.dataprepper.model.configuration.PluginSetting; -import org.opensearch.dataprepper.model.metric.JacksonExponentialHistogram; -import org.opensearch.dataprepper.model.metric.JacksonGauge; -import org.opensearch.dataprepper.model.metric.JacksonHistogram; -import org.opensearch.dataprepper.model.metric.JacksonSum; -import org.opensearch.dataprepper.model.metric.JacksonSummary; +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.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; +import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoCodec; import io.micrometer.core.instrument.Counter; + import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.util.ArrayList; -import java.util.Arrays; import java.util.Collection; -import java.util.List; import java.util.Map; -import java.util.stream.Collectors; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; @DataPrepperPlugin(name = "otel_metrics", deprecatedName = "otel_metrics_raw_processor", pluginType = Processor.class, pluginConfigurationType = OtelMetricsRawProcessorConfig.class) -public class OTelMetricsRawProcessor extends AbstractProcessor, Record> { +public class OTelMetricsRawProcessor extends AbstractProcessor, Record> { private static final Logger LOG = LoggerFactory.getLogger(OTelMetricsRawProcessor.class); public static final String RECORDS_DROPPED_METRICS_RAW = "recordsDroppedMetricsRaw"; @@ -52,240 +49,61 @@ public OTelMetricsRawProcessor(PluginSetting pluginSetting, final OtelMetricsRaw this.flattenAttributesFlag = otelMetricsRawProcessorConfig.getFlattenAttributesFlag(); } - @Override - public Collection> doExecute(Collection> records) { - Collection> recordsOut = new ArrayList<>(); - for (Record ets : records) { - for (ResourceMetrics rs : ets.getData().getResourceMetricsList()) { - final String schemaUrl = rs.getSchemaUrl(); - final Map resourceAttributes = OTelProtoCodec.getResourceAttributes(rs.getResource()); - final String serviceName = OTelProtoCodec.getServiceName(rs.getResource()).orElse(null); + private void modifyRecord(Record record, + boolean flattenAttributes, + boolean calcualteHistogramBuckets, + boolean calcualteExponentialHistogramBuckets) { + Event event = (Event)record.getData(); - for (InstrumentationLibraryMetrics is : rs.getInstrumentationLibraryMetricsList()) { - final Map ils = OTelProtoCodec.getInstrumentationLibraryAttributes(is.getInstrumentationLibrary()); - recordsOut.addAll(processMetricsList(is.getMetricsList(), serviceName, ils, resourceAttributes, schemaUrl)); - } + if (flattenAttributes) { + Map attributes = event.get(ATTRIBUTES_KEY, Map.class); - for (ScopeMetrics sm : rs.getScopeMetricsList()) { - final Map ils = OTelProtoCodec.getInstrumentationScopeAttributes(sm.getScope()); - recordsOut.addAll(processMetricsList(sm.getMetricsList(), serviceName, ils, resourceAttributes, schemaUrl)); - } + for (Map.Entry entry : attributes.entrySet()) { + event.put(entry.getKey(), entry.getValue()); } + event.delete(ATTRIBUTES_KEY); } - return recordsOut; - } - - private List> processMetricsList(final List metricsList, - final String serviceName, - final Map ils, - final Map resourceAttributes, - final String schemaUrl) { - List> recordsOut = new ArrayList<>(); - for (io.opentelemetry.proto.metrics.v1.Metric metric : metricsList) { - try { - if (metric.hasGauge()) { - recordsOut.addAll(mapGauge(metric, serviceName, ils, resourceAttributes, schemaUrl)); - } else if (metric.hasSum()) { - recordsOut.addAll(mapSum(metric, serviceName, ils, resourceAttributes, schemaUrl)); - } else if (metric.hasSummary()) { - recordsOut.addAll(mapSummary(metric, serviceName, ils, resourceAttributes, schemaUrl)); - } else if (metric.hasHistogram()) { - recordsOut.addAll(mapHistogram(metric, serviceName, ils, resourceAttributes, schemaUrl)); - } else if (metric.hasExponentialHistogram()) { - recordsOut.addAll(mapExponentialHistogram(metric, serviceName, ils, resourceAttributes, schemaUrl)); - } - } catch (Exception e) { - LOG.warn("Error while processing metrics", e); - recordsDroppedMetricsRawCounter.increment(); + if (!calcualteHistogramBuckets && event.get(BUCKETS_KEY, List.class) != null) { + event.delete(BUCKETS_KEY); + } + if (!calcualteExponentialHistogramBuckets) { + if (event.get(POSITIVE_BUCKETS_KEY, List.class) != null) { + event.delete(POSITIVE_BUCKETS_KEY); + } + if (event.get(NEGATIVE_BUCKETS_KEY, List.class) != null) { + event.delete(NEGATIVE_BUCKETS_KEY); } } - return recordsOut; - } - - private List> mapGauge(io.opentelemetry.proto.metrics.v1.Metric metric, - String serviceName, - final Map ils, - final Map resourceAttributes, - final String schemaUrl) { - return metric.getGauge().getDataPointsList().stream() - .map(dp -> JacksonGauge.builder() - .withUnit(metric.getUnit()) - .withName(metric.getName()) - .withDescription(metric.getDescription()) - .withStartTime(OTelProtoCodec.getStartTimeISO8601(dp)) - .withTime(OTelProtoCodec.getTimeISO8601(dp)) - .withServiceName(serviceName) - .withValue(OTelProtoCodec.getValueAsDouble(dp)) - .withAttributes(OTelProtoCodec.mergeAllAttributes( - Arrays.asList( - OTelProtoCodec.convertKeysOfDataPointAttributes(dp), - resourceAttributes, - ils - ) - )) - .withSchemaUrl(schemaUrl) - .withExemplars(OTelProtoCodec.convertExemplars(dp.getExemplarsList())) - .withFlags(dp.getFlags()) - .build(flattenAttributesFlag)) - .map(Record::new) - .collect(Collectors.toList()); } - private List> mapSum(final io.opentelemetry.proto.metrics.v1.Metric metric, - final String serviceName, - final Map ils, - final Map resourceAttributes, - final String schemaUrl) { - return metric.getSum().getDataPointsList().stream() - .map(dp -> JacksonSum.builder() - .withUnit(metric.getUnit()) - .withName(metric.getName()) - .withDescription(metric.getDescription()) - .withStartTime(OTelProtoCodec.getStartTimeISO8601(dp)) - .withTime(OTelProtoCodec.getTimeISO8601(dp)) - .withServiceName(serviceName) - .withIsMonotonic(metric.getSum().getIsMonotonic()) - .withValue(OTelProtoCodec.getValueAsDouble(dp)) - .withAggregationTemporality(metric.getSum().getAggregationTemporality().toString()) - .withAttributes(OTelProtoCodec.mergeAllAttributes( - Arrays.asList( - OTelProtoCodec.convertKeysOfDataPointAttributes(dp), - resourceAttributes, - ils - ) - )) - .withSchemaUrl(schemaUrl) - .withExemplars(OTelProtoCodec.convertExemplars(dp.getExemplarsList())) - .withFlags(dp.getFlags()) - .build(flattenAttributesFlag)) - .map(Record::new) - .collect(Collectors.toList()); - } - - private List> mapSummary(final io.opentelemetry.proto.metrics.v1.Metric metric, - final String serviceName, - final Map ils, - final Map resourceAttributes, - final String schemaUrl) { - return metric.getSummary().getDataPointsList().stream() - .map(dp -> JacksonSummary.builder() - .withUnit(metric.getUnit()) - .withName(metric.getName()) - .withDescription(metric.getDescription()) - .withStartTime(OTelProtoCodec.convertUnixNanosToISO8601(dp.getStartTimeUnixNano())) - .withTime(OTelProtoCodec.convertUnixNanosToISO8601(dp.getTimeUnixNano())) - .withServiceName(serviceName) - .withCount(dp.getCount()) - .withSum(dp.getSum()) - .withQuantiles(OTelProtoCodec.getQuantileValues(dp.getQuantileValuesList())) - .withQuantilesValueCount(dp.getQuantileValuesCount()) - .withAttributes(OTelProtoCodec.mergeAllAttributes( - Arrays.asList( - OTelProtoCodec.unpackKeyValueList(dp.getAttributesList()), - resourceAttributes, - ils - ) - )) - .withSchemaUrl(schemaUrl) - .withFlags(dp.getFlags()) - .build(flattenAttributesFlag)) - .map(Record::new) - .collect(Collectors.toList()); - } - - private List> mapHistogram(final io.opentelemetry.proto.metrics.v1.Metric metric, - final String serviceName, - final Map ils, - final Map resourceAttributes, - final String schemaUrl) { - return metric.getHistogram().getDataPointsList().stream() - .map(dp -> { - JacksonHistogram.Builder builder = JacksonHistogram.builder() - .withUnit(metric.getUnit()) - .withName(metric.getName()) - .withDescription(metric.getDescription()) - .withStartTime(OTelProtoCodec.convertUnixNanosToISO8601(dp.getStartTimeUnixNano())) - .withTime(OTelProtoCodec.convertUnixNanosToISO8601(dp.getTimeUnixNano())) - .withServiceName(serviceName) - .withSum(dp.getSum()) - .withCount(dp.getCount()) - .withBucketCount(dp.getBucketCountsCount()) - .withExplicitBoundsCount(dp.getExplicitBoundsCount()) - .withAggregationTemporality(metric.getHistogram().getAggregationTemporality().toString()) - .withBucketCountsList(dp.getBucketCountsList()) - .withExplicitBoundsList(dp.getExplicitBoundsList()) - .withAttributes(OTelProtoCodec.mergeAllAttributes( - Arrays.asList( - OTelProtoCodec.unpackKeyValueList(dp.getAttributesList()), - resourceAttributes, - ils - ) - )) - .withSchemaUrl(schemaUrl) - .withExemplars(OTelProtoCodec.convertExemplars(dp.getExemplarsList())) - .withFlags(dp.getFlags()); - if (otelMetricsRawProcessorConfig.getCalculateHistogramBuckets()) { - builder.withBuckets(OTelProtoCodec.createBuckets(dp.getBucketCountsList(), dp.getExplicitBoundsList())); - } - JacksonHistogram jh = builder.build(flattenAttributesFlag); - return jh; - - }) - .map(Record::new) - .collect(Collectors.toList()); - } - - private List> mapExponentialHistogram(io.opentelemetry.proto.metrics.v1.Metric metric, String serviceName, Map ils, Map resourceAttributes, String schemaUrl) { - return metric.getExponentialHistogram().getDataPointsList().stream() - .filter(dp -> { - if (otelMetricsRawProcessorConfig.getCalculateExponentialHistogramBuckets() && - otelMetricsRawProcessorConfig.getExponentialHistogramMaxAllowedScale() < Math.abs(dp.getScale())){ - LOG.error("Exponential histogram can not be processed since its scale of {} is bigger than the configured max of {}.", dp.getScale(), otelMetricsRawProcessorConfig.getExponentialHistogramMaxAllowedScale()); - return false; - } else { - return true; - } - }) - .map(dp -> { - JacksonExponentialHistogram.Builder builder = JacksonExponentialHistogram.builder() - .withUnit(metric.getUnit()) - .withName(metric.getName()) - .withDescription(metric.getDescription()) - .withStartTime(OTelProtoCodec.convertUnixNanosToISO8601(dp.getStartTimeUnixNano())) - .withTime(OTelProtoCodec.convertUnixNanosToISO8601(dp.getTimeUnixNano())) - .withServiceName(serviceName) - .withSum(dp.getSum()) - .withCount(dp.getCount()) - .withZeroCount(dp.getZeroCount()) - .withScale(dp.getScale()) - .withPositive(dp.getPositive().getBucketCountsList()) - .withPositiveOffset(dp.getPositive().getOffset()) - .withNegative(dp.getNegative().getBucketCountsList()) - .withNegativeOffset(dp.getNegative().getOffset()) - .withAggregationTemporality(metric.getHistogram().getAggregationTemporality().toString()) - .withAttributes(OTelProtoCodec.mergeAllAttributes( - Arrays.asList( - OTelProtoCodec.unpackKeyValueList(dp.getAttributesList()), - resourceAttributes, - ils - ) - )) - .withSchemaUrl(schemaUrl) - .withExemplars(OTelProtoCodec.convertExemplars(dp.getExemplarsList())) - .withFlags(dp.getFlags()); + @Override + public Collection> doExecute(Collection> records) { + Collection> recordsOut = new ArrayList<>(); + OTelProtoCodec.OTelProtoDecoder otelProtoDecoder = new OTelProtoCodec.OTelProtoDecoder(); + AtomicInteger droppedCounter = new AtomicInteger(0); + + for (Record rec : records) { + if ((rec.getData() instanceof Event)) { + Record newRecord = (Record)rec; + if (otelMetricsRawProcessorConfig.getFlattenAttributesFlag() || + !otelMetricsRawProcessorConfig.getCalculateHistogramBuckets() || + !otelMetricsRawProcessorConfig.getCalculateExponentialHistogramBuckets()) { + modifyRecord(newRecord, otelMetricsRawProcessorConfig.getFlattenAttributesFlag(), otelMetricsRawProcessorConfig.getCalculateHistogramBuckets(), otelMetricsRawProcessorConfig.getCalculateExponentialHistogramBuckets()); + } + recordsOut.add(newRecord); + } - if (otelMetricsRawProcessorConfig.getCalculateExponentialHistogramBuckets()) { - builder.withPositiveBuckets(OTelProtoCodec.createExponentialBuckets(dp.getPositive(), dp.getScale())); - builder.withNegativeBuckets(OTelProtoCodec.createExponentialBuckets(dp.getNegative(), dp.getScale())); - } + if (!(rec.getData() instanceof ExportMetricsServiceRequest)) { + continue; + } - return builder.build(flattenAttributesFlag); - }) - .map(Record::new) - .collect(Collectors.toList()); + ExportMetricsServiceRequest request = ((Record)rec).getData(); + recordsOut.addAll(otelProtoDecoder.parseExportMetricsServiceRequest(request, droppedCounter, otelMetricsRawProcessorConfig.getExponentialHistogramMaxAllowedScale(), otelMetricsRawProcessorConfig.getCalculateHistogramBuckets(), otelMetricsRawProcessorConfig.getCalculateExponentialHistogramBuckets(), flattenAttributesFlag)); + } + recordsDroppedMetricsRawCounter.increment(droppedCounter.get()); + return recordsOut; } - @Override public void prepareForShutdown() { } diff --git a/data-prepper-plugins/otel-metrics-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/otelmetrics/OtelMetricsRawProcessorConfig.java b/data-prepper-plugins/otel-metrics-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/otelmetrics/OtelMetricsRawProcessorConfig.java index b74460d0f7..9935cc9218 100644 --- a/data-prepper-plugins/otel-metrics-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/otelmetrics/OtelMetricsRawProcessorConfig.java +++ b/data-prepper-plugins/otel-metrics-raw-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/otelmetrics/OtelMetricsRawProcessorConfig.java @@ -4,6 +4,8 @@ */ package org.opensearch.dataprepper.plugins.processor.otelmetrics; + +import static org.opensearch.dataprepper.plugins.otel.codec.OTelProtoCodec.DEFAULT_EXPONENTIAL_HISTOGRAM_MAX_ALLOWED_SCALE; import com.fasterxml.jackson.annotation.JsonProperty; public class OtelMetricsRawProcessorConfig { @@ -15,7 +17,7 @@ public class OtelMetricsRawProcessorConfig { private Boolean calculateExponentialHistogramBuckets = true; - private Integer exponentialHistogramMaxAllowedScale = 10; + private Integer exponentialHistogramMaxAllowedScale = DEFAULT_EXPONENTIAL_HISTOGRAM_MAX_ALLOWED_SCALE; public Boolean getCalculateExponentialHistogramBuckets() { return calculateExponentialHistogramBuckets; diff --git a/data-prepper-plugins/otel-metrics-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/otelmetrics/MetricsPluginSumTest.java b/data-prepper-plugins/otel-metrics-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/otelmetrics/MetricsPluginSumTest.java index e202219ae1..9c6341f5da 100644 --- a/data-prepper-plugins/otel-metrics-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/otelmetrics/MetricsPluginSumTest.java +++ b/data-prepper-plugins/otel-metrics-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/otelmetrics/MetricsPluginSumTest.java @@ -24,11 +24,12 @@ import org.junit.runner.RunWith; import org.mockito.junit.MockitoJUnitRunner; import org.opensearch.dataprepper.model.configuration.PluginSetting; -import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.metric.JacksonMetric; import java.util.Arrays; import java.util.Collections; +import java.util.Collection; import java.util.List; import java.util.Map; @@ -102,8 +103,9 @@ public void test() throws JsonProcessingException { Record record = new Record<>(exportMetricRequest); - List> rec = (List>) rawProcessor.doExecute(Arrays.asList(record)); - Record firstRecord = rec.get(0); + Collection> records = Arrays.asList((Record)record); + List> outputRecords = (List>)rawProcessor.doExecute(records); + Record firstRecord = (Record)outputRecords.get(0); ObjectMapper objectMapper = new ObjectMapper(); Map map = objectMapper.readValue(firstRecord.getData().toJsonString(), Map.class); @@ -182,8 +184,11 @@ public void missingNameInvalidMetricTest() throws JsonProcessingException { Record record = new Record<>(exportMetricRequest); Record invalidRecord = new Record<>(exportMetricRequestWithInvalidMetric); - List> rec = (List>) rawProcessor.doExecute(Arrays.asList(record, invalidRecord)); - org.hamcrest.MatcherAssert.assertThat(rec.size(), equalTo(1)); + Collection> records = Arrays.asList((Record)record, invalidRecord); + List> outputRecords = (List>)rawProcessor.doExecute(records); + + //List> rec = (List>) rawProcessor.doExecute(Arrays.asList(record, invalidRecord)); + org.hamcrest.MatcherAssert.assertThat(outputRecords.size(), equalTo(1)); } private void assertSumProcessing(Map map) { diff --git a/data-prepper-plugins/otel-metrics-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/otelmetrics/MetricsPluginSummaryTest.java b/data-prepper-plugins/otel-metrics-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/otelmetrics/MetricsPluginSummaryTest.java index d915a48e37..234765e740 100644 --- a/data-prepper-plugins/otel-metrics-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/otelmetrics/MetricsPluginSummaryTest.java +++ b/data-prepper-plugins/otel-metrics-raw-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/otelmetrics/MetricsPluginSummaryTest.java @@ -21,11 +21,12 @@ import org.junit.runner.RunWith; import org.mockito.junit.MockitoJUnitRunner; import org.opensearch.dataprepper.model.configuration.PluginSetting; -import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.metric.JacksonMetric; import java.util.Arrays; import java.util.Collections; +import java.util.Collection; import java.util.List; import java.util.Map; @@ -83,8 +84,9 @@ public void testSummaryProcessing() throws JsonProcessingException { Record record = new Record<>(exportMetricRequest); - List> rec = (List>) rawProcessor.doExecute(Arrays.asList(record)); - Record firstRecord = rec.get(0); + Collection> records = Arrays.asList((Record)record); + List> outputRecords = (List>)rawProcessor.doExecute(records); + Record firstRecord = (Record)outputRecords.get(0); ObjectMapper objectMapper = new ObjectMapper(); Map map = objectMapper.readValue(firstRecord.getData().toJsonString(), Map.class); diff --git a/data-prepper-plugins/otel-metrics-source/build.gradle b/data-prepper-plugins/otel-metrics-source/build.gradle index abe038c645..6372395a81 100644 --- a/data-prepper-plugins/otel-metrics-source/build.gradle +++ b/data-prepper-plugins/otel-metrics-source/build.gradle @@ -13,6 +13,7 @@ dependencies { implementation project(':data-prepper-plugins:blocking-buffer') implementation libs.commons.codec implementation project(':data-prepper-plugins:armeria-common') + implementation project(':data-prepper-plugins:otel-proto-common') testImplementation project(':data-prepper-api').sourceSets.test.output implementation libs.opentelemetry.proto implementation libs.commons.io 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 8da1ad63f7..0177a57584 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 @@ -70,7 +70,11 @@ public void export(final ExportMetricsServiceRequest request, final StreamObserv private void processRequest(final ExportMetricsServiceRequest request, final StreamObserver responseObserver) { try { - buffer.write(new Record<>(request), bufferWriteTimeoutInMillis); + if (buffer.isByteBuffer()) { + buffer.writeBytes(request.toByteArray(), null, bufferWriteTimeoutInMillis); + } else { + buffer.write(new Record<>(request), bufferWriteTimeoutInMillis); + } } catch (Exception e) { if (ServiceRequestContext.current().isTimedOut()) { LOG.warn("Exception writing to buffer but request already timed out.", e); 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 fcfd9524d9..33c4023e67 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 @@ -32,6 +32,8 @@ import org.opensearch.dataprepper.model.plugin.PluginFactory; import org.opensearch.dataprepper.model.record.Record; 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.certificate.CertificateProvider; import org.opensearch.dataprepper.plugins.certificate.model.Certificate; import org.opensearch.dataprepper.plugins.health.HealthGrpcService; @@ -62,6 +64,7 @@ public class OTelMetricsSource implements Source> 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); + for (Record record: records) { + final JacksonEvent event = JacksonEvent.fromEvent(record.getData()); + eventConsumer.accept(new Record<>(event)); + } + } + +} 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 16f596c989..dba17b0851 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 @@ -8,6 +8,7 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.protobuf.ByteString; +import io.opentelemetry.proto.collector.metrics.v1.ExportMetricsServiceRequest; import io.opentelemetry.proto.collector.logs.v1.ExportLogsServiceRequest; import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest; import io.opentelemetry.proto.common.v1.AnyValue; @@ -19,6 +20,9 @@ import io.opentelemetry.proto.metrics.v1.ExponentialHistogramDataPoint; import io.opentelemetry.proto.metrics.v1.NumberDataPoint; import io.opentelemetry.proto.metrics.v1.SummaryDataPoint; +import io.opentelemetry.proto.metrics.v1.InstrumentationLibraryMetrics; +import io.opentelemetry.proto.metrics.v1.ResourceMetrics; +import io.opentelemetry.proto.metrics.v1.ScopeMetrics; import io.opentelemetry.proto.resource.v1.Resource; import io.opentelemetry.proto.trace.v1.InstrumentationLibrarySpans; import io.opentelemetry.proto.trace.v1.ResourceSpans; @@ -26,6 +30,7 @@ import io.opentelemetry.proto.trace.v1.Status; import org.apache.commons.codec.DecoderException; import org.apache.commons.codec.binary.Hex; +import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.log.JacksonOtelLog; import org.opensearch.dataprepper.model.log.OpenTelemetryLog; import org.opensearch.dataprepper.model.metric.Bucket; @@ -34,6 +39,12 @@ import org.opensearch.dataprepper.model.metric.DefaultQuantile; import org.opensearch.dataprepper.model.metric.Exemplar; import org.opensearch.dataprepper.model.metric.Quantile; +import org.opensearch.dataprepper.model.metric.JacksonExponentialHistogram; +import org.opensearch.dataprepper.model.metric.JacksonGauge; +import org.opensearch.dataprepper.model.metric.JacksonHistogram; +import org.opensearch.dataprepper.model.metric.JacksonSum; +import org.opensearch.dataprepper.model.metric.JacksonSummary; +import org.opensearch.dataprepper.model.metric.Metric; import org.opensearch.dataprepper.model.trace.DefaultLink; import org.opensearch.dataprepper.model.trace.DefaultSpanEvent; import org.slf4j.Logger; @@ -57,6 +68,7 @@ import java.util.Objects; import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -67,6 +79,7 @@ public class OTelProtoCodec { private static final Logger LOG = LoggerFactory.getLogger(OTelProtoCodec.class); + public static final int DEFAULT_EXPONENTIAL_HISTOGRAM_MAX_ALLOWED_SCALE = 10; private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); private static final long NANO_MULTIPLIER = 1_000 * 1_000 * 1_000; @@ -147,6 +160,7 @@ public static long timeISO8601ToNanos(final String timeISO08601) { } public static class OTelProtoDecoder { + public List parseExportTraceServiceRequest(final ExportTraceServiceRequest exportTraceServiceRequest) { return exportTraceServiceRequest.getResourceSpansList().stream() .flatMap(rs -> parseResourceSpans(rs).stream()).collect(Collectors.toList()); @@ -445,6 +459,268 @@ protected Optional getServiceName(final Resource resource) { && !keyValue.getValue().getStringValue().isEmpty() ).findFirst().map(i -> i.getValue().getStringValue()); } + + public Collection> parseExportMetricsServiceRequest( + final ExportMetricsServiceRequest request, + AtomicInteger droppedCounter, + final Integer exponentialHistogramMaxAllowedScale, + final boolean calculateHistogramBuckets, + final boolean calculateExponentialHistogramBuckets, + final boolean flattenAttributes) { + Collection> recordsOut = new ArrayList<>(); + for (ResourceMetrics rs : request.getResourceMetricsList()) { + final String schemaUrl = rs.getSchemaUrl(); + final Map resourceAttributes = OTelProtoCodec.getResourceAttributes(rs.getResource()); + final String serviceName = OTelProtoCodec.getServiceName(rs.getResource()).orElse(null); + + 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)); + } + + 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)); + } + } + return recordsOut; + } + + private List> processMetricsList( + final List metricsList, + final String serviceName, + final Map ils, + final Map resourceAttributes, + final String schemaUrl, + AtomicInteger droppedCounter, + final Integer exponentialHistogramMaxAllowedScale, + final boolean calculateHistogramBuckets, + final boolean calculateExponentialHistogramBuckets, + final boolean flattenAttributes) { + List> recordsOut = new ArrayList<>(); + for (io.opentelemetry.proto.metrics.v1.Metric metric : metricsList) { + try { + if (metric.hasGauge()) { + recordsOut.addAll(mapGauge(metric, serviceName, ils, resourceAttributes, schemaUrl, flattenAttributes)); + } else if (metric.hasSum()) { + recordsOut.addAll(mapSum(metric, serviceName, ils, resourceAttributes, schemaUrl, flattenAttributes)); + } else if (metric.hasSummary()) { + recordsOut.addAll(mapSummary(metric, serviceName, ils, resourceAttributes, schemaUrl, flattenAttributes)); + } else if (metric.hasHistogram()) { + recordsOut.addAll(mapHistogram(metric, serviceName, ils, resourceAttributes, schemaUrl, calculateHistogramBuckets, flattenAttributes)); + } else if (metric.hasExponentialHistogram()) { + recordsOut.addAll(mapExponentialHistogram(metric, serviceName, ils, resourceAttributes, schemaUrl, exponentialHistogramMaxAllowedScale, calculateExponentialHistogramBuckets, flattenAttributes)); + } + } catch (Exception e) { + LOG.warn("Error while processing metrics", e); + droppedCounter.incrementAndGet(); + } + } + return recordsOut; + } + + private List> mapGauge( + io.opentelemetry.proto.metrics.v1.Metric metric, + String serviceName, + final Map ils, + final Map resourceAttributes, + final String schemaUrl, + final boolean flattenAttributes) { + return metric.getGauge().getDataPointsList().stream() + .map(dp -> JacksonGauge.builder() + .withUnit(metric.getUnit()) + .withName(metric.getName()) + .withDescription(metric.getDescription()) + .withStartTime(OTelProtoCodec.getStartTimeISO8601(dp)) + .withTime(OTelProtoCodec.getTimeISO8601(dp)) + .withServiceName(serviceName) + .withValue(OTelProtoCodec.getValueAsDouble(dp)) + .withAttributes(OTelProtoCodec.mergeAllAttributes( + Arrays.asList( + OTelProtoCodec.convertKeysOfDataPointAttributes(dp), + resourceAttributes, + ils + ) + )) + .withSchemaUrl(schemaUrl) + .withExemplars(OTelProtoCodec.convertExemplars(dp.getExemplarsList())) + .withFlags(dp.getFlags()) + .build(flattenAttributes)) + .map(Record::new) + .collect(Collectors.toList()); + } + + private List> mapSum( + final io.opentelemetry.proto.metrics.v1.Metric metric, + final String serviceName, + final Map ils, + final Map resourceAttributes, + final String schemaUrl, + final boolean flattenAttributes) { + return metric.getSum().getDataPointsList().stream() + .map(dp -> JacksonSum.builder() + .withUnit(metric.getUnit()) + .withName(metric.getName()) + .withDescription(metric.getDescription()) + .withStartTime(OTelProtoCodec.getStartTimeISO8601(dp)) + .withTime(OTelProtoCodec.getTimeISO8601(dp)) + .withServiceName(serviceName) + .withIsMonotonic(metric.getSum().getIsMonotonic()) + .withValue(OTelProtoCodec.getValueAsDouble(dp)) + .withAggregationTemporality(metric.getSum().getAggregationTemporality().toString()) + .withAttributes(OTelProtoCodec.mergeAllAttributes( + Arrays.asList( + OTelProtoCodec.convertKeysOfDataPointAttributes(dp), + resourceAttributes, + ils + ) + )) + .withSchemaUrl(schemaUrl) + .withExemplars(OTelProtoCodec.convertExemplars(dp.getExemplarsList())) + .withFlags(dp.getFlags()) + .build(flattenAttributes)) + .map(Record::new) + .collect(Collectors.toList()); + } + + private List> mapSummary( + final io.opentelemetry.proto.metrics.v1.Metric metric, + final String serviceName, + final Map ils, + final Map resourceAttributes, + final String schemaUrl, + final boolean flattenAttributes) { + return metric.getSummary().getDataPointsList().stream() + .map(dp -> JacksonSummary.builder() + .withUnit(metric.getUnit()) + .withName(metric.getName()) + .withDescription(metric.getDescription()) + .withStartTime(OTelProtoCodec.convertUnixNanosToISO8601(dp.getStartTimeUnixNano())) + .withTime(OTelProtoCodec.convertUnixNanosToISO8601(dp.getTimeUnixNano())) + .withServiceName(serviceName) + .withCount(dp.getCount()) + .withSum(dp.getSum()) + .withQuantiles(OTelProtoCodec.getQuantileValues(dp.getQuantileValuesList())) + .withQuantilesValueCount(dp.getQuantileValuesCount()) + .withAttributes(OTelProtoCodec.mergeAllAttributes( + Arrays.asList( + OTelProtoCodec.unpackKeyValueList(dp.getAttributesList()), + resourceAttributes, + ils + ) + )) + .withSchemaUrl(schemaUrl) + .withFlags(dp.getFlags()) + .build(flattenAttributes)) + .map(Record::new) + .collect(Collectors.toList()); + } + + private List> mapHistogram( + final io.opentelemetry.proto.metrics.v1.Metric metric, + final String serviceName, + final Map ils, + final Map resourceAttributes, + final String schemaUrl, + final boolean calculateHistogramBuckets, + final boolean flattenAttributes) { + return metric.getHistogram().getDataPointsList().stream() + .map(dp -> { + JacksonHistogram.Builder builder = JacksonHistogram.builder() + .withUnit(metric.getUnit()) + .withName(metric.getName()) + .withDescription(metric.getDescription()) + .withStartTime(OTelProtoCodec.convertUnixNanosToISO8601(dp.getStartTimeUnixNano())) + .withTime(OTelProtoCodec.convertUnixNanosToISO8601(dp.getTimeUnixNano())) + .withServiceName(serviceName) + .withSum(dp.getSum()) + .withCount(dp.getCount()) + .withBucketCount(dp.getBucketCountsCount()) + .withExplicitBoundsCount(dp.getExplicitBoundsCount()) + .withAggregationTemporality(metric.getHistogram().getAggregationTemporality().toString()) + .withBucketCountsList(dp.getBucketCountsList()) + .withExplicitBoundsList(dp.getExplicitBoundsList()) + .withAttributes(OTelProtoCodec.mergeAllAttributes( + Arrays.asList( + OTelProtoCodec.unpackKeyValueList(dp.getAttributesList()), + resourceAttributes, + ils + ) + )) + .withSchemaUrl(schemaUrl) + .withExemplars(OTelProtoCodec.convertExemplars(dp.getExemplarsList())) + .withFlags(dp.getFlags()); + if (calculateHistogramBuckets) { + builder.withBuckets(OTelProtoCodec.createBuckets(dp.getBucketCountsList(), dp.getExplicitBoundsList())); + } + JacksonHistogram jh = builder.build(flattenAttributes); + return jh; + + }) + .map(Record::new) + .collect(Collectors.toList()); + } + + private List> mapExponentialHistogram( + final io.opentelemetry.proto.metrics.v1.Metric metric, + final String serviceName, + final Map ils, + final Map resourceAttributes, + final String schemaUrl, + final Integer exponentialHistogramMaxAllowedScale, + final boolean calculateExponentialHistogramBuckets, + final boolean flattenAttributes) { + return metric.getExponentialHistogram() + .getDataPointsList() + .stream() + .filter(dp -> { + if (calculateExponentialHistogramBuckets && + exponentialHistogramMaxAllowedScale < Math.abs(dp.getScale())){ + LOG.error("Exponential histogram can not be processed since its scale of {} is bigger than the configured max of {}.", dp.getScale(), exponentialHistogramMaxAllowedScale); + return false; + } else { + return true; + } + }) + .map(dp -> { + JacksonExponentialHistogram.Builder builder = JacksonExponentialHistogram.builder() + .withUnit(metric.getUnit()) + .withName(metric.getName()) + .withDescription(metric.getDescription()) + .withStartTime(OTelProtoCodec.convertUnixNanosToISO8601(dp.getStartTimeUnixNano())) + .withTime(OTelProtoCodec.convertUnixNanosToISO8601(dp.getTimeUnixNano())) + .withServiceName(serviceName) + .withSum(dp.getSum()) + .withCount(dp.getCount()) + .withZeroCount(dp.getZeroCount()) + .withScale(dp.getScale()) + .withPositive(dp.getPositive().getBucketCountsList()) + .withPositiveOffset(dp.getPositive().getOffset()) + .withNegative(dp.getNegative().getBucketCountsList()) + .withNegativeOffset(dp.getNegative().getOffset()) + .withAggregationTemporality(metric.getHistogram().getAggregationTemporality().toString()) + .withAttributes(OTelProtoCodec.mergeAllAttributes( + Arrays.asList( + OTelProtoCodec.unpackKeyValueList(dp.getAttributesList()), + resourceAttributes, + ils + ) + )) + .withSchemaUrl(schemaUrl) + .withExemplars(OTelProtoCodec.convertExemplars(dp.getExemplarsList())) + .withFlags(dp.getFlags()); + + if (calculateExponentialHistogramBuckets) { + builder.withPositiveBuckets(OTelProtoCodec.createExponentialBuckets(dp.getPositive(), dp.getScale())); + builder.withNegativeBuckets(OTelProtoCodec.createExponentialBuckets(dp.getNegative(), dp.getScale())); + } + + return builder.build(flattenAttributes); + }) + .map(Record::new) + .collect(Collectors.toList()); + } + } public static class OTelProtoEncoder { 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 afc4cf2ab3..6c9a167ad5 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 @@ -12,6 +12,7 @@ import com.google.protobuf.util.JsonFormat; import io.opentelemetry.proto.collector.logs.v1.ExportLogsServiceRequest; import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest; +import io.opentelemetry.proto.collector.metrics.v1.ExportMetricsServiceRequest; import io.opentelemetry.proto.common.v1.AnyValue; import io.opentelemetry.proto.common.v1.ArrayValue; import io.opentelemetry.proto.common.v1.InstrumentationLibrary; @@ -34,6 +35,12 @@ import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.model.log.OpenTelemetryLog; import org.opensearch.dataprepper.model.metric.Bucket; +import org.opensearch.dataprepper.model.metric.Metric; +import org.opensearch.dataprepper.model.metric.JacksonMetric; +import org.opensearch.dataprepper.model.metric.JacksonGauge; +import org.opensearch.dataprepper.model.metric.JacksonSum; +import org.opensearch.dataprepper.model.metric.JacksonHistogram; +import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.trace.DefaultLink; import org.opensearch.dataprepper.model.trace.DefaultSpanEvent; import org.opensearch.dataprepper.model.trace.DefaultTraceGroupFields; @@ -56,12 +63,14 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.Collection; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Random; import java.util.UUID; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; import static org.assertj.core.api.Assertions.entry; @@ -84,9 +93,10 @@ public class OTelProtoCodecTest { private static final String TEST_REQUEST_BOTH_SPAN_TYPES_JSON_FILE = "test-request-both-span-types.json"; private static final String TEST_REQUEST_NO_SPANS_JSON_FILE = "test-request-no-spans.json"; private static final String TEST_SPAN_EVENT_JSON_FILE = "test-span-event.json"; - + private static final String TEST_REQUEST_GAUGE_METRICS_JSON_FILE = "test-gauge-metrics.json"; + private static final String TEST_REQUEST_SUM_METRICS_JSON_FILE = "test-sum-metrics.json"; + private static final String TEST_REQUEST_HISTOGRAM_METRICS_JSON_FILE = "test-histogram-metrics.json"; private static final String TEST_REQUEST_LOGS_JSON_FILE = "test-request-log.json"; - private static final String TEST_REQUEST_LOGS_IS_JSON_FILE = "test-request-log-is.json"; @@ -124,6 +134,12 @@ private ExportLogsServiceRequest buildExportLogsServiceRequestFromJsonFile(Strin return builder.build(); } + private ExportMetricsServiceRequest buildExportMetricsServiceRequestFromJsonFile(String requestJsonFileName) throws IOException { + final ExportMetricsServiceRequest.Builder builder = ExportMetricsServiceRequest.newBuilder(); + JsonFormat.parser().merge(getFileAsJsonString(requestJsonFileName), builder); + return builder.build(); + } + private String getFileAsJsonString(String requestJsonFileName) throws IOException { final StringBuilder jsonBuilder = new StringBuilder(); try (final InputStream inputStream = Objects.requireNonNull( @@ -460,6 +476,71 @@ public void testParseExportLogsServiceRequest_InstrumentationLibrarySpans() thro validateSpans(spans); } + @Test + 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); + + validateGaugeMetricRequest(metrics); + } + + @Test + 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); + validateSumMetricRequest(metrics); + } + + @Test + 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); + validateHistogramMetricRequest(metrics); + } + + private void validateGaugeMetricRequest(Collection> metrics) { + assertThat(metrics.size(), equalTo(1)); + Record record = ((List>)metrics).get(0); + JacksonMetric metric = (JacksonMetric) record.getData(); + assertThat(metric.getKind(), equalTo(Metric.KIND.GAUGE.toString())); + assertThat(metric.getUnit(), equalTo("1")); + assertThat(metric.getName(), equalTo("counter-int")); + JacksonGauge gauge = (JacksonGauge)metric; + assertThat(gauge.getValue(), equalTo(123.0)); + } + + private void validateSumMetricRequest(Collection> metrics) { + assertThat(metrics.size(), equalTo(1)); + Record record = ((List>)metrics).get(0); + JacksonMetric metric = (JacksonMetric) record.getData(); + assertThat(metric.getKind(), equalTo(Metric.KIND.SUM.toString())); + assertThat(metric.getUnit(), equalTo("1")); + assertThat(metric.getName(), equalTo("sum-int")); + JacksonSum sum = (JacksonSum)metric; + assertThat(sum.getValue(), equalTo(456.0)); + } + + private void validateHistogramMetricRequest(Collection> metrics) { + assertThat(metrics.size(), equalTo(1)); + Record record = ((List>)metrics).get(0); + JacksonMetric metric = (JacksonMetric) record.getData(); + assertThat(metric.getKind(), equalTo(Metric.KIND.HISTOGRAM.toString())); + assertThat(metric.getUnit(), equalTo("1")); + assertThat(metric.getName(), equalTo("histogram-int")); + JacksonHistogram histogram = (JacksonHistogram)metric; + assertThat(histogram.getSum(), equalTo(100.0)); + assertThat(histogram.getCount(), equalTo(30L)); + assertThat(histogram.getExemplars(), equalTo(Collections.emptyList())); + assertThat(histogram.getExplicitBoundsList(), equalTo(List.of(1.0, 2.0, 3.0, 4.0))); + assertThat(histogram.getExplicitBoundsCount(), equalTo(4)); + assertThat(histogram.getBucketCountsList(), equalTo(List.of(3L, 5L, 15L, 6L, 1L))); + assertThat(histogram.getBucketCount(), equalTo(5)); + assertThat(histogram.getAggregationTemporality(), equalTo("AGGREGATION_TEMPORALITY_CUMULATIVE")); + } + } @Nested diff --git a/data-prepper-plugins/otel-proto-common/src/test/resources/test-gauge-metrics.json b/data-prepper-plugins/otel-proto-common/src/test/resources/test-gauge-metrics.json new file mode 100644 index 0000000000..abca7b0b29 --- /dev/null +++ b/data-prepper-plugins/otel-proto-common/src/test/resources/test-gauge-metrics.json @@ -0,0 +1,44 @@ +{ + "resourceMetrics": [ + { + "resource": { + "attributes": [ + { + "key": "resource-attr", + "value": { + "stringValue": "resource-attr-val-1" + } + } + ] + }, + "scopeMetrics": [ + { + "scope": {}, + "metrics": [ + { + "name": "counter-int", + "unit": 1, + "gauge": { + "dataPoints": [ + { + "attributes": [ + { + "key": "label-1", + "value": { + "stringValue": "label-value-1" + } + } + ], + "startTimeUnixNano": "1581452773000000789", + "timeUnixNano": "1581452773000000789", + "asInt": "123" + } + ] + } + } + ] + } + ] + } + ] +} diff --git a/data-prepper-plugins/otel-proto-common/src/test/resources/test-histogram-metrics.json b/data-prepper-plugins/otel-proto-common/src/test/resources/test-histogram-metrics.json new file mode 100644 index 0000000000..1220de6214 --- /dev/null +++ b/data-prepper-plugins/otel-proto-common/src/test/resources/test-histogram-metrics.json @@ -0,0 +1,50 @@ +{ + "resourceMetrics": [ + { + "resource": { + "attributes": [ + { + "key": "resource-attr", + "value": { + "stringValue": "resource-attr-val-1" + } + } + ] + }, + "scopeMetrics": [ + { + "scope": {}, + "metrics": [ + { + "name": "histogram-int", + "unit": 1, + "histogram": { + "dataPoints": [ + { + "attributes": [ + { + "key": "label-1", + "value": { + "stringValue": "label-value-1" + } + } + ], + "startTimeUnixNano": "1581452773000000789", + "timeUnixNano": "1581452773000000789", + "count": "30", + "sum": "100", + "bucket_counts": [3, 5, 15, 6, 1], + "explicit_bounds": [1.0, 2.0, 3.0, 4.0], + "exemplars": [] + } + ], + "aggregationTemporality":"2" + } + } + ] + } + ] + } + ] +} + diff --git a/data-prepper-plugins/otel-proto-common/src/test/resources/test-sum-metrics.json b/data-prepper-plugins/otel-proto-common/src/test/resources/test-sum-metrics.json new file mode 100644 index 0000000000..97d3560cc6 --- /dev/null +++ b/data-prepper-plugins/otel-proto-common/src/test/resources/test-sum-metrics.json @@ -0,0 +1,45 @@ +{ + "resourceMetrics": [ + { + "resource": { + "attributes": [ + { + "key": "resource-attr", + "value": { + "stringValue": "resource-attr-val-1" + } + } + ] + }, + "scopeMetrics": [ + { + "scope": {}, + "metrics": [ + { + "name": "sum-int", + "unit": 1, + "sum": { + "dataPoints": [ + { + "attributes": [ + { + "key": "label-1", + "value": { + "stringValue": "label-value-1" + } + } + ], + "startTimeUnixNano": "1581452773000000789", + "timeUnixNano": "1581452773000000789", + "asInt": "456" + } + ] + } + } + ] + } + ] + } + ] +} + From ca6f86aae8607d3098860e9b27ca1c0ea94ef604 Mon Sep 17 00:00:00 2001 From: kkondaka <41027584+kkondaka@users.noreply.github.com> Date: Thu, 9 Nov 2023 21:09:45 -0800 Subject: [PATCH 48/53] Add splitExportTraceServiceRequest API to OTelProtoDecoder (#3600) * Add splitExportTraceServiceRequest API to OTelProtoDecoder Signed-off-by: Krishna Kondaka * Renamed the API Signed-off-by: Krishna Kondaka * Fixed code and modified test case Signed-off-by: Krishna Kondaka * Fixed check style test Signed-off-by: Krishna Kondaka --------- Signed-off-by: Krishna Kondaka Co-authored-by: Krishna Kondaka --- .../plugins/otel/codec/OTelProtoCodec.java | 99 +++++++ .../otel/codec/OTelProtoCodecTest.java | 54 ++++ .../test-request-multiple-traces.json | 249 ++++++++++++++++++ 3 files changed, 402 insertions(+) create mode 100644 data-prepper-plugins/otel-proto-common/src/test/resources/test-request-multiple-traces.json 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 dba17b0851..29a58be6df 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 @@ -166,6 +166,26 @@ public List parseExportTraceServiceRequest(final ExportTraceServiceRequest .flatMap(rs -> parseResourceSpans(rs).stream()).collect(Collectors.toList()); } + public Map splitExportTraceServiceRequestByTraceId(final ExportTraceServiceRequest exportTraceServiceRequest) { + Map result = new HashMap<>(); + Map resultBuilderMap = new HashMap<>(); + for (final ResourceSpans resourceSpans: exportTraceServiceRequest.getResourceSpansList()) { + for (Map.Entry entry: splitResourceSpansByTraceId(resourceSpans).entrySet()) { + String traceId = entry.getKey(); + + if (resultBuilderMap.containsKey(traceId)) { + resultBuilderMap.get(traceId).addResourceSpans(entry.getValue()); + } else { + resultBuilderMap.put(traceId, ExportTraceServiceRequest.newBuilder().addResourceSpans(entry.getValue())); + } + } + } + for (Map.Entry entry: resultBuilderMap.entrySet()) { + result.put(entry.getKey(), entry.getValue().build()); + } + return result; + } + public List parseExportLogsServiceRequest(final ExportLogsServiceRequest exportLogsServiceRequest) { return exportLogsServiceRequest.getResourceLogsList().stream() .flatMap(rs -> parseResourceLogs(rs).stream()).collect(Collectors.toList()); @@ -199,6 +219,38 @@ protected Collection parseResourceLogs(ResourceLogs rs) { return Stream.concat(mappedInstrumentationLibraryLogs, mappedScopeListLogs).collect(Collectors.toList()); } + protected Map splitResourceSpansByTraceId(final ResourceSpans resourceSpans) { + final Resource resource = resourceSpans.getResource(); + Map result = new HashMap<>(); + Map resultBuilderMap = new HashMap<>(); + + if (resourceSpans.getScopeSpansList().size() > 0) { + for (Map.Entry> entry: splitScopeSpansByTraceId(resourceSpans.getScopeSpansList()).entrySet()) { + ResourceSpans.Builder b = ResourceSpans.newBuilder().setResource(resource).addAllScopeSpans(entry.getValue()); + resultBuilderMap.put(entry.getKey(), b); + } + } + + if (resourceSpans.getInstrumentationLibrarySpansList().size() > 0) { + for (Map.Entry> entry: splitInstrumentationLibrarySpansByTraceId(resourceSpans.getInstrumentationLibrarySpansList()).entrySet()) { + ResourceSpans.Builder resourceSpansBuilder; + String traceId = entry.getKey(); + if (resultBuilderMap.containsKey(traceId)) { + resourceSpansBuilder = resultBuilderMap.get(traceId); + } else { + resourceSpansBuilder = ResourceSpans.newBuilder().setResource(resource); + resultBuilderMap.put(traceId, resourceSpansBuilder); + } + resourceSpansBuilder.addAllInstrumentationLibrarySpans(entry.getValue()); + } + } + for (Map.Entry entry: resultBuilderMap.entrySet()) { + result.put(entry.getKey(), entry.getValue().build()); + } + + return result; + } + protected List parseResourceSpans(final ResourceSpans resourceSpans) { final String serviceName = getServiceName(resourceSpans.getResource()).orElse(null); final Map resourceAttributes = getResourceAttributes(resourceSpans.getResource()); @@ -223,6 +275,21 @@ private List parseScopeSpans(final List scopeSpansList, final .collect(Collectors.toList()); } + private Map> splitScopeSpansByTraceId(final List scopeSpansList) { + Map> result = new HashMap<>(); + for (ScopeSpans ss: scopeSpansList) { + for (Map.Entry> entry: splitSpansByTraceId(ss.getSpansList()).entrySet()) { + ScopeSpans.Builder scopeSpansBuilder = ScopeSpans.newBuilder().setScope(ss.getScope()).addAllSpans(entry.getValue()); + String traceId = entry.getKey(); + if (!result.containsKey(traceId)) { + result.put(traceId, new ArrayList<>()); + } + result.get(traceId).add(scopeSpansBuilder.build()); + } + } + return result; + } + private List parseInstrumentationLibrarySpans(final List instrumentationLibrarySpansList, final String serviceName, final Map resourceAttributes) { return instrumentationLibrarySpansList.stream() @@ -233,6 +300,38 @@ private List parseInstrumentationLibrarySpans(final List> splitInstrumentationLibrarySpansByTraceId(final List instrumentationLibrarySpansList) { + Map> result = new HashMap<>(); + for (InstrumentationLibrarySpans is: instrumentationLibrarySpansList) { + for (Map.Entry> entry: splitSpansByTraceId(is.getSpansList()).entrySet()) { + String traceId = entry.getKey(); + InstrumentationLibrarySpans.Builder ilSpansBuilder = InstrumentationLibrarySpans.newBuilder().setInstrumentationLibrary(is.getInstrumentationLibrary()).addAllSpans(entry.getValue()); + if (!result.containsKey(traceId)) { + result.put(traceId, new ArrayList<>()); + } + result.get(traceId).add(ilSpansBuilder.build()); + } + } + return result; + } + + + private Map> splitSpansByTraceId(final List spans) { + Map> result = new HashMap<>(); + for (io.opentelemetry.proto.trace.v1.Span span: spans) { + String traceId = convertByteStringToString(span.getTraceId()); + List spanList; + if (result.containsKey(traceId)) { + spanList = result.get(traceId); + } else { + spanList = new ArrayList<>(); + result.put(traceId, spanList); + } + spanList.add(span); + } + return result; + } + private List parseSpans(final List spans, final T scope, final Function> scopeAttributesGetter, final String serviceName, final Map resourceAttributes) { 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 6c9a167ad5..a68bf57e6d 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 @@ -80,6 +80,7 @@ import static org.hamcrest.CoreMatchers.notNullValue; import static org.hamcrest.CoreMatchers.nullValue; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -98,6 +99,7 @@ public class OTelProtoCodecTest { private static final String TEST_REQUEST_HISTOGRAM_METRICS_JSON_FILE = "test-histogram-metrics.json"; private static final String TEST_REQUEST_LOGS_JSON_FILE = "test-request-log.json"; private static final String TEST_REQUEST_LOGS_IS_JSON_FILE = "test-request-log-is.json"; + private static final String TEST_REQUEST_MULTIPLE_TRACES_FILE = "test-request-multiple-traces.json"; private static final Long TIME = TimeUnit.MILLISECONDS.toNanos(ZonedDateTime.of( @@ -153,6 +155,58 @@ private String getFileAsJsonString(String requestJsonFileName) throws IOExceptio @Nested class OTelProtoDecoderTest { + @Test + public void testSplitExportTraceServiceRequestWithMultipleTraces() throws Exception { + final ExportTraceServiceRequest exportTraceServiceRequest = buildExportTraceServiceRequestFromJsonFile(TEST_REQUEST_MULTIPLE_TRACES_FILE); + final Map map = decoderUnderTest.splitExportTraceServiceRequestByTraceId(exportTraceServiceRequest); + assertThat(map.size(), is(equalTo(3))); + for (Map.Entry entry: map.entrySet()) { + String expectedTraceId = new String(Hex.decodeHex(entry.getKey()), StandardCharsets.UTF_8); + ExportTraceServiceRequest request = entry.getValue(); + if (expectedTraceId.equals("TRACEID1")) { + assertThat(request.getResourceSpansList().size(), equalTo(1)); + ResourceSpans rs = request.getResourceSpansList().get(0); + assertThat(rs.getScopeSpansList().size(), equalTo(1)); + assertThat(rs.getInstrumentationLibrarySpansList().size(), equalTo(0)); + ScopeSpans ss = rs.getScopeSpansList().get(0); + assertThat(ss.getSpansList().size(), equalTo(1)); + io.opentelemetry.proto.trace.v1.Span span = ss.getSpansList().get(0); + String spanId = span.getSpanId().toStringUtf8(); + assertTrue(spanId.equals("TRACEID1-SPAN1")); + } else if (expectedTraceId.equals("TRACEID2")) { + assertThat(request.getResourceSpansList().size(), equalTo(1)); + ResourceSpans rs = request.getResourceSpansList().get(0); + assertThat(rs.getScopeSpansList().size(), equalTo(2)); + assertThat(rs.getInstrumentationLibrarySpansList().size(), equalTo(2)); + + ScopeSpans ss = rs.getScopeSpansList().get(0); + assertThat(ss.getSpansList().size(), equalTo(1)); + io.opentelemetry.proto.trace.v1.Span span = ss.getSpansList().get(0); + String spanId = span.getSpanId().toStringUtf8(); + assertTrue(spanId.equals("TRACEID2-SPAN1")); + + ss = rs.getScopeSpansList().get(1); + assertThat(ss.getSpansList().size(), equalTo(1)); + span = ss.getSpansList().get(0); + spanId = span.getSpanId().toStringUtf8(); + assertTrue(spanId.equals("TRACEID2-SPAN2")); + + } else if (expectedTraceId.equals("TRACEID3")) { + assertThat(request.getResourceSpansList().size(), equalTo(1)); + ResourceSpans rs = request.getResourceSpansList().get(0); + assertThat(rs.getScopeSpansList().size(), equalTo(1)); + assertThat(rs.getInstrumentationLibrarySpansList().size(), equalTo(0)); + ScopeSpans ss = rs.getScopeSpansList().get(0); + assertThat(ss.getSpansList().size(), equalTo(1)); + io.opentelemetry.proto.trace.v1.Span span = ss.getSpansList().get(0); + String spanId = span.getSpanId().toStringUtf8(); + assertTrue(spanId.equals("TRACEID3-SPAN1")); + } else { + assertTrue("Failed".equals("Unknown TraceId")); + } + } + } + @Test public void testParseExportTraceServiceRequest() throws IOException { final ExportTraceServiceRequest exportTraceServiceRequest = buildExportTraceServiceRequestFromJsonFile(TEST_REQUEST_TRACE_JSON_FILE); diff --git a/data-prepper-plugins/otel-proto-common/src/test/resources/test-request-multiple-traces.json b/data-prepper-plugins/otel-proto-common/src/test/resources/test-request-multiple-traces.json new file mode 100644 index 0000000000..461a5c935a --- /dev/null +++ b/data-prepper-plugins/otel-proto-common/src/test/resources/test-request-multiple-traces.json @@ -0,0 +1,249 @@ +{ + "resourceSpans": [ + { + "resource": { + "attributes": [ + { + "key": "service.name", + "value": { + "stringValue": "analytics-service1" + } + }, + { + "key": "telemetry.sdk.language", + "value": { + "stringValue": "java" + } + }, + { + "key": "telemetry.sdk.name", + "value": { + "stringValue": "opentelemetry" + } + }, + { + "key": "telemetry.sdk.version", + "value": { + "stringValue": "0.8.0-SNAPSHOT" + } + }, + { + "key": "array", + "value": { + "arrayValue": { + "values": [ + { + "stringValue": "test string" + }, + { + "boolValue": false + }, + { + "intValue": 0 + } + ] + } + } + }, + { + "key": "kvList", + "value": { + "kvlistValue": { + "values": [ + { + "key": "key1", + "value": { + "stringValue": "value1" + } + }, + { + "key": "key2", + "value": { + "stringValue": "value2" + } + } + ] + } + } + } + ], + "droppedAttributesCount": 0 + }, + "scopeSpans": [ + { + "scope": { + "name": "io.opentelemetry.auto.spring-webmvc-3.1", + "version": "" + }, + "spans": [ + { + "traceId": "VFJBQ0VJRDE=", + "spanId": "VFJBQ0VJRDEtU1BBTjE=", + "traceState": "", + "parentSpanId": "yxwHNNFJQP0=", + "name": "LoggingController.save", + "kind": "SPAN_KIND_INTERNAL", + "startTimeUnixNano": "1597902043168792500", + "endTimeUnixNano": "1597902043215953100", + "attributes": [], + "droppedAttributesCount": 0, + "events": [], + "droppedEventsCount": 0, + "links": [], + "droppedLinksCount": 0, + "status": { + "code": "STATUS_CODE_OK", + "message": "" + } + }, + { + "traceId": "VFJBQ0VJRDI=", + "spanId": "VFJBQ0VJRDItU1BBTjE=", + "traceState": "", + "parentSpanId": "yxwHNNFJQP0=", + "name": "LoggingController.save", + "kind": "SPAN_KIND_INTERNAL", + "startTimeUnixNano": "1597902043168792500", + "endTimeUnixNano": "1597902043215953100", + "attributes": [], + "droppedAttributesCount": 0, + "events": [], + "droppedEventsCount": 0, + "links": [], + "droppedLinksCount": 0, + "status": { + "code": "STATUS_CODE_OK", + "message": "" + } + } + ] + }, + { + "scope": { + "name": "io.opentelemetry.auto.apache-httpasyncclient-4.0", + "version": "" + }, + "spans": [ + { + "traceId": "VFJBQ0VJRDI=", + "spanId": "VFJBQ0VJRDItU1BBTjI=", + "traceState": "", + "parentSpanId": "XYZAgv/Pv40=", + "name": "HTTP PUT", + "kind": "SPAN_KIND_CLIENT", + "startTimeUnixNano": "1597902043175204700", + "endTimeUnixNano": "1597902043205117100", + "attributes": [ + { + "key": "http.status_code", + "value": { + "intValue": "200" + } + }, + { + "key": "http.url", + "value": { + "stringValue": "/logs/_doc/service_1?timeout\\u003d1m" + } + }, + { + "key": "http.method", + "value": { + "stringValue": "PUT" + } + } + ], + "droppedAttributesCount": 0, + "events": [], + "droppedEventsCount": 0, + "links": [], + "droppedLinksCount": 0, + "status": { + "code": "STATUS_CODE_OK", + "message": "" + } + }, + { + "traceId": "VFJBQ0VJRDM=", + "spanId": "VFJBQ0VJRDMtU1BBTjE=", + "traceState": "", + "parentSpanId": "yxwHNNFJQP0=", + "name": "LoggingController.save", + "kind": "SPAN_KIND_INTERNAL", + "startTimeUnixNano": "1597902043168792500", + "endTimeUnixNano": "1597902043215953100", + "attributes": [], + "droppedAttributesCount": 0, + "events": [], + "droppedEventsCount": 0, + "links": [], + "droppedLinksCount": 0, + "status": { + "code": "STATUS_CODE_OK", + "message": "" + } + } + ] + } + ], + "instrumentationLibrarySpans": [ + { + "instrumentationLibrary": { + "name": "io.opentelemetry.auto.spring-webmvc-3.1", + "version": "" + }, + "spans": [ + { + "traceId": "VFJBQ0VJRDI=", + "spanId": "VFJBQ0VJRDItU1BBTjE=", + "traceState": "", + "parentSpanId": "yxwHNNFJQP0=", + "name": "LoggingController.save", + "kind": "SPAN_KIND_INTERNAL", + "startTimeUnixNano": "1597902043168792500", + "endTimeUnixNano": "1597902043215953100", + "attributes": [], + "droppedAttributesCount": 0, + "events": [], + "droppedEventsCount": 0, + "links": [], + "droppedLinksCount": 0, + "status": { + "code": "STATUS_CODE_OK", + "message": "" + } + } + ] + }, + { + "instrumentationLibrary": { + "name": "io.opentelemetry.auto.spring-webmvc-3.1", + "version": "" + }, + "spans": [ + { + "traceId": "VFJBQ0VJRDI=", + "spanId": "VFJBQ0VJRDItU1BBTjI=", + "traceState": "", + "parentSpanId": "yxwHNNFJQP0=", + "name": "LoggingController.save", + "kind": "SPAN_KIND_INTERNAL", + "startTimeUnixNano": "1597902043168792500", + "endTimeUnixNano": "1597902043215953100", + "attributes": [], + "droppedAttributesCount": 0, + "events": [], + "droppedEventsCount": 0, + "links": [], + "droppedLinksCount": 0, + "status": { + "code": "STATUS_CODE_OK", + "message": "" + } + } + ] + } + ] + } + ] +} From 30d6c2c9b7cae4585a305950daba2885678ebfe7 Mon Sep 17 00:00:00 2001 From: David Venable Date: Fri, 10 Nov 2023 08:45:05 -0800 Subject: [PATCH 49/53] =?UTF-8?q?Disable=20the=20circuit=20breaker=20for?= =?UTF-8?q?=20buffers=20that=20write=20data=20off-heap=20only=E2=80=A6=20(?= =?UTF-8?q?#3619)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Disable the circuit breaker for buffers that write data off-heap only. This is currently only the Kafka buffer. Resolves #3616 Signed-off-by: David Venable --- .../dataprepper/model/buffer/Buffer.java | 11 +++ .../model/buffer/DelegatingBuffer.java | 5 ++ .../dataprepper/model/buffer/BufferTest.java | 38 ++++++--- .../model/buffer/DelegatingBufferTest.java | 9 ++ .../parser/MultiBufferDecorator.java | 49 +++++++++++ .../parser/PipelineTransformer.java | 24 +++--- .../plugins/MultiBufferDecorator.java | 83 ------------------- .../dataprepper/TestDataProvider.java | 9 +- .../MultiBufferDecoratorTest.java | 2 +- .../parser/PipelineTransformerTests.java | 22 +++++ .../plugins/TestOffHeapBuffer.java | 48 +++++++++++ .../single_pipeline_valid_off_heap_buffer.yml | 7 ++ .../plugins/kafka/buffer/KafkaBuffer.java | 5 ++ .../plugins/kafka/buffer/KafkaBufferTest.java | 6 ++ 14 files changed, 207 insertions(+), 111 deletions(-) create mode 100644 data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/MultiBufferDecorator.java delete mode 100644 data-prepper-core/src/main/java/org/opensearch/dataprepper/plugins/MultiBufferDecorator.java rename data-prepper-core/src/test/java/org/opensearch/dataprepper/{plugins => parser}/MultiBufferDecoratorTest.java (99%) create mode 100644 data-prepper-core/src/test/java/org/opensearch/dataprepper/plugins/TestOffHeapBuffer.java create mode 100644 data-prepper-core/src/test/resources/single_pipeline_valid_off_heap_buffer.yml diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/Buffer.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/Buffer.java index f8e3137285..b1ac56e8a8 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/Buffer.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/Buffer.java @@ -92,6 +92,17 @@ default Duration getDrainTimeout() { return Duration.ZERO; } + /** + * Indicates if writes to this buffer are also in some way written + * onto the JVM heap. If writes do go on heap, this should false + * which is the default. + * + * @return True if this buffer does not write to the JVM heap. + */ + default boolean isWrittenOffHeapOnly() { + return false; + } + /** * shuts down the buffer */ diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/DelegatingBuffer.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/DelegatingBuffer.java index adc6b76b9d..84f57883ea 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/DelegatingBuffer.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/buffer/DelegatingBuffer.java @@ -78,6 +78,11 @@ public Duration getDrainTimeout() { return delegateBuffer.getDrainTimeout(); } + @Override + public boolean isWrittenOffHeapOnly() { + return delegateBuffer.isWrittenOffHeapOnly(); + } + @Override public void shutdown() { delegateBuffer.shutdown(); diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/BufferTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/BufferTest.java index b7fc8505d4..505a00cc4a 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/BufferTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/BufferTest.java @@ -5,43 +5,55 @@ package org.opensearch.dataprepper.model.buffer; -import org.junit.Assert; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.model.event.Event; import org.opensearch.dataprepper.model.record.Record; import java.time.Duration; +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.Mockito.spy; -public class BufferTest { +class BufferTest { + + private Buffer createObjectUnderTest() { + return spy(Buffer.class); + } @Test - public void testGetDrainTimeout() { - final Buffer> buffer = spy(Buffer.class); + void testGetDrainTimeout() { + final Buffer> buffer = createObjectUnderTest(); - Assert.assertEquals(Duration.ZERO, buffer.getDrainTimeout()); + assertEquals(Duration.ZERO, buffer.getDrainTimeout()); } @Test - public void testShutdown() { - final Buffer> buffer = spy(Buffer.class); + void testShutdown() { + final Buffer> buffer = createObjectUnderTest(); buffer.shutdown(); } @Test - public void testIsByteBuffer() { - final Buffer> buffer = spy(Buffer.class); + void testIsByteBuffer() { + final Buffer> buffer = createObjectUnderTest(); - Assert.assertEquals(false, buffer.isByteBuffer()); + assertEquals(false, buffer.isByteBuffer()); + } + + @Test + void isWrittenOffHeapOnly_returns_false_by_default() { + assertThat(createObjectUnderTest().isWrittenOffHeapOnly(), equalTo(false)); } @Test - public void testWriteBytes() { - final Buffer> buffer = spy(Buffer.class); + void testWriteBytes() { + final Buffer> buffer = createObjectUnderTest(); byte[] bytes = new byte[2]; - Assert.assertThrows(UnsupportedOperationException.class, () -> buffer.writeBytes(bytes, "", 10)); + assertThrows(UnsupportedOperationException.class, () -> buffer.writeBytes(bytes, "", 10)); } diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/DelegatingBufferTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/DelegatingBufferTest.java index 6934104593..09e300c206 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/DelegatingBufferTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/buffer/DelegatingBufferTest.java @@ -193,6 +193,15 @@ void getDrainTimeout_returns_inner_getDrainTimeout() { equalTo(drainTimeout)); } + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void isWrittenOffHeapOnly_returns_inner_isWrittenOffHeapOnly(final boolean isWrittenOffHeapOnly) { + when(innerBuffer.isWrittenOffHeapOnly()).thenReturn(isWrittenOffHeapOnly); + + assertThat(createObjectUnderTest().isWrittenOffHeapOnly(), + equalTo(isWrittenOffHeapOnly)); + } + @Test void shutdown_calls_inner_shutdown() { createObjectUnderTest().shutdown(); diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/MultiBufferDecorator.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/MultiBufferDecorator.java new file mode 100644 index 0000000000..9fb7a7a201 --- /dev/null +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/MultiBufferDecorator.java @@ -0,0 +1,49 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.parser; + +import org.opensearch.dataprepper.model.buffer.Buffer; +import org.opensearch.dataprepper.model.buffer.DelegatingBuffer; +import org.opensearch.dataprepper.model.record.Record; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; + +/** + * Buffer decorator created for pipelines that make use of multiple buffers, such as PeerForwarder-enabled pipelines. The decorator + * acts as a pass-through to the primary buffer for most methods except those that rely on a combination of the primary + * and second. For example, isEmpty depends on all buffers being empty. + * + * @since 2.0 + */ +class MultiBufferDecorator> extends DelegatingBuffer implements Buffer { + private final List allBuffers; + + MultiBufferDecorator(final Buffer primaryBuffer, final List secondaryBuffers) { + super(primaryBuffer); + allBuffers = new ArrayList<>(1 + secondaryBuffers.size()); + allBuffers.add(primaryBuffer); + allBuffers.addAll(secondaryBuffers); + } + + @Override + public boolean isEmpty() { + return allBuffers.stream().allMatch(Buffer::isEmpty); + } + + @Override + public Duration getDrainTimeout() { + return allBuffers.stream() + .map(Buffer::getDrainTimeout) + .reduce(Duration.ZERO, Duration::plus); + } + + @Override + public void shutdown() { + allBuffers.forEach(Buffer::shutdown); + } +} diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/PipelineTransformer.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/PipelineTransformer.java index d520fe5062..64944aa47a 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/PipelineTransformer.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/parser/PipelineTransformer.java @@ -28,7 +28,6 @@ import org.opensearch.dataprepper.pipeline.PipelineConnector; import org.opensearch.dataprepper.pipeline.router.Router; import org.opensearch.dataprepper.pipeline.router.RouterFactory; -import org.opensearch.dataprepper.plugins.MultiBufferDecorator; import org.opensearch.dataprepper.sourcecoordination.SourceCoordinatorFactory; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -147,15 +146,7 @@ private void buildPipelineFromConfiguration( final MultiBufferDecorator multiBufferDecorator = new MultiBufferDecorator(pipelineDefinedBuffer, secondaryBuffers); - final Buffer buffer; - if(source instanceof PipelineConnector) { - buffer = multiBufferDecorator; - } else { - buffer = circuitBreakerManager.getGlobalCircuitBreaker() - .map(circuitBreaker -> new CircuitBreakingBuffer<>(multiBufferDecorator, circuitBreaker)) - .map(b -> (Buffer)b) - .orElseGet(() -> multiBufferDecorator); - } + final Buffer buffer = applyCircuitBreakerToBuffer(source, multiBufferDecorator); final Router router = routerFactory.createRouter(pipelineConfiguration.getRoutes()); @@ -313,4 +304,17 @@ List getSecondaryBuffers() { .map(innerEntry -> innerEntry.getValue()) .collect(Collectors.toList()); } + + private Buffer applyCircuitBreakerToBuffer(final Source source, final Buffer buffer) { + if (source instanceof PipelineConnector) + return buffer; + + if(buffer.isWrittenOffHeapOnly()) + return buffer; + + return circuitBreakerManager.getGlobalCircuitBreaker() + .map(circuitBreaker -> new CircuitBreakingBuffer<>(buffer, circuitBreaker)) + .map(b -> (Buffer) b) + .orElseGet(() -> buffer); + } } diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugins/MultiBufferDecorator.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugins/MultiBufferDecorator.java deleted file mode 100644 index cd199a7486..0000000000 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/plugins/MultiBufferDecorator.java +++ /dev/null @@ -1,83 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins; - -import org.opensearch.dataprepper.model.CheckpointState; -import org.opensearch.dataprepper.model.buffer.Buffer; -import org.opensearch.dataprepper.model.record.Record; - -import java.time.Duration; -import java.util.Collection; -import java.util.List; -import java.util.Map; -import java.util.concurrent.TimeoutException; -import java.util.stream.Stream; - -/** - * Buffer decorator created for pipelines that make use of multiple buffers, such as PeerForwarder-enabled pipelines. The decorator - * acts as a pass-through to the primary buffer for all methods except isEmpty, which verifies that all buffers are empty. - * - * @since 2.0 - */ -public class MultiBufferDecorator> implements Buffer { - private final Buffer primaryBuffer; - private final List secondaryBuffers; - - public MultiBufferDecorator(final Buffer primaryBuffer, final List secondaryBuffers) { - this.primaryBuffer = primaryBuffer; - this.secondaryBuffers = secondaryBuffers; - } - - @Override - public void write(final T record, final int timeoutInMillis) throws TimeoutException { - primaryBuffer.write(record, timeoutInMillis); - } - - @Override - public void writeAll(final Collection records, final int timeoutInMillis) throws Exception { - primaryBuffer.writeAll(records, timeoutInMillis); - } - - @Override - public void writeBytes(final byte[] bytes, final String key, int timeoutInMillis) throws Exception { - primaryBuffer.writeBytes(bytes, key, timeoutInMillis); - } - - @Override - public Map.Entry, CheckpointState> read(final int timeoutInMillis) { - return primaryBuffer.read(timeoutInMillis); - } - - @Override - public void checkpoint(final CheckpointState checkpointState) { - primaryBuffer.checkpoint(checkpointState); - } - - @Override - public boolean isByteBuffer() { - return primaryBuffer.isByteBuffer(); - } - - @Override - public boolean isEmpty() { - return primaryBuffer.isEmpty() && secondaryBuffers.stream() - .map(Buffer::isEmpty) - .allMatch(result -> result == true); - } - - @Override - public Duration getDrainTimeout() { - return Stream.concat(Stream.of(primaryBuffer), secondaryBuffers.stream()) - .map(Buffer::getDrainTimeout) - .reduce(Duration.ZERO, Duration::plus); - } - - @Override - public void shutdown() { - primaryBuffer.shutdown(); - secondaryBuffers.forEach(Buffer::shutdown); - } -} diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/TestDataProvider.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/TestDataProvider.java index 958a3331b3..c3051c7fc0 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/TestDataProvider.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/TestDataProvider.java @@ -5,14 +5,14 @@ package org.opensearch.dataprepper; -import org.opensearch.dataprepper.model.configuration.PluginModel; -import org.opensearch.dataprepper.model.configuration.PluginSetting; -import org.opensearch.dataprepper.model.configuration.SinkModel; -import org.opensearch.dataprepper.parser.model.PipelineConfiguration; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.DeserializationFeature; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; +import org.opensearch.dataprepper.model.configuration.PluginModel; +import org.opensearch.dataprepper.model.configuration.PluginSetting; +import org.opensearch.dataprepper.model.configuration.SinkModel; +import org.opensearch.dataprepper.parser.model.PipelineConfiguration; import java.io.File; import java.io.IOException; @@ -38,6 +38,7 @@ public class TestDataProvider { public static final String VALID_MULTIPLE_PIPELINE_CONFIG_FILE = "src/test/resources/valid_multiple_pipeline_configuration.yml"; public static final String VALID_PIPELINE_CONFIG_FILE_WITH_EXTENSIONS = "src/test/resources/valid_pipeline_configuration_with_extensions.yml"; public static final String VALID_SINGLE_PIPELINE_EMPTY_SOURCE_PLUGIN_FILE = "src/test/resources/single_pipeline_valid_empty_source_plugin_settings.yml"; + public static final String VALID_OFF_HEAP_FILE = "src/test/resources/single_pipeline_valid_off_heap_buffer.yml"; public static final String CONNECTED_PIPELINE_ROOT_SOURCE_INCORRECT = "src/test/resources/connected_pipeline_incorrect_root_source.yml"; public static final String CONNECTED_PIPELINE_CHILD_PIPELINE_INCORRECT = "src/test/resources/connected_pipeline_incorrect_child_pipeline.yml"; public static final String CYCLE_MULTIPLE_PIPELINE_CONFIG_FILE = "src/test/resources/cyclic_multiple_pipeline_configuration.yml"; diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugins/MultiBufferDecoratorTest.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/MultiBufferDecoratorTest.java similarity index 99% rename from data-prepper-core/src/test/java/org/opensearch/dataprepper/plugins/MultiBufferDecoratorTest.java rename to data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/MultiBufferDecoratorTest.java index 003e47c024..fa405154a1 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugins/MultiBufferDecoratorTest.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/MultiBufferDecoratorTest.java @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.dataprepper.plugins; +package org.opensearch.dataprepper.parser; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/PipelineTransformerTests.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/PipelineTransformerTests.java index 8182b4ab37..611dfaaf77 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/PipelineTransformerTests.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/parser/PipelineTransformerTests.java @@ -54,6 +54,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; @@ -303,6 +304,27 @@ void parseConfiguration_uses_CircuitBreaking_buffer_when_circuit_breakers_applie verify(dataPrepperConfiguration).getPipelineExtensions(); } + @Test + void parseConfiguration_uses_unwrapped_buffer_when_circuit_breakers_applied_but_Buffer_is_off_heap() { + final PipelineTransformer objectUnderTest = + createObjectUnderTest(TestDataProvider.VALID_OFF_HEAP_FILE); + + final Map pipelineMap = objectUnderTest.transformConfiguration(); + + assertThat(pipelineMap.size(), equalTo(1)); + assertThat(pipelineMap, hasKey("test-pipeline-1")); + final Pipeline pipeline = pipelineMap.get("test-pipeline-1"); + assertThat(pipeline, notNullValue()); + assertThat(pipeline.getBuffer(), notNullValue()); + assertThat(pipeline.getBuffer(), CoreMatchers.not(instanceOf(CircuitBreakingBuffer.class))); + + verifyNoInteractions(circuitBreakerManager); + verify(dataPrepperConfiguration).getProcessorShutdownTimeout(); + verify(dataPrepperConfiguration).getSinkShutdownTimeout(); + verify(dataPrepperConfiguration).getPeerForwarderConfiguration(); + verify(dataPrepperConfiguration).getPipelineExtensions(); + } + @Test void parseConfiguration_uses_unwrapped_buffer_when_no_circuit_breakers_are_applied() { when(circuitBreakerManager.getGlobalCircuitBreaker()) diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugins/TestOffHeapBuffer.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugins/TestOffHeapBuffer.java new file mode 100644 index 0000000000..5075c210ac --- /dev/null +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/plugins/TestOffHeapBuffer.java @@ -0,0 +1,48 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins; + +import org.opensearch.dataprepper.model.CheckpointState; +import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; +import org.opensearch.dataprepper.model.buffer.Buffer; +import org.opensearch.dataprepper.model.record.Record; + +import java.util.Collection; +import java.util.Map; +import java.util.concurrent.TimeoutException; + +@DataPrepperPlugin(name = "test_off_heap", pluginType = Buffer.class) +public class TestOffHeapBuffer implements Buffer { + @Override + public void write(Record record, int timeoutInMillis) throws TimeoutException { + + } + + @Override + public void writeAll(Collection records, int timeoutInMillis) throws Exception { + + } + + @Override + public Map.Entry read(int timeoutInMillis) { + return null; + } + + @Override + public void checkpoint(CheckpointState checkpointState) { + + } + + @Override + public boolean isWrittenOffHeapOnly() { + return true; + } + + @Override + public boolean isEmpty() { + return false; + } +} diff --git a/data-prepper-core/src/test/resources/single_pipeline_valid_off_heap_buffer.yml b/data-prepper-core/src/test/resources/single_pipeline_valid_off_heap_buffer.yml new file mode 100644 index 0000000000..e7ed4c35e1 --- /dev/null +++ b/data-prepper-core/src/test/resources/single_pipeline_valid_off_heap_buffer.yml @@ -0,0 +1,7 @@ +test-pipeline-1: + source: + stdin: + buffer: + test_off_heap: + sink: + - stdout: 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 296571bf1c..fd8f7365da 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 @@ -141,6 +141,11 @@ public Duration getDrainTimeout() { return drainTimeout; } + @Override + public boolean isWrittenOffHeapOnly() { + return true; + } + @Override public void shutdown() { shutdownInProgress.set(true); 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 70665dcd72..9a984115d1 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 @@ -356,6 +356,12 @@ void test_kafkaBuffer_getDrainTimeout() { verify(bufferConfig).getDrainTimeout(); } + @Test + void isWrittenOffHeapOnly_returns_true() { + assertThat(createObjectUnderTest().isWrittenOffHeapOnly(), + equalTo(true)); + } + @Test public void testShutdown_Successful() throws InterruptedException { kafkaBuffer = createObjectUnderTest(); From 9f0c92afd885ff2467f7d19e98c8b7e8da82a429 Mon Sep 17 00:00:00 2001 From: David Venable Date: Fri, 10 Nov 2023 09:57:00 -0800 Subject: [PATCH 50/53] Updated the Router_ThreeRoutesIT test with a longer timeout. Also, use Awaitility's during() method to verify that certain data never reaches a sink that it never should reach. (#3624) Signed-off-by: David Venable --- .../dataprepper/integration/Router_ThreeRoutesIT.java | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/data-prepper-core/src/integrationTest/java/org/opensearch/dataprepper/integration/Router_ThreeRoutesIT.java b/data-prepper-core/src/integrationTest/java/org/opensearch/dataprepper/integration/Router_ThreeRoutesIT.java index 18a711a8e6..e8c1c08563 100644 --- a/data-prepper-core/src/integrationTest/java/org/opensearch/dataprepper/integration/Router_ThreeRoutesIT.java +++ b/data-prepper-core/src/integrationTest/java/org/opensearch/dataprepper/integration/Router_ThreeRoutesIT.java @@ -79,7 +79,7 @@ void sending_alpha_and_beta_events_sends_to_sinks_that_take_both(final String so inMemorySourceAccessor.submit(TESTING_KEY, allEvents); - await().atMost(400, TimeUnit.MILLISECONDS) + await().atMost(2, TimeUnit.SECONDS) .untilAsserted(() -> { assertThat(inMemorySinkAccessor.get(sourceKeyToReceiveAll), not(empty())); }); @@ -101,7 +101,7 @@ void sending_alpha_and_beta_events_sends_to_both_sinks() { inMemorySourceAccessor.submit(TESTING_KEY, allEvents); - await().atMost(400, TimeUnit.MILLISECONDS) + await().atMost(2, TimeUnit.SECONDS) .untilAsserted(() -> { assertThat(inMemorySinkAccessor.get(ALPHA_SOURCE_KEY), not(empty())); assertThat(inMemorySinkAccessor.get(BETA_SOURCE_KEY), not(empty())); @@ -136,9 +136,9 @@ void sending_non_alpha_beta_events_never_reaches_sink(final String sourceKey) th inMemorySourceAccessor.submit(TESTING_KEY, randomEvents); - Thread.sleep(1000); - - assertThat(inMemorySinkAccessor.get(sourceKey), empty()); + await().during(1200, TimeUnit.MILLISECONDS) + .pollDelay(50, TimeUnit.MILLISECONDS) + .until(() -> inMemorySinkAccessor.get(sourceKey), empty()); } private List> createEvents(final String value, final int numberToCreate) { From 95dd099b3762699c8f1e75106c235b4d24ac8516 Mon Sep 17 00:00:00 2001 From: Taylor Gray Date: Fri, 10 Nov 2023 12:56:00 -0600 Subject: [PATCH 51/53] =?UTF-8?q?Add=20metric=20for=20shards=20actively=20?= =?UTF-8?q?being=20processed,=20lower=20ownership=20timeo=E2=80=A6=20(#362?= =?UTF-8?q?9)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add metric for shards actively being processed, lower ownership timeout from 10 minutes to 5 minutes for ddb source Signed-off-by: Taylor Gray --- .../source/dynamodb/export/DataFileCheckpointer.java | 4 +++- .../source/dynamodb/stream/StreamCheckpointer.java | 4 +++- .../source/dynamodb/stream/StreamScheduler.java | 10 +++++++++- .../source/dynamodb/stream/ShardConsumerTest.java | 5 +++-- .../source/dynamodb/stream/StreamSchedulerTest.java | 11 +++++++++++ 5 files changed, 29 insertions(+), 5 deletions(-) diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileCheckpointer.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileCheckpointer.java index f5fdc882f3..42f5463d52 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileCheckpointer.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileCheckpointer.java @@ -21,6 +21,8 @@ public class DataFileCheckpointer { private static final Logger LOG = LoggerFactory.getLogger(DataFileCheckpointer.class); + static final Duration CHECKPOINT_OWNERSHIP_TIMEOUT_INCREASE = Duration.ofMinutes(5); + private final EnhancedSourceCoordinator enhancedSourceCoordinator; @@ -48,7 +50,7 @@ private void setProgressState(int lineNumber) { public void checkpoint(int lineNumber) { LOG.debug("Checkpoint data file " + dataFilePartition.getKey() + " with line number " + lineNumber); setProgressState(lineNumber); - enhancedSourceCoordinator.saveProgressStateForPartition(dataFilePartition, null); + enhancedSourceCoordinator.saveProgressStateForPartition(dataFilePartition, CHECKPOINT_OWNERSHIP_TIMEOUT_INCREASE); } /** diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamCheckpointer.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamCheckpointer.java index 9e3113732a..85c7c9c69c 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamCheckpointer.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamCheckpointer.java @@ -22,6 +22,8 @@ public class StreamCheckpointer { private static final Logger LOG = LoggerFactory.getLogger(StreamCheckpointer.class); + static final Duration CHECKPOINT_OWNERSHIP_TIMEOUT_INCREASE = Duration.ofMinutes(5); + private final EnhancedSourceCoordinator coordinator; private final StreamPartition streamPartition; @@ -52,7 +54,7 @@ private void setSequenceNumber(String sequenceNumber) { public void checkpoint(String sequenceNumber) { LOG.debug("Checkpoint shard " + streamPartition.getShardId() + " with sequenceNumber " + sequenceNumber); setSequenceNumber(sequenceNumber); - coordinator.saveProgressStateForPartition(streamPartition, null); + coordinator.saveProgressStateForPartition(streamPartition, CHECKPOINT_OWNERSHIP_TIMEOUT_INCREASE); } /** diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamScheduler.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamScheduler.java index ea799bfc24..660d1b82bf 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamScheduler.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamScheduler.java @@ -47,6 +47,7 @@ public class StreamScheduler implements Runnable { private static final int DELAY_TO_GET_CHILD_SHARDS_MILLIS = 1_500; static final String ACTIVE_CHANGE_EVENT_CONSUMERS = "activeChangeEventConsumers"; + static final String SHARDS_IN_PROCESSING = "activeShardsInProcessing"; private final AtomicInteger numOfWorkers = new AtomicInteger(0); private final EnhancedSourceCoordinator coordinator; @@ -55,6 +56,7 @@ public class StreamScheduler implements Runnable { private final ShardManager shardManager; private final PluginMetrics pluginMetrics; private final AtomicLong activeChangeEventConsumers; + private final AtomicLong shardsInProcessing; private final AcknowledgementSetManager acknowledgementSetManager; private final DynamoDBSourceConfig dynamoDBSourceConfig; @@ -74,6 +76,7 @@ public StreamScheduler(final EnhancedSourceCoordinator coordinator, executor = Executors.newFixedThreadPool(MAX_JOB_COUNT); activeChangeEventConsumers = pluginMetrics.gauge(ACTIVE_CHANGE_EVENT_CONSUMERS, new AtomicLong()); + shardsInProcessing = pluginMetrics.gauge(SHARDS_IN_PROCESSING, new AtomicLong()); } private void processStreamPartition(StreamPartition streamPartition) { @@ -98,11 +101,15 @@ private void processStreamPartition(StreamPartition streamPartition) { CompletableFuture runConsumer = CompletableFuture.runAsync(shardConsumer, executor); if (acknowledgmentsEnabled) { - runConsumer.whenComplete((v, ex) -> numOfWorkers.decrementAndGet()); + runConsumer.whenComplete((v, ex) -> { + numOfWorkers.decrementAndGet(); + shardsInProcessing.decrementAndGet(); + }); } else { runConsumer.whenComplete(completeConsumer(streamPartition)); } numOfWorkers.incrementAndGet(); + shardsInProcessing.incrementAndGet(); } else { // If failed to create a new consumer. coordinator.completePartition(streamPartition); @@ -153,6 +160,7 @@ private BiConsumer completeConsumer(StreamPartition streamPartition) { return (v, ex) -> { if (!dynamoDBSourceConfig.isAcknowledgmentsEnabled()) { numOfWorkers.decrementAndGet(); + shardsInProcessing.decrementAndGet(); } if (ex == null) { LOG.info("Shard consumer for {} is completed", streamPartition.getShardId()); diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerTest.java index ec9df68189..4da0969933 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerTest.java @@ -54,6 +54,7 @@ import static org.mockito.Mockito.when; import static org.opensearch.dataprepper.plugins.source.dynamodb.stream.ShardConsumer.BUFFER_TIMEOUT; import static org.opensearch.dataprepper.plugins.source.dynamodb.stream.ShardConsumer.DEFAULT_BUFFER_BATCH_SIZE; +import static org.opensearch.dataprepper.plugins.source.dynamodb.stream.StreamCheckpointer.CHECKPOINT_OWNERSHIP_TIMEOUT_INCREASE; @ExtendWith(MockitoExtension.class) class ShardConsumerTest { @@ -169,7 +170,7 @@ void test_run_shardConsumer_correctly() throws Exception { verify(bufferAccumulator, times(total)).add(any(org.opensearch.dataprepper.model.record.Record.class)); verify(bufferAccumulator).flush(); // Should complete the consumer as reach to end of shard - verify(coordinator).saveProgressStateForPartition(any(StreamPartition.class), eq(null)); + verify(coordinator).saveProgressStateForPartition(any(StreamPartition.class), eq(CHECKPOINT_OWNERSHIP_TIMEOUT_INCREASE)); } @Test @@ -203,7 +204,7 @@ void test_run_shardConsumer_with_acknowledgments_correctly() throws Exception { verify(bufferAccumulator).flush(); // Should complete the consumer as reach to end of shard - verify(coordinator).saveProgressStateForPartition(any(StreamPartition.class), eq(null)); + verify(coordinator).saveProgressStateForPartition(any(StreamPartition.class), eq(CHECKPOINT_OWNERSHIP_TIMEOUT_INCREASE)); verify(acknowledgementSet).complete(); } diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamSchedulerTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamSchedulerTest.java index 4c6dfdcd0d..149ca2eca0 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamSchedulerTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/StreamSchedulerTest.java @@ -44,6 +44,7 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.opensearch.dataprepper.plugins.source.dynamodb.stream.StreamScheduler.ACTIVE_CHANGE_EVENT_CONSUMERS; +import static org.opensearch.dataprepper.plugins.source.dynamodb.stream.StreamScheduler.SHARDS_IN_PROCESSING; @ExtendWith(MockitoExtension.class) class StreamSchedulerTest { @@ -79,6 +80,9 @@ class StreamSchedulerTest { @Mock private AtomicLong activeShardConsumers; + @Mock + private AtomicLong activeShardsInProcessing; + private final String tableName = UUID.randomUUID().toString(); private final String tableArn = "arn:aws:dynamodb:us-west-2:123456789012:table/" + tableName; @@ -108,6 +112,7 @@ void setup() { lenient().when(shardManager.getChildShardIds(anyString(), anyString())).thenReturn(List.of(shardId)); when(pluginMetrics.gauge(eq(ACTIVE_CHANGE_EVENT_CONSUMERS), any(AtomicLong.class))).thenReturn(activeShardConsumers); + when(pluginMetrics.gauge(eq(SHARDS_IN_PROCESSING), any(AtomicLong.class))).thenReturn(activeShardsInProcessing); } @@ -135,6 +140,9 @@ public void test_normal_run() throws InterruptedException { // Should mask the stream partition as completed. verify(coordinator).completePartition(any(StreamPartition.class)); + verify(activeShardsInProcessing).incrementAndGet(); + verify(activeShardsInProcessing).decrementAndGet(); + executorService.shutdownNow(); } @@ -174,6 +182,9 @@ public void test_normal_run_with_acknowledgments() throws InterruptedException { // Should mask the stream partition as completed. verify(coordinator).completePartition(any(StreamPartition.class)); + verify(activeShardsInProcessing).incrementAndGet(); + verify(activeShardsInProcessing).decrementAndGet(); + executorService.shutdownNow(); } From 7df58f4f61cc90e46a3621312e532f9438f3a0d7 Mon Sep 17 00:00:00 2001 From: David Venable Date: Fri, 10 Nov 2023 11:15:41 -0800 Subject: [PATCH 52/53] Attempting to fix the flaky tests in AwsCloudMapPeerListProviderTest. This uses a higher wait for changes, and refreshes at sub-second intervals for testing. (#3628) Signed-off-by: David Venable --- .../AwsCloudMapPeerListProvider.java | 17 +++++++++-------- .../AwsCloudMapPeerListProviderTest.java | 19 ++++++++++--------- 2 files changed, 19 insertions(+), 17 deletions(-) diff --git a/data-prepper-core/src/main/java/org/opensearch/dataprepper/peerforwarder/discovery/AwsCloudMapPeerListProvider.java b/data-prepper-core/src/main/java/org/opensearch/dataprepper/peerforwarder/discovery/AwsCloudMapPeerListProvider.java index 5ddca5f10e..1b86adebfe 100644 --- a/data-prepper-core/src/main/java/org/opensearch/dataprepper/peerforwarder/discovery/AwsCloudMapPeerListProvider.java +++ b/data-prepper-core/src/main/java/org/opensearch/dataprepper/peerforwarder/discovery/AwsCloudMapPeerListProvider.java @@ -23,6 +23,7 @@ import software.amazon.awssdk.services.servicediscovery.model.DiscoverInstancesResponse; import software.amazon.awssdk.services.servicediscovery.model.HttpInstanceSummary; +import java.time.Duration; import java.util.List; import java.util.Map; import java.util.Objects; @@ -48,7 +49,7 @@ class AwsCloudMapPeerListProvider implements PeerListProvider, AutoCloseable { private final String serviceName; private final Map queryParameters; private final AwsCloudMapDynamicEndpointGroup endpointGroup; - private final int timeToRefreshSeconds; + private final Duration timeToRefresh; private final Backoff backoff; private final EventLoop eventLoop; private final String domainName; @@ -58,18 +59,18 @@ class AwsCloudMapPeerListProvider implements PeerListProvider, AutoCloseable { final String namespaceName, final String serviceName, final Map queryParameters, - final int timeToRefreshSeconds, + final Duration timeToRefresh, final Backoff backoff, final PluginMetrics pluginMetrics) { this.awsServiceDiscovery = Objects.requireNonNull(awsServiceDiscovery); this.namespaceName = Objects.requireNonNull(namespaceName); this.serviceName = Objects.requireNonNull(serviceName); this.queryParameters = Objects.requireNonNull(queryParameters); - this.timeToRefreshSeconds = timeToRefreshSeconds; + this.timeToRefresh = timeToRefresh; this.backoff = Objects.requireNonNull(backoff); - if (timeToRefreshSeconds < 1) - throw new IllegalArgumentException("timeToRefreshSeconds must be positive. Actual: " + timeToRefreshSeconds); + if (timeToRefresh.isNegative() || timeToRefresh.isZero()) + throw new IllegalArgumentException("timeToRefreshSeconds must be positive. Actual: " + timeToRefresh); eventLoop = CommonPools.workerGroup().next(); LOG.info("Using AWS CloudMap for Peer Forwarding. namespace='{}', serviceName='{}'", @@ -92,7 +93,7 @@ static AwsCloudMapPeerListProvider createPeerListProvider(final PeerForwarderCon final Map queryParameters = peerForwarderConfiguration.getAwsCloudMapQueryParameters(); final Backoff standardBackoff = Backoff.exponential(ONE_SECOND, TWENTY_SECONDS).withJitter(TWENTY_PERCENT); - final int timeToRefreshSeconds = 20; + final Duration timeToRefresh = Duration.ofSeconds(20); final PluginMetrics awsSdkMetrics = PluginMetrics.fromNames("sdk", "aws"); @@ -108,7 +109,7 @@ static AwsCloudMapPeerListProvider createPeerListProvider(final PeerForwarderCon namespace, serviceName, queryParameters, - timeToRefreshSeconds, + timeToRefresh, standardBackoff, pluginMetrics); } @@ -175,7 +176,7 @@ private void discoverInstances() { LOG.warn("Failed to update endpoints.", ex); } finally { scheduledDiscovery = eventLoop.schedule(this::discoverInstances, - timeToRefreshSeconds, TimeUnit.SECONDS); + timeToRefresh.toMillis(), TimeUnit.MILLISECONDS); } } diff --git a/data-prepper-core/src/test/java/org/opensearch/dataprepper/peerforwarder/discovery/AwsCloudMapPeerListProviderTest.java b/data-prepper-core/src/test/java/org/opensearch/dataprepper/peerforwarder/discovery/AwsCloudMapPeerListProviderTest.java index 0c7c23093a..67429e27ef 100644 --- a/data-prepper-core/src/test/java/org/opensearch/dataprepper/peerforwarder/discovery/AwsCloudMapPeerListProviderTest.java +++ b/data-prepper-core/src/test/java/org/opensearch/dataprepper/peerforwarder/discovery/AwsCloudMapPeerListProviderTest.java @@ -5,11 +5,9 @@ package org.opensearch.dataprepper.peerforwarder.discovery; -import org.opensearch.dataprepper.metrics.PluginMetrics; import com.linecorp.armeria.client.Endpoint; import com.linecorp.armeria.client.retry.Backoff; import org.apache.commons.lang3.RandomStringUtils; -import org.awaitility.Awaitility; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Nested; @@ -19,11 +17,13 @@ import org.mockito.ArgumentCaptor; import org.mockito.ArgumentMatchers; import org.mockito.InOrder; +import org.opensearch.dataprepper.metrics.PluginMetrics; import software.amazon.awssdk.services.servicediscovery.ServiceDiscoveryAsyncClient; import software.amazon.awssdk.services.servicediscovery.model.DiscoverInstancesRequest; import software.amazon.awssdk.services.servicediscovery.model.DiscoverInstancesResponse; import software.amazon.awssdk.services.servicediscovery.model.HttpInstanceSummary; +import java.time.Duration; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; @@ -38,6 +38,7 @@ import java.util.stream.Collectors; import java.util.stream.IntStream; +import static org.awaitility.Awaitility.await; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.notNullValue; import static org.hamcrest.CoreMatchers.nullValue; @@ -59,7 +60,7 @@ class AwsCloudMapPeerListProviderTest { private String namespaceName; private String serviceName; private Map queryParameters; - private int timeToRefreshSeconds; + private Duration timeToRefresh; private Backoff backoff; private PluginMetrics pluginMetrics; private List objectsToClose; @@ -71,7 +72,7 @@ void setUp() { serviceName = RandomStringUtils.randomAlphabetic(10); queryParameters = generateRandomStringMap(); - timeToRefreshSeconds = 1; + timeToRefresh = Duration.ofMillis(200); backoff = mock(Backoff.class); pluginMetrics = mock(PluginMetrics.class); @@ -85,7 +86,7 @@ void tearDown() { private AwsCloudMapPeerListProvider createObjectUnderTest() { final AwsCloudMapPeerListProvider objectUnderTest = - new AwsCloudMapPeerListProvider(awsServiceDiscovery, namespaceName, serviceName, queryParameters, timeToRefreshSeconds, backoff, pluginMetrics); + new AwsCloudMapPeerListProvider(awsServiceDiscovery, namespaceName, serviceName, queryParameters, timeToRefresh, backoff, pluginMetrics); objectsToClose.add(objectUnderTest); return objectUnderTest; } @@ -133,7 +134,7 @@ void constructor_throws_with_null_Backoff() { @ParameterizedTest @ValueSource(ints = {Integer.MIN_VALUE, -10, -1, 0}) void constructor_throws_with_non_positive_timeToRefreshSeconds(final int badTimeToRefresh) { - timeToRefreshSeconds = badTimeToRefresh; + timeToRefresh = Duration.ofSeconds(badTimeToRefresh); assertThrows(IllegalArgumentException.class, this::createObjectUnderTest); @@ -374,7 +375,7 @@ private void waitUntilDiscoverInstancesCalledAtLeastOnce() { */ private void waitUntilDiscoverInstancesCalledAtLeast(final int timesCalled) { final long waitTimeMillis = (long) timesCalled * WAIT_TIME_MULTIPLIER_MILLIS; - Awaitility.waitAtMost(waitTimeMillis, TimeUnit.MILLISECONDS) + await().atMost(waitTimeMillis, TimeUnit.MILLISECONDS) .pollDelay(100, TimeUnit.MILLISECONDS) .untilAsserted(() -> then(awsServiceDiscovery) .should(atLeast(timesCalled)) @@ -388,7 +389,7 @@ private void waitUntilDiscoverInstancesCalledAtLeast(final int timesCalled) { * @param objectUnderTest The object to wait for. */ private void waitUntilPeerListPopulated(final AwsCloudMapPeerListProvider objectUnderTest) { - Awaitility.waitAtMost(2, TimeUnit.SECONDS) + await().atMost(5, TimeUnit.SECONDS) .pollDelay(100, TimeUnit.MILLISECONDS) .untilAsserted(() -> { final List actualPeers = objectUnderTest.getPeerList(); @@ -411,7 +412,7 @@ private static Map generateRandomStringMap() { final Map map = new HashMap<>(); IntStream.range(0, random.nextInt(5) + 1) - .mapToObj(num -> map.put(UUID.randomUUID().toString(), UUID.randomUUID().toString())); + .forEach(num -> map.put(UUID.randomUUID().toString(), UUID.randomUUID().toString())); return map; } From d2ad8aed36f30006547fc3e441682db4f82665b7 Mon Sep 17 00:00:00 2001 From: Taylor Gray Date: Fri, 10 Nov 2023 13:27:29 -0600 Subject: [PATCH 53/53] Start unit tests for the OpenSearch sink testing the document_version error cases (#3599) Signed-off-by: Taylor Gray --- .../expression/ExpressionEvaluator.java | 2 +- .../expression/ExpressionEvaluatorTest.java | 2 +- .../GenericExpressionEvaluator.java | 2 +- .../GenericExpressionEvaluatorTest.java | 4 +- .../sink/opensearch/OpenSearchSinkIT.java | 4 +- .../sink/opensearch/OpenSearchSink.java | 13 +- .../opensearch/index/IndexConfiguration.java | 6 +- .../OpenSearchSinkConfigurationTests.java | 2 +- .../sink/opensearch/OpenSearchSinkTest.java | 296 ++++++++++++++++++ .../index/IndexConfigurationTests.java | 4 +- 10 files changed, 318 insertions(+), 17 deletions(-) create mode 100644 data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkTest.java diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/expression/ExpressionEvaluator.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/expression/ExpressionEvaluator.java index ea786ce5bd..7dc5930816 100644 --- a/data-prepper-api/src/main/java/org/opensearch/dataprepper/expression/ExpressionEvaluator.java +++ b/data-prepper-api/src/main/java/org/opensearch/dataprepper/expression/ExpressionEvaluator.java @@ -35,5 +35,5 @@ default Boolean evaluateConditional(final String statement, final Event context) Boolean isValidExpressionStatement(final String statement); - Boolean isValidFormatExpressions(final String format); + Boolean isValidFormatExpression(final String format); } \ No newline at end of file diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/expression/ExpressionEvaluatorTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/expression/ExpressionEvaluatorTest.java index 6ce1af8660..9b76fbc807 100644 --- a/data-prepper-api/src/test/java/org/opensearch/dataprepper/expression/ExpressionEvaluatorTest.java +++ b/data-prepper-api/src/test/java/org/opensearch/dataprepper/expression/ExpressionEvaluatorTest.java @@ -25,7 +25,7 @@ public Boolean isValidExpressionStatement(final String statement) { } @Override - public Boolean isValidFormatExpressions(String format) { + public Boolean isValidFormatExpression(String format) { return true; } } diff --git a/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluator.java b/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluator.java index 89d8b15a84..6653e5c7b4 100644 --- a/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluator.java +++ b/data-prepper-expression/src/main/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluator.java @@ -55,7 +55,7 @@ public Boolean isValidExpressionStatement(final String statement) { } @Override - public Boolean isValidFormatExpressions(final String format) { + public Boolean isValidFormatExpression(final String format) { int fromIndex = 0; int position = 0; while ((position = format.indexOf("${", fromIndex)) != -1) { diff --git a/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluatorTest.java b/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluatorTest.java index 0b8ac2ee03..a91e7fe368 100644 --- a/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluatorTest.java +++ b/data-prepper-expression/src/test/java/org/opensearch/dataprepper/expression/GenericExpressionEvaluatorTest.java @@ -132,14 +132,14 @@ void isValidExpressionStatement_returns_false_when_parse_throws() { "abc-${invalid, false" }) void isValidFormatExpressionsReturnsCorrectResult(final String format, final Boolean expectedResult) { - assertThat(statementEvaluator.isValidFormatExpressions(format), equalTo(expectedResult)); + assertThat(statementEvaluator.isValidFormatExpression(format), equalTo(expectedResult)); } @ParameterizedTest @ValueSource(strings = {"abc-${anyS(=tring}"}) void isValidFormatExpressionsReturnsFalseWhenIsValidKeyAndValidExpressionIsFalse(final String format) { doThrow(RuntimeException.class).when(parser).parse(anyString()); - assertThat(statementEvaluator.isValidFormatExpressions(format), equalTo(false)); + assertThat(statementEvaluator.isValidFormatExpression(format), equalTo(false)); } } 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 97c80fe3bc..a5237f21de 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 @@ -682,7 +682,7 @@ public void testOpenSearchBulkActionsCreateWithExpression() throws IOException, Event event = (Event) testRecords.get(0).getData(); event.getMetadata().setAttribute("action", "create"); final String actionFormatExpression = "${getMetadata(\"action\")}"; - when(expressionEvaluator.isValidFormatExpressions(actionFormatExpression)).thenReturn(true); + when(expressionEvaluator.isValidFormatExpression(actionFormatExpression)).thenReturn(true); when(expressionEvaluator.isValidExpressionStatement("getMetadata(\"action\")")).thenReturn(true); when(expressionEvaluator.evaluate("getMetadata(\"action\")", event)).thenReturn(event.getMetadata().getAttribute("action")); pluginSetting.getSettings().put(IndexConfiguration.ACTION, actionFormatExpression); @@ -715,7 +715,7 @@ public void testOpenSearchBulkActionsCreateWithInvalidExpression() throws IOExce Event event = (Event) testRecords.get(0).getData(); event.getMetadata().setAttribute("action", "unknown"); final String actionFormatExpression = "${getMetadata(\"action\")}"; - when(expressionEvaluator.isValidFormatExpressions(actionFormatExpression)).thenReturn(true); + when(expressionEvaluator.isValidFormatExpression(actionFormatExpression)).thenReturn(true); when(expressionEvaluator.isValidExpressionStatement("getMetadata(\"action\")")).thenReturn(true); when(expressionEvaluator.evaluate("getMetadata(\"action\")", event)).thenReturn(event.getMetadata().getAttribute("action")); pluginSetting.getSettings().put(IndexConfiguration.ACTION, actionFormatExpression); 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 b1edff4802..19676a6a2a 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 @@ -371,12 +371,17 @@ public void doOutput(final Collection> records) { versionExpressionEvaluationResult = event.formatString(versionExpression, expressionEvaluator); version = Long.valueOf(event.formatString(versionExpression, expressionEvaluator)); } catch (final NumberFormatException e) { - LOG.warn("Unable to convert the result of evaluating document_version '{}' to Long for an Event. The evaluation result '{}' must be a valid Long type", versionExpression, versionExpressionEvaluationResult); - logFailureForDlqObjects(List.of(createDlqObjectFromEvent(event, indexName, e.getMessage())), e); + final String errorMessage = String.format( + "Unable to convert the result of evaluating document_version '%s' to Long for an Event. The evaluation result '%s' must be a valid Long type", versionExpression, versionExpressionEvaluationResult + ); + LOG.error(errorMessage); + logFailureForDlqObjects(List.of(createDlqObjectFromEvent(event, indexName, errorMessage)), e); dynamicDocumentVersionDroppedEvents.increment(); } catch (final RuntimeException e) { - LOG.error("There was an exception when evaluating the document_version '{}'. Check the dlq if configured to see details about the affected Event: {}", versionExpression, e.getMessage()); - logFailureForDlqObjects(List.of(createDlqObjectFromEvent(event, indexName, e.getMessage())), e); + final String errorMessage = String.format( + "There was an exception when evaluating the document_version '%s': %s", versionExpression, e.getMessage()); + LOG.error(errorMessage + " Check the dlq if configured to see more details about the affected Event"); + logFailureForDlqObjects(List.of(createDlqObjectFromEvent(event, indexName, errorMessage)), e); dynamicDocumentVersionDroppedEvents.increment(); } } diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/IndexConfiguration.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/IndexConfiguration.java index c6ad893578..363e074c65 100644 --- a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/IndexConfiguration.java +++ b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/IndexConfiguration.java @@ -232,7 +232,7 @@ public static IndexConfiguration readIndexConfig(final PluginSetting pluginSetti final String versionType = pluginSetting.getStringOrDefault(DOCUMENT_VERSION_TYPE, null); builder = builder.withVersionExpression(versionExpression); - if (versionExpression != null && (!expressionEvaluator.isValidFormatExpressions(versionExpression))) { + if (versionExpression != null && (!expressionEvaluator.isValidFormatExpression(versionExpression))) { throw new InvalidPluginConfigurationException("document_version {} is not a valid format expression."); } @@ -546,7 +546,7 @@ public Builder withIsmPolicyFile(final String ismPolicyFile) { public Builder withAction(final String action, final ExpressionEvaluator expressionEvaluator) { checkArgument((EnumUtils.isValidEnumIgnoreCase(OpenSearchBulkActions.class, action) || - (action.contains("${") && expressionEvaluator.isValidFormatExpressions(action))), "action \"" + action + "\" is invalid. action must be one of the following: " + Arrays.stream(OpenSearchBulkActions.values()).collect(Collectors.toList())); + (action.contains("${") && expressionEvaluator.isValidFormatExpression(action))), "action \"" + action + "\" is invalid. action must be one of the following: " + Arrays.stream(OpenSearchBulkActions.values()).collect(Collectors.toList())); this.action = action; return this; } @@ -556,7 +556,7 @@ public Builder withActions(final List> actions, final Expres String action = (String)actionMap.get("type"); if (action != null) { checkArgument((EnumUtils.isValidEnumIgnoreCase(OpenSearchBulkActions.class, action) || - (action.contains("${") && expressionEvaluator.isValidFormatExpressions(action))), "action \"" + action + "\". action must be one of the following: " + Arrays.stream(OpenSearchBulkActions.values()).collect(Collectors.toList())); + (action.contains("${") && expressionEvaluator.isValidFormatExpression(action))), "action \"" + action + "\". action must be one of the following: " + Arrays.stream(OpenSearchBulkActions.values()).collect(Collectors.toList())); } } this.actions = actions; diff --git a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkConfigurationTests.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkConfigurationTests.java index 78f3efc516..cb9dfbe898 100644 --- a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkConfigurationTests.java +++ b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkConfigurationTests.java @@ -142,7 +142,7 @@ public void testReadESConfigWithBulkActionCreateExpression() { pluginSetting.setPipelineName(PIPELINE_NAME); expressionEvaluator = mock(ExpressionEvaluator.class); - when(expressionEvaluator.isValidFormatExpressions(actionFormatExpression)).thenReturn(true); + when(expressionEvaluator.isValidFormatExpression(actionFormatExpression)).thenReturn(true); final OpenSearchSinkConfiguration openSearchSinkConfiguration = OpenSearchSinkConfiguration.readESConfig(pluginSetting, expressionEvaluator); diff --git a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkTest.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkTest.java new file mode 100644 index 0000000000..237f4b9a2f --- /dev/null +++ b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkTest.java @@ -0,0 +1,296 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.opensearch; + +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.DistributionSummary; +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.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.MockedConstruction; +import org.mockito.MockedStatic; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.client.RestHighLevelClient; +import org.opensearch.client.opensearch.OpenSearchClient; +import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; +import org.opensearch.dataprepper.expression.ExpressionEvaluator; +import org.opensearch.dataprepper.metrics.MetricNames; +import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.configuration.PluginSetting; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.event.EventHandle; +import org.opensearch.dataprepper.model.event.JacksonEvent; +import org.opensearch.dataprepper.model.failures.DlqObject; +import org.opensearch.dataprepper.model.plugin.PluginFactory; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.sink.SinkContext; +import org.opensearch.dataprepper.plugins.sink.opensearch.dlq.FailedDlqData; +import org.opensearch.dataprepper.plugins.sink.opensearch.index.DocumentBuilder; +import org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexConfiguration; +import org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexManager; +import org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexManagerFactory; +import org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexType; +import org.opensearch.dataprepper.plugins.sink.opensearch.index.TemplateStrategy; +import org.opensearch.dataprepper.plugins.sink.opensearch.index.TemplateType; + +import java.io.IOException; +import java.util.List; +import java.util.Optional; +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.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockConstruction; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.model.sink.SinkLatencyMetrics.EXTERNAL_LATENCY; +import static org.opensearch.dataprepper.model.sink.SinkLatencyMetrics.INTERNAL_LATENCY; +import static org.opensearch.dataprepper.plugins.sink.opensearch.OpenSearchSink.BULKREQUEST_ERRORS; +import static org.opensearch.dataprepper.plugins.sink.opensearch.OpenSearchSink.BULKREQUEST_LATENCY; +import static org.opensearch.dataprepper.plugins.sink.opensearch.OpenSearchSink.BULKREQUEST_SIZE_BYTES; +import static org.opensearch.dataprepper.plugins.sink.opensearch.OpenSearchSink.DYNAMIC_INDEX_DROPPED_EVENTS; +import static org.opensearch.dataprepper.plugins.sink.opensearch.OpenSearchSink.INVALID_ACTION_ERRORS; +import static org.opensearch.dataprepper.plugins.sink.opensearch.OpenSearchSink.INVALID_VERSION_EXPRESSION_DROPPED_EVENTS; +import static org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexConfiguration.DEFAULT_BULK_SIZE; +import static org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexConfiguration.DEFAULT_FLUSH_TIMEOUT; + +@ExtendWith(MockitoExtension.class) +public class OpenSearchSinkTest { + + @Mock + private IndexManagerFactory indexManagerFactory; + + @Mock + private OpenSearchClient openSearchClient; + + @Mock + private PluginFactory pluginFactory; + + @Mock + private SinkContext sinkContext; + + @Mock + private PluginSetting pluginSetting; + + @Mock + private ExpressionEvaluator expressionEvaluator; + + @Mock + private AwsCredentialsSupplier awsCredentialsSupplier; + + @Mock + private OpenSearchSinkConfiguration openSearchSinkConfiguration; + + @Mock + private IndexConfiguration indexConfiguration; + + @Mock + private IndexManager indexManager; + + @Mock + private PluginMetrics pluginMetrics; + + @Mock + private Timer bulkRequestTimer; + + @Mock + private Counter bulkRequestErrorsCounter; + + @Mock + private Counter invalidActionErrorsCounter; + + @Mock + private Counter dynamicIndexDroppedEvents; + + @Mock + private DistributionSummary bulkRequestSizeBytesSummary; + + @Mock + private Counter dynamicDocumentVersionDroppedEvents; + + @BeforeEach + void setup() { + when(pluginSetting.getPipelineName()).thenReturn(UUID.randomUUID().toString()); + when(pluginSetting.getName()).thenReturn(UUID.randomUUID().toString()); + + final RetryConfiguration retryConfiguration = mock(RetryConfiguration.class); + when(retryConfiguration.getDlq()).thenReturn(Optional.empty()); + when(retryConfiguration.getDlqFile()).thenReturn(null); + + final ConnectionConfiguration connectionConfiguration = mock(ConnectionConfiguration.class); + final RestHighLevelClient restHighLevelClient = mock(RestHighLevelClient.class); + when(connectionConfiguration.createClient(awsCredentialsSupplier)).thenReturn(restHighLevelClient); + when(connectionConfiguration.createOpenSearchClient(restHighLevelClient, awsCredentialsSupplier)).thenReturn(openSearchClient); + + when(indexConfiguration.getAction()).thenReturn("index"); + when(indexConfiguration.getDocumentId()).thenReturn(null); + when(indexConfiguration.getDocumentIdField()).thenReturn(null); + when(indexConfiguration.getRoutingField()).thenReturn(null); + when(indexConfiguration.getActions()).thenReturn(null); + when(indexConfiguration.getDocumentRootKey()).thenReturn(null); + when(indexConfiguration.getVersionType()).thenReturn(null); + when(indexConfiguration.getVersionExpression()).thenReturn(null); + when(indexConfiguration.getIndexAlias()).thenReturn(UUID.randomUUID().toString()); + when(indexConfiguration.getTemplateType()).thenReturn(TemplateType.V1); + when(indexConfiguration.getIndexType()).thenReturn(IndexType.CUSTOM); + when(indexConfiguration.getBulkSize()).thenReturn(DEFAULT_BULK_SIZE); + when(indexConfiguration.getFlushTimeout()).thenReturn(DEFAULT_FLUSH_TIMEOUT); + + when(openSearchSinkConfiguration.getIndexConfiguration()).thenReturn(indexConfiguration); + when(openSearchSinkConfiguration.getRetryConfiguration()).thenReturn(retryConfiguration); + when(openSearchSinkConfiguration.getConnectionConfiguration()).thenReturn(connectionConfiguration); + + when(pluginMetrics.counter(MetricNames.RECORDS_IN)).thenReturn(mock(Counter.class)); + when(pluginMetrics.timer(MetricNames.TIME_ELAPSED)).thenReturn(mock(Timer.class)); + when(pluginMetrics.summary(INTERNAL_LATENCY)).thenReturn(mock(DistributionSummary.class)); + when(pluginMetrics.summary(EXTERNAL_LATENCY)).thenReturn(mock(DistributionSummary.class)); + when(pluginMetrics.timer(BULKREQUEST_LATENCY)).thenReturn(bulkRequestTimer); + when(pluginMetrics.counter(BULKREQUEST_ERRORS)).thenReturn(bulkRequestErrorsCounter); + when(pluginMetrics.counter(INVALID_ACTION_ERRORS)).thenReturn(invalidActionErrorsCounter); + when(pluginMetrics.counter(DYNAMIC_INDEX_DROPPED_EVENTS)).thenReturn(dynamicIndexDroppedEvents); + when(pluginMetrics.counter(INVALID_VERSION_EXPRESSION_DROPPED_EVENTS)).thenReturn(dynamicDocumentVersionDroppedEvents); + when(pluginMetrics.summary(BULKREQUEST_SIZE_BYTES)).thenReturn(bulkRequestSizeBytesSummary); + + when(sinkContext.getTagsTargetKey()).thenReturn(null); + when(sinkContext.getIncludeKeys()).thenReturn(null); + when(sinkContext.getExcludeKeys()).thenReturn(null); + } + + private OpenSearchSink createObjectUnderTest() throws IOException { + try (final MockedStatic openSearchSinkConfigurationMockedStatic = mockStatic(OpenSearchSinkConfiguration.class); + final MockedStatic pluginMetricsMockedStatic = mockStatic(PluginMetrics.class); + final MockedConstruction indexManagerFactoryMockedConstruction = mockConstruction(IndexManagerFactory.class, (mock, context) -> { + indexManagerFactory = mock; + })) { + pluginMetricsMockedStatic.when(() -> PluginMetrics.fromPluginSetting(pluginSetting)).thenReturn(pluginMetrics); + openSearchSinkConfigurationMockedStatic.when(() -> OpenSearchSinkConfiguration.readESConfig(pluginSetting, expressionEvaluator)) + .thenReturn(openSearchSinkConfiguration); + return new OpenSearchSink(pluginSetting, pluginFactory, sinkContext, expressionEvaluator, awsCredentialsSupplier); + } + } + + @Test + void doOutput_with_invalid_version_expression_catches_NumberFormatException_and_creates_DLQObject() throws IOException { + + final String versionExpression = UUID.randomUUID().toString(); + when(indexConfiguration.getVersionExpression()).thenReturn(versionExpression); + + final Event event = mock(JacksonEvent.class); + final String document = UUID.randomUUID().toString(); + when(event.toJsonString()).thenReturn(document); + final EventHandle eventHandle = mock(EventHandle.class); + when(event.getEventHandle()).thenReturn(eventHandle); + final String index = UUID.randomUUID().toString(); + when(event.formatString(versionExpression, expressionEvaluator)).thenReturn("not_a_number"); + when(event.formatString(indexConfiguration.getIndexAlias(), expressionEvaluator)).thenReturn(index); + final Record eventRecord = new Record<>(event); + + final OpenSearchSink objectUnderTest = createObjectUnderTest(); + when(indexManagerFactory.getIndexManager(any(IndexType.class), eq(openSearchClient), any(RestHighLevelClient.class), eq(openSearchSinkConfiguration), any(TemplateStrategy.class), any())) + .thenReturn(indexManager); + doNothing().when(indexManager).setupIndex(); + objectUnderTest.initialize(); + + when(indexManager.getIndexName(anyString())).thenReturn(index); + + final DlqObject dlqObject = mock(DlqObject.class); + + final DlqObject.Builder dlqObjectBuilder = mock(DlqObject.Builder.class); + final ArgumentCaptor failedDlqData = ArgumentCaptor.forClass(FailedDlqData.class); + when(dlqObjectBuilder.withEventHandle(eventHandle)).thenReturn(dlqObjectBuilder); + when(dlqObjectBuilder.withFailedData(failedDlqData.capture())).thenReturn(dlqObjectBuilder); + when(dlqObjectBuilder.withPluginName(pluginSetting.getName())).thenReturn(dlqObjectBuilder); + when(dlqObjectBuilder.withPluginId(pluginSetting.getName())).thenReturn(dlqObjectBuilder); + when(dlqObjectBuilder.withPipelineName(pluginSetting.getPipelineName())).thenReturn(dlqObjectBuilder); + + when(dlqObject.getFailedData()).thenReturn(mock(FailedDlqData.class)); + doNothing().when(dlqObject).releaseEventHandle(false); + when(dlqObjectBuilder.build()).thenReturn(dlqObject); + + try (final MockedStatic documentBuilderMockedStatic = mockStatic(DocumentBuilder.class); + final MockedStatic dlqObjectMockedStatic = mockStatic(DlqObject.class)) { + documentBuilderMockedStatic.when(() -> DocumentBuilder.build(eq(event), eq(null), eq(null), eq(null), eq(null))) + .thenReturn(UUID.randomUUID().toString()); + + dlqObjectMockedStatic.when(DlqObject::builder).thenReturn(dlqObjectBuilder); + objectUnderTest.doOutput(List.of(eventRecord)); + } + + final FailedDlqData failedDlqDataResult = failedDlqData.getValue(); + assertThat(failedDlqDataResult, notNullValue()); + assertThat(failedDlqDataResult.getDocument(), equalTo(document)); + assertThat(failedDlqDataResult.getIndex(), equalTo(index)); + assertThat(failedDlqDataResult.getMessage().startsWith("Unable to convert the result of evaluating document_version"), equalTo(true)); + + verify(dynamicDocumentVersionDroppedEvents).increment(); + } + + @Test + void doOutput_with_invalid_version_expression_result_catches_RuntimeException_and_creates_DLQObject() throws IOException { + + final String versionExpression = UUID.randomUUID().toString(); + when(indexConfiguration.getVersionExpression()).thenReturn(versionExpression); + + final Event event = mock(JacksonEvent.class); + final String document = UUID.randomUUID().toString(); + when(event.toJsonString()).thenReturn(document); + final EventHandle eventHandle = mock(EventHandle.class); + when(event.getEventHandle()).thenReturn(eventHandle); + final String index = UUID.randomUUID().toString(); + when(event.formatString(versionExpression, expressionEvaluator)).thenThrow(RuntimeException.class); + when(event.formatString(indexConfiguration.getIndexAlias(), expressionEvaluator)).thenReturn(index); + final Record eventRecord = new Record<>(event); + + final OpenSearchSink objectUnderTest = createObjectUnderTest(); + when(indexManagerFactory.getIndexManager(any(IndexType.class), eq(openSearchClient), any(RestHighLevelClient.class), eq(openSearchSinkConfiguration), any(TemplateStrategy.class), any())) + .thenReturn(indexManager); + doNothing().when(indexManager).setupIndex(); + objectUnderTest.initialize(); + + when(indexManager.getIndexName(anyString())).thenReturn(index); + + final DlqObject dlqObject = mock(DlqObject.class); + + final DlqObject.Builder dlqObjectBuilder = mock(DlqObject.Builder.class); + final ArgumentCaptor failedDlqData = ArgumentCaptor.forClass(FailedDlqData.class); + when(dlqObjectBuilder.withEventHandle(eventHandle)).thenReturn(dlqObjectBuilder); + when(dlqObjectBuilder.withFailedData(failedDlqData.capture())).thenReturn(dlqObjectBuilder); + when(dlqObjectBuilder.withPluginName(pluginSetting.getName())).thenReturn(dlqObjectBuilder); + when(dlqObjectBuilder.withPluginId(pluginSetting.getName())).thenReturn(dlqObjectBuilder); + when(dlqObjectBuilder.withPipelineName(pluginSetting.getPipelineName())).thenReturn(dlqObjectBuilder); + + when(dlqObject.getFailedData()).thenReturn(mock(FailedDlqData.class)); + doNothing().when(dlqObject).releaseEventHandle(false); + when(dlqObjectBuilder.build()).thenReturn(dlqObject); + + try (final MockedStatic documentBuilderMockedStatic = mockStatic(DocumentBuilder.class); + final MockedStatic dlqObjectMockedStatic = mockStatic(DlqObject.class)) { + documentBuilderMockedStatic.when(() -> DocumentBuilder.build(eq(event), eq(null), eq(null), eq(null), eq(null))) + .thenReturn(UUID.randomUUID().toString()); + + dlqObjectMockedStatic.when(DlqObject::builder).thenReturn(dlqObjectBuilder); + objectUnderTest.doOutput(List.of(eventRecord)); + } + + final FailedDlqData failedDlqDataResult = failedDlqData.getValue(); + assertThat(failedDlqDataResult, notNullValue()); + assertThat(failedDlqDataResult.getDocument(), equalTo(document)); + assertThat(failedDlqDataResult.getIndex(), equalTo(index)); + assertThat(failedDlqDataResult.getMessage().startsWith("There was an exception when evaluating the document_version"), equalTo(true)); + + verify(dynamicDocumentVersionDroppedEvents).increment(); + } +} diff --git a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/IndexConfigurationTests.java b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/IndexConfigurationTests.java index b0878f53e9..99251fb956 100644 --- a/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/IndexConfigurationTests.java +++ b/data-prepper-plugins/opensearch/src/test/java/org/opensearch/dataprepper/plugins/sink/opensearch/index/IndexConfigurationTests.java @@ -481,7 +481,7 @@ public void testReadIndexConfig_emptyDocumentRootKey() { public void testReadIndexConfig_withValidDocumentVersionExpression(final String versionExpression) { final ExpressionEvaluator expressionEvaluator = mock(ExpressionEvaluator.class); - when(expressionEvaluator.isValidFormatExpressions(versionExpression)).thenReturn(true); + when(expressionEvaluator.isValidFormatExpression(versionExpression)).thenReturn(true); final Map metadata = initializeConfigMetaData( IndexType.CUSTOM.getValue(), "foo", null, null, null, null, null); @@ -500,7 +500,7 @@ public void testReadIndexConfig_withInvalidDocumentVersionExpression_throws_Inva final String versionExpression = UUID.randomUUID().toString(); final ExpressionEvaluator expressionEvaluator = mock(ExpressionEvaluator.class); - when(expressionEvaluator.isValidFormatExpressions(versionExpression)).thenReturn(false); + when(expressionEvaluator.isValidFormatExpression(versionExpression)).thenReturn(false); final Map metadata = initializeConfigMetaData( IndexType.CUSTOM.getValue(), "foo", null, null, null, null, null);