Skip to content

Commit 1f9e20f

Browse files
mzurita-amzmzurita-amz
committed
Add endpoint override support to CloudWatch Logs sink
Signed-off-by: Manuel Mangas Zurita <mzurita@amazon.com>
1 parent 13049d3 commit 1f9e20f

7 files changed

Lines changed: 138 additions & 15 deletions

File tree

‎data-prepper-plugins/cloudwatch-logs/README.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ pipeline:
3434
X-Custom-Header: "custom-value"
3535
X-Request-ID: "request-123"
3636
X-Source: "dataprepper"
37+
endpoint: "https://logs.us-west-2.amazonaws.com"
3738
```
3839

3940
## AWS Configuration
@@ -62,6 +63,8 @@ pipeline:
6263

6364
- `header_overrides` (Optional) : A string map representing custom HTTP headers to include in CloudWatch Logs requests.
6465

66+
- `endpoint` (Optional) : A string representing a custom CloudWatch Logs endpoint URL to override the default service endpoint.
67+
6568
## Buffer Type Configuration
6669

6770
- `buffer_type` (Optional) : A string representing the type of buffer to use to hold onto events. Currently only supports `in_memory`.

‎data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/CloudWatchLogsSink.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ public CloudWatchLogsSink(final PluginSetting pluginSetting,
7171
if (awsConfig == null && awsCredentialsSupplier == null) {
7272
throw new RuntimeException("Missing awsConfig and awsCredentialsSupplier");
7373
}
74-
CloudWatchLogsClient cloudWatchLogsClient = CloudWatchLogsClientFactory.createCwlClient(awsConfig, awsCredentialsSupplier, headerOverrides);
74+
CloudWatchLogsClient cloudWatchLogsClient = CloudWatchLogsClientFactory.createCwlClient(awsConfig, awsCredentialsSupplier, headerOverrides, cloudWatchLogsSinkConfig.getEndpoint());
7575
if (cloudWatchLogsClient == null) {
7676
throw new RuntimeException("cloudWatchLogsClient is null");
7777
}

‎data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsClientFactory.java‎

Lines changed: 29 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,9 @@
1313
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
1414
import software.amazon.awssdk.regions.Region;
1515
import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClient;
16+
import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClientBuilder;
17+
18+
import java.net.URI;
1619

1720
import java.util.HashMap;
1821
import java.util.Map;
@@ -32,11 +35,13 @@ private CloudWatchLogsClientFactory() {
3235
* @param awsConfig AwsConfig specifying region, roles, and header overrides.
3336
* @param awsCredentialsSupplier AwsCredentialsSupplier Interface for which to create CredentialsProvider for Client config.
3437
* @param customHeaders Map of custom headers to include in requests. Can be null.
38+
* @param endpoint Optional endpoint URL to override the default CloudWatch Logs endpoint.
3539
* @return CloudWatchLogsClient used to interact with CloudWatch Logs services.
3640
*/
3741
public static CloudWatchLogsClient createCwlClient(final AwsConfig awsConfig,
3842
final AwsCredentialsSupplier awsCredentialsSupplier,
39-
final Map<String, String> customHeaders) {
43+
final Map<String, String> customHeaders,
44+
final String endpoint) {
4045
final AwsCredentialsOptions awsCredentialsOptions = awsConfig != null
4146
? convertToCredentialOptions(awsConfig)
4247
: AwsCredentialsOptions.defaultOptions();
@@ -47,11 +52,30 @@ public static CloudWatchLogsClient createCwlClient(final AwsConfig awsConfig,
4752
return null;
4853
}
4954

50-
return CloudWatchLogsClient.builder()
55+
CloudWatchLogsClientBuilder clientBuilder = CloudWatchLogsClient.builder()
5156
.region(region)
5257
.credentialsProvider(awsCredentialsProvider)
53-
.overrideConfiguration(createOverrideConfiguration(customHeaders))
54-
.build();
58+
.overrideConfiguration(createOverrideConfiguration(customHeaders));
59+
60+
if (endpoint != null && !endpoint.trim().isEmpty()) {
61+
clientBuilder.endpointOverride(URI.create(endpoint));
62+
}
63+
64+
return clientBuilder.build();
65+
}
66+
67+
/**
68+
* Generates a CloudWatchLogs Client based on STS role ARN system credentials.
69+
*
70+
* @param awsConfig AwsConfig specifying region, roles, and header overrides.
71+
* @param awsCredentialsSupplier AwsCredentialsSupplier Interface for which to create CredentialsProvider for Client config.
72+
* @param customHeaders Map of custom headers to include in requests. Can be null.
73+
* @return CloudWatchLogsClient used to interact with CloudWatch Logs services.
74+
*/
75+
public static CloudWatchLogsClient createCwlClient(final AwsConfig awsConfig,
76+
final AwsCredentialsSupplier awsCredentialsSupplier,
77+
final Map<String, String> customHeaders) {
78+
return createCwlClient(awsConfig, awsCredentialsSupplier, customHeaders, null);
5579
}
5680

5781
/**
@@ -63,7 +87,7 @@ public static CloudWatchLogsClient createCwlClient(final AwsConfig awsConfig,
6387
*/
6488
public static CloudWatchLogsClient createCwlClient(final AwsConfig awsConfig,
6589
final AwsCredentialsSupplier awsCredentialsSupplier) {
66-
return createCwlClient(awsConfig, awsCredentialsSupplier, new HashMap<>());
90+
return createCwlClient(awsConfig, awsCredentialsSupplier, new HashMap<>(), null);
6791
}
6892

6993
private static ClientOverrideConfiguration createOverrideConfiguration(final Map<String, String> customHeaders) {

‎data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/config/CloudWatchLogsSinkConfig.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,9 @@ public class CloudWatchLogsSinkConfig {
5757
@ValidCustomHeaders
5858
private Map<String, String> headerOverrides = new HashMap<>();
5959

60+
@JsonProperty("endpoint")
61+
private String endpoint;
62+
6063
public AwsConfig getAwsConfig() {
6164
return awsConfig;
6265
}
@@ -89,4 +92,8 @@ public Map<String, String> getHeaderOverrides() {
8992
return headerOverrides;
9093
}
9194

95+
public String getEndpoint() {
96+
return endpoint;
97+
}
98+
9299
}

‎data-prepper-plugins/cloudwatch-logs/src/test/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/CloudWatchLogsSinkTest.java‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ Collection<Record<Event>> getMockedRecords() {
100100
void WHEN_sink_is_initialized_THEN_sink_is_ready_returns_true() {
101101
try(MockedStatic<CloudWatchLogsClientFactory> mockedStatic = mockStatic(CloudWatchLogsClientFactory.class)) {
102102
mockedStatic.when(() -> CloudWatchLogsClientFactory.createCwlClient(any(AwsConfig.class),
103-
any(AwsCredentialsSupplier.class), any()))
103+
any(AwsCredentialsSupplier.class), any(), any()))
104104
.thenReturn(mockClient);
105105

106106
CloudWatchLogsSink testCloudWatchSink = getTestCloudWatchSink();
@@ -115,7 +115,7 @@ void WHEN_awsConfig_and_awsCredentialsSupplier_null_THEN_should_throw() {
115115
when(mockCloudWatchLogsSinkConfig.getAwsConfig()).thenReturn(null);
116116
try(MockedStatic<CloudWatchLogsClientFactory> mockedStatic = mockStatic(CloudWatchLogsClientFactory.class)) {
117117
mockedStatic.when(() -> CloudWatchLogsClientFactory.createCwlClient(any(AwsConfig.class),
118-
any(AwsCredentialsSupplier.class), any()))
118+
any(AwsCredentialsSupplier.class), any(), any()))
119119
.thenReturn(mockClient);
120120

121121
assertThrows(RuntimeException.class, ()-> getTestCloudWatchSink());
@@ -126,7 +126,7 @@ void WHEN_awsConfig_and_awsCredentialsSupplier_null_THEN_should_throw() {
126126
void WHEN_given_sample_empty_records_THEN_records_are_processed() {
127127
try(MockedStatic<CloudWatchLogsClientFactory> mockedStatic = mockStatic(CloudWatchLogsClientFactory.class)) {
128128
mockedStatic.when(() -> CloudWatchLogsClientFactory.createCwlClient(any(AwsConfig.class),
129-
any(AwsCredentialsSupplier.class), any()))
129+
any(AwsCredentialsSupplier.class), any(), any()))
130130
.thenReturn(mockClient);
131131

132132
CloudWatchLogsSink testCloudWatchSink = getTestCloudWatchSink();
@@ -145,7 +145,7 @@ void WHEN_given_sample_empty_records_THEN_records_are_processed() {
145145
void WHEN_given_sample_empty_records_THEN_records_are_not_processed() {
146146
try(MockedStatic<CloudWatchLogsClientFactory> mockedStatic = mockStatic(CloudWatchLogsClientFactory.class)) {
147147
mockedStatic.when(() -> CloudWatchLogsClientFactory.createCwlClient(any(AwsConfig.class),
148-
any(AwsCredentialsSupplier.class), any()))
148+
any(AwsCredentialsSupplier.class), any(), any()))
149149
.thenReturn(mockClient);
150150

151151
CloudWatchLogsSink testCloudWatchSink = getTestCloudWatchSink();
@@ -166,15 +166,16 @@ void WHEN_header_overrides_is_empty_THEN_empty_map_is_passed_to_client_factory()
166166

167167
try(MockedStatic<CloudWatchLogsClientFactory> mockedStatic = mockStatic(CloudWatchLogsClientFactory.class)) {
168168
mockedStatic.when(() -> CloudWatchLogsClientFactory.createCwlClient(any(AwsConfig.class),
169-
any(AwsCredentialsSupplier.class), any()))
169+
any(AwsCredentialsSupplier.class), any(), any()))
170170
.thenReturn(mockClient);
171171

172172
CloudWatchLogsSink testCloudWatchSink = getTestCloudWatchSink();
173173

174174
mockedStatic.verify(() -> CloudWatchLogsClientFactory.createCwlClient(
175175
eq(mockAwsConfig),
176176
eq(mockCredentialSupplier),
177-
eq(emptyHeaders)));
177+
eq(emptyHeaders),
178+
any()));
178179
}
179180
}
180181

@@ -184,15 +185,16 @@ void WHEN_header_overrides_is_provided_THEN_headers_are_passed_to_client_factory
184185

185186
try(MockedStatic<CloudWatchLogsClientFactory> mockedStatic = mockStatic(CloudWatchLogsClientFactory.class)) {
186187
mockedStatic.when(() -> CloudWatchLogsClientFactory.createCwlClient(any(AwsConfig.class),
187-
any(AwsCredentialsSupplier.class), any()))
188+
any(AwsCredentialsSupplier.class), any(), any()))
188189
.thenReturn(mockClient);
189190

190191
CloudWatchLogsSink testCloudWatchSink = getTestCloudWatchSink();
191192

192193
mockedStatic.verify(() -> CloudWatchLogsClientFactory.createCwlClient(
193194
eq(mockAwsConfig),
194195
eq(mockCredentialSupplier),
195-
eq(mockHeaderOverrides)));
196+
eq(mockHeaderOverrides),
197+
any()));
196198
}
197199
}
198200

@@ -202,7 +204,7 @@ void WHEN_sink_initialization_with_header_overrides_THEN_sink_is_ready() {
202204

203205
try(MockedStatic<CloudWatchLogsClientFactory> mockedStatic = mockStatic(CloudWatchLogsClientFactory.class)) {
204206
mockedStatic.when(() -> CloudWatchLogsClientFactory.createCwlClient(any(AwsConfig.class),
205-
any(AwsCredentialsSupplier.class), any()))
207+
any(AwsCredentialsSupplier.class), any(), any()))
206208
.thenReturn(mockClient);
207209

208210
CloudWatchLogsSink testCloudWatchSink = getTestCloudWatchSink();

‎data-prepper-plugins/cloudwatch-logs/src/test/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/client/CloudWatchLogsClientFactoryTest.java‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClient;
2222
import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClientBuilder;
2323

24+
import java.net.URI;
25+
2426

2527
import java.util.HashMap;
2628
import java.util.List;
@@ -184,4 +186,76 @@ void GIVEN_empty_custom_headers_SHOULD_create_client_without_custom_headers() {
184186
assertThat(actualConfig.headers().isEmpty(), is(true));
185187
}
186188
}
189+
190+
@Test
191+
void GIVEN_endpoint_SHOULD_create_client_with_endpoint_override() {
192+
// Arrange
193+
final String customEndpoint = "https://logs.us-west-2.amazonaws.com";
194+
when(mockAwsCredentialsSupplier.getProvider(any())).thenReturn(mockAwsCredentialsProvider);
195+
196+
final CloudWatchLogsClientBuilder cloudWatchLogsClientBuilder = mock(CloudWatchLogsClientBuilder.class);
197+
when(cloudWatchLogsClientBuilder.region(any())).thenReturn(cloudWatchLogsClientBuilder);
198+
when(cloudWatchLogsClientBuilder.credentialsProvider(any())).thenReturn(cloudWatchLogsClientBuilder);
199+
when(cloudWatchLogsClientBuilder.overrideConfiguration(any(ClientOverrideConfiguration.class))).thenReturn(cloudWatchLogsClientBuilder);
200+
when(cloudWatchLogsClientBuilder.endpointOverride(any(URI.class))).thenReturn(cloudWatchLogsClientBuilder);
201+
202+
try(final MockedStatic<CloudWatchLogsClient> cloudWatchLogsClientMockedStatic = mockStatic(CloudWatchLogsClient.class)) {
203+
cloudWatchLogsClientMockedStatic.when(CloudWatchLogsClient::builder)
204+
.thenReturn(cloudWatchLogsClientBuilder);
205+
206+
// Act
207+
CloudWatchLogsClientFactory.createCwlClient(mockAwsConfig, mockAwsCredentialsSupplier, new HashMap<>(), customEndpoint);
208+
209+
// Assert
210+
final ArgumentCaptor<URI> endpointCaptor = ArgumentCaptor.forClass(URI.class);
211+
verify(cloudWatchLogsClientBuilder).endpointOverride(endpointCaptor.capture());
212+
213+
final URI actualEndpoint = endpointCaptor.getValue();
214+
assertThat(actualEndpoint.toString(), equalTo(customEndpoint));
215+
}
216+
}
217+
218+
@Test
219+
void GIVEN_null_endpoint_SHOULD_create_client_without_endpoint_override() {
220+
// Arrange
221+
when(mockAwsCredentialsSupplier.getProvider(any())).thenReturn(mockAwsCredentialsProvider);
222+
223+
final CloudWatchLogsClientBuilder cloudWatchLogsClientBuilder = mock(CloudWatchLogsClientBuilder.class);
224+
when(cloudWatchLogsClientBuilder.region(any())).thenReturn(cloudWatchLogsClientBuilder);
225+
when(cloudWatchLogsClientBuilder.credentialsProvider(any())).thenReturn(cloudWatchLogsClientBuilder);
226+
when(cloudWatchLogsClientBuilder.overrideConfiguration(any(ClientOverrideConfiguration.class))).thenReturn(cloudWatchLogsClientBuilder);
227+
228+
try(final MockedStatic<CloudWatchLogsClient> cloudWatchLogsClientMockedStatic = mockStatic(CloudWatchLogsClient.class)) {
229+
cloudWatchLogsClientMockedStatic.when(CloudWatchLogsClient::builder)
230+
.thenReturn(cloudWatchLogsClientBuilder);
231+
232+
// Act
233+
CloudWatchLogsClientFactory.createCwlClient(mockAwsConfig, mockAwsCredentialsSupplier, new HashMap<>(), null);
234+
235+
// Assert - verify endpointOverride is never called
236+
verify(cloudWatchLogsClientBuilder, org.mockito.Mockito.never()).endpointOverride(any(URI.class));
237+
}
238+
}
239+
240+
@Test
241+
void GIVEN_empty_endpoint_SHOULD_create_client_without_endpoint_override() {
242+
// Arrange
243+
when(mockAwsCredentialsSupplier.getProvider(any())).thenReturn(mockAwsCredentialsProvider);
244+
245+
final CloudWatchLogsClientBuilder cloudWatchLogsClientBuilder = mock(CloudWatchLogsClientBuilder.class);
246+
when(cloudWatchLogsClientBuilder.region(any())).thenReturn(cloudWatchLogsClientBuilder);
247+
when(cloudWatchLogsClientBuilder.credentialsProvider(any())).thenReturn(cloudWatchLogsClientBuilder);
248+
when(cloudWatchLogsClientBuilder.overrideConfiguration(any(ClientOverrideConfiguration.class))).thenReturn(cloudWatchLogsClientBuilder);
249+
250+
try(final MockedStatic<CloudWatchLogsClient> cloudWatchLogsClientMockedStatic = mockStatic(CloudWatchLogsClient.class)) {
251+
cloudWatchLogsClientMockedStatic.when(CloudWatchLogsClient::builder)
252+
.thenReturn(cloudWatchLogsClientBuilder);
253+
254+
// Act
255+
CloudWatchLogsClientFactory.createCwlClient(mockAwsConfig, mockAwsCredentialsSupplier, new HashMap<>(), "");
256+
257+
// Assert - verify endpointOverride is never called
258+
verify(cloudWatchLogsClientBuilder, org.mockito.Mockito.never()).endpointOverride(any(URI.class));
259+
}
260+
}
187261
}

‎data-prepper-plugins/cloudwatch-logs/src/test/java/org/opensearch/dataprepper/plugins/sink/cloudwatch_logs/config/CloudWatchLogsSinkConfigTest.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -309,4 +309,17 @@ void GIVEN_default_config_WHEN_accessed_THEN_header_overrides_should_be_empty()
309309
assertThat(config.getHeaderOverrides(), aMapWithSize(0));
310310
assertThat(config.getHeaderOverrides().size(), lessThanOrEqualTo(10));
311311
}
312+
313+
@Test
314+
void GIVEN_new_sink_config_WHEN_get_endpoint_called_SHOULD_return_null() {
315+
assertThat(new CloudWatchLogsSinkConfig().getEndpoint(), equalTo(null));
316+
}
317+
318+
@Test
319+
void GIVEN_endpoint_configured_SHOULD_return_the_configured_value() throws NoSuchFieldException, IllegalAccessException {
320+
String testEndpoint = "https://logs.us-west-2.amazonaws.com";
321+
322+
ReflectivelySetField.setField(cloudWatchLogsSinkConfig.getClass(), cloudWatchLogsSinkConfig, "endpoint", testEndpoint);
323+
assertThat(cloudWatchLogsSinkConfig.getEndpoint(), equalTo(testEndpoint));
324+
}
312325
}

0 commit comments

Comments
 (0)