Skip to content

Commit 97cd930

Browse files
daixbadaixba
authored
Add CMK encryption support to DynamoDB export (#3592)
Signed-off-by: Aiden Dai <daixb@amazon.com>
1 parent 15d8247 commit 97cd930

10 files changed

Lines changed: 81 additions & 41 deletions

File tree

‎data-prepper-plugins/dynamodb-source/README.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@ source:
5252

5353
* s3_bucket (Required): The destination bucket to store the exported data files
5454
* s3_prefix (Optional): Custom prefix.
55+
* s3_sse_kms_key_id (Optional): A AWS KMS Customer Managed Key (CMK) to encrypt the export data files. The key id will
56+
be the ARN of the Key, e.g. arn:aws:kms:us-west-2:123456789012:key/0a4bc22f-bb96-4ad4-80ca-63b12b3ec147
5557

5658
### Stream Configurations
5759

‎data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/DynamoDBService.java‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,12 @@ public void init() {
171171
Instant startTime = Instant.now();
172172

173173
if (tableInfo.getMetadata().isExportRequired()) {
174-
createExportPartition(tableInfo.getTableArn(), startTime, tableInfo.getMetadata().getExportBucket(), tableInfo.getMetadata().getExportPrefix());
174+
createExportPartition(
175+
tableInfo.getTableArn(),
176+
startTime,
177+
tableInfo.getMetadata().getExportBucket(),
178+
tableInfo.getMetadata().getExportPrefix(),
179+
tableInfo.getMetadata().getExportKmsKeyId());
175180
}
176181

177182
if (tableInfo.getMetadata().isStreamRequired()) {
@@ -209,11 +214,12 @@ public void init() {
209214
* @param bucket Export bucket
210215
* @param prefix Export Prefix
211216
*/
212-
private void createExportPartition(String tableArn, Instant exportTime, String bucket, String prefix) {
217+
private void createExportPartition(String tableArn, Instant exportTime, String bucket, String prefix, String kmsKeyId) {
213218
ExportProgressState exportProgressState = new ExportProgressState();
214219
exportProgressState.setBucket(bucket);
215220
exportProgressState.setPrefix(prefix);
216221
exportProgressState.setExportTime(exportTime.toString()); // information purpose
222+
exportProgressState.setKmsKeyId(kmsKeyId);
217223
ExportPartition exportPartition = new ExportPartition(tableArn, exportTime, Optional.of(exportProgressState));
218224
coordinator.createPartition(exportPartition);
219225
}
@@ -310,6 +316,7 @@ private TableInfo getTableInfo(TableConfig tableConfig) {
310316
.streamStartPosition(streamStartPosition)
311317
.exportBucket(tableConfig.getExportConfig() == null ? null : tableConfig.getExportConfig().getS3Bucket())
312318
.exportPrefix(tableConfig.getExportConfig() == null ? null : tableConfig.getExportConfig().getS3Prefix())
319+
.exportKmsKeyId(tableConfig.getExportConfig() == null ? null : tableConfig.getExportConfig().getS3SseKmsKeyId())
313320
.build();
314321
return new TableInfo(tableConfig.getTableArn(), metadata);
315322
}

‎data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/configuration/ExportConfig.java‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,21 +6,25 @@
66
package org.opensearch.dataprepper.plugins.source.dynamodb.configuration;
77

88
import com.fasterxml.jackson.annotation.JsonProperty;
9+
import jakarta.validation.constraints.AssertTrue;
910
import jakarta.validation.constraints.NotBlank;
11+
import software.amazon.awssdk.arns.Arn;
1012
import software.amazon.awssdk.regions.Region;
1113

1214
public class ExportConfig {
1315

1416
@JsonProperty("s3_bucket")
1517
@NotBlank(message = "Bucket Name is required for export")
1618
private String s3Bucket;
17-
1819
@JsonProperty("s3_prefix")
1920
private String s3Prefix;
2021

2122
@JsonProperty("s3_region")
2223
private String s3Region;
2324

25+
@JsonProperty("s3_sse_kms_key_id")
26+
private String s3SseKmsKeyId;
27+
2428
public String getS3Bucket() {
2529
return s3Bucket;
2630
}
@@ -33,4 +37,15 @@ public Region getAwsRegion() {
3337
return s3Region != null ? Region.of(s3Region) : null;
3438
}
3539

40+
public String getS3SseKmsKeyId() {
41+
return s3SseKmsKeyId;
42+
}
43+
44+
@AssertTrue(message = "KMS Key ID must be a valid one.")
45+
boolean isKmsKeyIdValid() {
46+
// If key id is provided, it should be in a format like
47+
// arn:aws:kms:us-west-2:123456789012:key/0a4bc22f-bb96-4ad3-80ca-63b12b3ec147
48+
return s3SseKmsKeyId == null || Arn.fromString(s3SseKmsKeyId).resourceAsString() != null;
49+
}
50+
3651
}

‎data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/coordination/state/ExportProgressState.java‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,13 @@ public class ExportProgressState {
2121
@JsonProperty("prefix")
2222
private String prefix;
2323

24+
@JsonProperty("kmsKeyId")
25+
private String kmsKeyId;
26+
2427
@JsonProperty("exportTime")
2528
private String exportTime;
2629

27-
30+
2831
public String getExportArn() {
2932
return exportArn;
3033
}
@@ -64,4 +67,12 @@ public String getStatus() {
6467
public void setStatus(String status) {
6568
this.status = status;
6669
}
70+
71+
public String getKmsKeyId() {
72+
return kmsKeyId;
73+
}
74+
75+
public void setKmsKeyId(String kmsKeyId) {
76+
this.kmsKeyId = kmsKeyId;
77+
}
6778
}

‎data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoader.java‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,10 @@
77

88
import org.opensearch.dataprepper.buffer.common.BufferAccumulator;
99
import org.opensearch.dataprepper.metrics.PluginMetrics;
10+
import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
1011
import org.opensearch.dataprepper.model.buffer.Buffer;
1112
import org.opensearch.dataprepper.model.event.Event;
1213
import org.opensearch.dataprepper.model.record.Record;
13-
import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSet;
1414
import org.opensearch.dataprepper.plugins.source.dynamodb.converter.ExportRecordConverter;
1515
import org.opensearch.dataprepper.plugins.source.dynamodb.model.TableInfo;
1616
import org.slf4j.Logger;
@@ -169,14 +169,12 @@ public void run() {
169169
int lineCount = 0;
170170
int lastLineProcessed = 0;
171171

172-
try {
173-
InputStream inputStream = objectReader.readFile(bucketName, key);
174-
GZIPInputStream gzipInputStream = new GZIPInputStream(inputStream);
175-
BufferedReader reader = new BufferedReader(new InputStreamReader(gzipInputStream));
172+
try (InputStream inputStream = objectReader.readFile(bucketName, key);
173+
GZIPInputStream gzipInputStream = new GZIPInputStream(inputStream);
174+
BufferedReader reader = new BufferedReader(new InputStreamReader(gzipInputStream))) {
176175

177176
String line;
178177
while ((line = reader.readLine()) != null) {
179-
180178
if (shouldStop) {
181179
checkpointer.checkpoint(lastLineProcessed);
182180
LOG.debug("Should Stop flag is set to True, looks like shutdown has triggered");
@@ -213,9 +211,6 @@ public void run() {
213211
}
214212

215213
lines.clear();
216-
reader.close();
217-
gzipInputStream.close();
218-
inputStream.close();
219214

220215
LOG.info("Completed loading s3://{}/{} to buffer", bucketName, key);
221216

@@ -226,7 +221,7 @@ public void run() {
226221
} catch (Exception e) {
227222
checkpointer.checkpoint(lineCount);
228223

229-
String errorMessage = String.format("Loading of s3://%s/%s completed with Exception: %S", bucketName, key, e.getMessage());
224+
String errorMessage = String.format("Loading of s3://%s/%s completed with Exception: %s", bucketName, key, e.getMessage());
230225
throw new RuntimeException(errorMessage);
231226
}
232227
}

‎data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -156,8 +156,7 @@ private BiConsumer<String, Throwable> completeExport(ExportPartition exportParti
156156
ExportProgressState state = exportPartition.getProgressState().get();
157157
String bucketName = state.getBucket();
158158
String exportArn = state.getExportArn();
159-
160-
159+
161160
String manifestKey = exportTaskManager.getExportManifest(exportArn);
162161
LOG.debug("Export manifest summary file is " + manifestKey);
163162

@@ -189,6 +188,7 @@ private void createDataFilePartitions(String exportArn, String bucketName, Map<S
189188
DataFileProgressState progressState = new DataFileProgressState();
190189
progressState.setTotal(size);
191190
progressState.setLoaded(0);
191+
192192
totalFiles.addAndGet(1);
193193
totalRecords.addAndGet(size);
194194
DataFilePartition partition = new DataFilePartition(exportArn, bucketName, key, Optional.of(progressState));
@@ -260,7 +260,7 @@ private String getOrCreateExportArn(ExportPartition exportPartition) {
260260

261261
LOG.info("Try to submit a new export job for table {} with export time {}", exportPartition.getTableArn(), exportPartition.getExportTime());
262262
// submit a new export request
263-
String exportArn = exportTaskManager.submitExportJob(exportPartition.getTableArn(), state.getBucket(), state.getPrefix(), exportPartition.getExportTime());
263+
String exportArn = exportTaskManager.submitExportJob(exportPartition.getTableArn(), state.getBucket(), state.getPrefix(), state.getKmsKeyId(), exportPartition.getExportTime());
264264

265265
// Update state with export Arn in the coordination table.
266266
// So that it won't be submitted again after a restart.

‎data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportTaskManager.java‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import software.amazon.awssdk.services.dynamodb.model.ExportFormat;
1515
import software.amazon.awssdk.services.dynamodb.model.ExportTableToPointInTimeRequest;
1616
import software.amazon.awssdk.services.dynamodb.model.ExportTableToPointInTimeResponse;
17+
import software.amazon.awssdk.services.dynamodb.model.S3SseAlgorithm;
1718

1819
import java.time.Instant;
1920

@@ -30,12 +31,15 @@ public ExportTaskManager(DynamoDbClient dynamoDBClient) {
3031
this.dynamoDBClient = dynamoDBClient;
3132
}
3233

33-
public String submitExportJob(String tableArn, String bucketName, String prefix, Instant exportTime) {
34+
public String submitExportJob(String tableArn, String bucket, String prefix, String kmsKeyId, Instant exportTime) {
35+
S3SseAlgorithm algorithm = kmsKeyId == null || kmsKeyId.isEmpty() ? S3SseAlgorithm.AES256 : S3SseAlgorithm.KMS;
3436
// No needs to use a client token here.
3537
ExportTableToPointInTimeRequest req = ExportTableToPointInTimeRequest.builder()
3638
.tableArn(tableArn)
37-
.s3Bucket(bucketName)
39+
.s3Bucket(bucket)
3840
.s3Prefix(prefix)
41+
.s3SseAlgorithm(algorithm)
42+
.s3SseKmsKeyId(kmsKeyId)
3943
.exportFormat(DEFAULT_EXPORT_FORMAT)
4044
.exportTime(exportTime)
4145
.build();
@@ -46,7 +50,6 @@ public String submitExportJob(String tableArn, String bucketName, String prefix,
4650

4751
String exportArn = response.exportDescription().exportArn();
4852
String status = response.exportDescription().exportStatusAsString();
49-
5053
LOG.debug("Export Job submitted with ARN {} and status {}", exportArn, status);
5154
return exportArn;
5255
} catch (SdkException e) {

‎data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ManifestFileReader.java‎

Lines changed: 11 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,13 @@
55

66
package org.opensearch.dataprepper.plugins.source.dynamodb.export;
77

8-
import com.fasterxml.jackson.core.JsonProcessingException;
98
import com.fasterxml.jackson.core.type.TypeReference;
109
import com.fasterxml.jackson.databind.ObjectMapper;
1110
import org.opensearch.dataprepper.plugins.source.dynamodb.model.ExportSummary;
1211
import org.slf4j.Logger;
1312
import org.slf4j.LoggerFactory;
1413

1514
import java.io.BufferedReader;
16-
import java.io.IOException;
1715
import java.io.InputStream;
1816
import java.io.InputStreamReader;
1917
import java.util.HashMap;
@@ -39,22 +37,18 @@ public ManifestFileReader(S3ObjectReader objectReader) {
3937

4038
public ExportSummary parseSummaryFile(String bucket, String key) {
4139
LOG.debug("Try to read the manifest summary file");
42-
InputStream object = objectReader.readFile(bucket, key);
4340

44-
BufferedReader reader = new BufferedReader(new InputStreamReader(object));
45-
try {
41+
try (InputStream object = objectReader.readFile(bucket, key);
42+
BufferedReader reader = new BufferedReader(new InputStreamReader(object))) {
43+
// Only one line
4644
String line = reader.readLine();
4745
LOG.debug("Manifest summary: {}", line);
4846
ExportSummary summaryInfo = MAPPER.readValue(line, ExportSummary.class);
4947
return summaryInfo;
5048

51-
} catch (JsonProcessingException e) {
49+
} catch (Exception e) {
5250
LOG.error("Failed to parse the summary info due to {}", e.getMessage());
5351
throw new RuntimeException(e);
54-
55-
} catch (IOException e) {
56-
LOG.error("IO Exception due to {}", e.getMessage());
57-
throw new RuntimeException(e);
5852
}
5953

6054
}
@@ -63,11 +57,10 @@ public Map<String, Integer> parseDataFile(String bucket, String key) {
6357
LOG.info("Try to read the manifest data file");
6458

6559
Map<String, Integer> result = new HashMap<>();
66-
InputStream object = objectReader.readFile(bucket, key);
67-
BufferedReader reader = new BufferedReader(new InputStreamReader(object));
68-
69-
String line;
70-
try {
60+
61+
try (InputStream object = objectReader.readFile(bucket, key);
62+
BufferedReader reader = new BufferedReader(new InputStreamReader(object))) {
63+
String line;
7164
while ((line = reader.readLine()) != null) {
7265
// An example line as below:
7366
// {"itemCount":46331,"md5Checksum":"a0k21IY3eelgr2PuWJLjJw==","etag":"51f9f394903c5d682321c6211aae8b6a-1","dataFileS3Key":"test-table-export/AWSDynamoDB/01692350182719-6de2c037/data/fpgzwz7ome3s7a5gqn2mu3ogtq.json.gz"}
@@ -76,8 +69,9 @@ public Map<String, Integer> parseDataFile(String bucket, String key) {
7669
result.put(map.get(DATA_FILE_S3_KEY), Integer.valueOf(map.get(DATA_FILE_ITEM_COUNT_KEY)));
7770

7871
}
79-
} catch (IOException e) {
80-
LOG.error("IO Exception due to {}", e.getMessage());
72+
} catch (Exception e) {
73+
LOG.error("Exception due to {}", e.getMessage());
74+
throw new RuntimeException(e);
8175
}
8276

8377
return result;

‎data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/model/TableMetadata.java‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@ public class TableMetadata {
1919
private static final String REQUIRE_EXPORT_KEY = "export";
2020
private static final String REQUIRE_STREAM_KEY = "stream";
2121

22-
2322
private final String partitionKeyAttributeName;
2423

2524
private final String sortKeyAttributeName;
@@ -36,6 +35,8 @@ public class TableMetadata {
3635

3736
private final String exportPrefix;
3837

38+
private final String exportKmsKeyId;
39+
3940
private TableMetadata(Builder builder) {
4041
this.partitionKeyAttributeName = builder.partitionKeyAttributeName;
4142
this.sortKeyAttributeName = builder.sortKeyAttributeName;
@@ -45,6 +46,7 @@ private TableMetadata(Builder builder) {
4546
this.exportBucket = builder.exportBucket;
4647
this.exportPrefix = builder.exportPrefix;
4748
this.streamStartPosition = builder.streamStartPosition;
49+
this.exportKmsKeyId = builder.exportKmsKeyId;
4850

4951
}
5052

@@ -70,6 +72,8 @@ public static class Builder {
7072

7173
private String exportPrefix;
7274

75+
private String exportKmsKeyId;
76+
7377
private StreamStartPosition streamStartPosition;
7478

7579

@@ -108,6 +112,11 @@ public Builder exportPrefix(String exportPrefix) {
108112
return this;
109113
}
110114

115+
public Builder exportKmsKeyId(String exportKmsKeyId) {
116+
this.exportKmsKeyId = exportKmsKeyId;
117+
return this;
118+
}
119+
111120
public Builder streamStartPosition(StreamStartPosition streamStartPosition) {
112121
this.streamStartPosition = streamStartPosition;
113122
return this;
@@ -173,4 +182,8 @@ public String getExportBucket() {
173182
public String getExportPrefix() {
174183
return exportPrefix;
175184
}
185+
186+
public String getExportKmsKeyId() {
187+
return exportKmsKeyId;
188+
}
176189
}

‎data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/DataFileLoaderTest.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@ class DataFileLoaderTest {
9393

9494
private final Random random = new Random();
9595

96-
private final int total = random.nextInt(10);
96+
private final int total = random.nextInt(10) + 1;
9797

9898
@BeforeEach
9999
void setup() {
@@ -153,8 +153,8 @@ void test_run_loadFile_correctly() {
153153
try (
154154
final MockedStatic<BufferAccumulator> bufferAccumulatorMockedStatic = mockStatic(BufferAccumulator.class);
155155
final MockedConstruction<ExportRecordConverter> recordConverterMockedConstruction = mockConstruction(ExportRecordConverter.class, (mock, context) -> {
156-
exportRecordConverter = mock;
157-
})) {
156+
exportRecordConverter = mock;
157+
})) {
158158
bufferAccumulatorMockedStatic.when(() -> BufferAccumulator.create(buffer, DEFAULT_BUFFER_BATCH_SIZE, BUFFER_TIMEOUT)).thenReturn(bufferAccumulator);
159159
loader = DataFileLoader.builder(objectReader, pluginMetrics, buffer)
160160
.bucketName(bucketName)

0 commit comments

Comments
 (0)