|
| 1 | +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. |
| 2 | +// SPDX-License-Identifier: Apache-2.0 |
| 3 | + |
| 4 | +using Amazon.DynamoDBStreams; |
| 5 | +using Amazon.DynamoDBStreams.Model; |
| 6 | +using Amazon.Lambda.DynamoDBEvents; |
| 7 | +using Amazon.Lambda.Model; |
| 8 | +using Amazon.Lambda.TestTool.Services; |
| 9 | +using Amazon.Runtime; |
| 10 | +using System.Text.Json; |
| 11 | + |
| 12 | +namespace Amazon.Lambda.TestTool.Processes.DynamoDBStreamsEventSource; |
| 13 | + |
| 14 | +/// <summary> |
| 15 | +/// IHostedService that will run continually polling a DynamoDB Stream for records and invoking the connected |
| 16 | +/// Lambda function with the polled records. |
| 17 | +/// </summary> |
| 18 | +public class DynamoDBStreamsEventSourceBackgroundService : BackgroundService |
| 19 | +{ |
| 20 | + private static readonly JsonSerializerOptions _jsonOptions = new JsonSerializerOptions |
| 21 | + { |
| 22 | + PropertyNamingPolicy = JsonNamingPolicy.CamelCase |
| 23 | + }; |
| 24 | + |
| 25 | + private readonly ILogger<DynamoDBStreamsEventSourceProcess> _logger; |
| 26 | + private readonly IAmazonDynamoDBStreams _streamsClient; |
| 27 | + private readonly ILambdaClient _lambdaClient; |
| 28 | + private readonly DynamoDBStreamsEventSourceBackgroundServiceConfig _config; |
| 29 | + |
| 30 | + /// <summary> |
| 31 | + /// Constructs instance of <see cref="DynamoDBStreamsEventSourceBackgroundService"/>. |
| 32 | + /// </summary> |
| 33 | + public DynamoDBStreamsEventSourceBackgroundService( |
| 34 | + ILogger<DynamoDBStreamsEventSourceProcess> logger, |
| 35 | + IAmazonDynamoDBStreams streamsClient, |
| 36 | + DynamoDBStreamsEventSourceBackgroundServiceConfig config, |
| 37 | + ILambdaClient lambdaClient) |
| 38 | + { |
| 39 | + _logger = logger; |
| 40 | + _streamsClient = streamsClient; |
| 41 | + _config = config; |
| 42 | + _lambdaClient = lambdaClient; |
| 43 | + } |
| 44 | + |
| 45 | + /// <summary> |
| 46 | + /// Execute the DynamoDBStreamsEventSourceBackgroundService. |
| 47 | + /// </summary> |
| 48 | + protected override async Task ExecuteAsync(CancellationToken stoppingToken) |
| 49 | + { |
| 50 | + _logger.LogInformation("Starting DynamoDB Streams poller for table: {tableName}", _config.TableName); |
| 51 | + |
| 52 | + while (!stoppingToken.IsCancellationRequested) |
| 53 | + { |
| 54 | + try |
| 55 | + { |
| 56 | + var streamArn = await GetStreamArnForTable(stoppingToken); |
| 57 | + if (streamArn == null) |
| 58 | + { |
| 59 | + _logger.LogWarning("No stream found for table {tableName}. Retrying in 5 seconds.", _config.TableName); |
| 60 | + await Task.Delay(5000, stoppingToken); |
| 61 | + continue; |
| 62 | + } |
| 63 | + |
| 64 | + await PollStream(streamArn, stoppingToken); |
| 65 | + } |
| 66 | + catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested) |
| 67 | + { |
| 68 | + return; |
| 69 | + } |
| 70 | + catch (TaskCanceledException) when (stoppingToken.IsCancellationRequested) |
| 71 | + { |
| 72 | + return; |
| 73 | + } |
| 74 | + catch (Exception e) |
| 75 | + { |
| 76 | + _logger.LogWarning(e, "Exception occurred in DynamoDB Streams poller for {tableName}: {message}", _config.TableName, e.Message); |
| 77 | + await Task.Delay(3000, stoppingToken); |
| 78 | + } |
| 79 | + } |
| 80 | + } |
| 81 | + |
| 82 | + private async Task<string?> GetStreamArnForTable(CancellationToken stoppingToken) |
| 83 | + { |
| 84 | + var response = await _streamsClient.ListStreamsAsync(new ListStreamsRequest |
| 85 | + { |
| 86 | + TableName = _config.TableName |
| 87 | + }, stoppingToken); |
| 88 | + |
| 89 | + // Use the first active stream for the table |
| 90 | + return response.Streams.FirstOrDefault()?.StreamArn; |
| 91 | + } |
| 92 | + |
| 93 | + private async Task PollStream(string streamArn, CancellationToken stoppingToken) |
| 94 | + { |
| 95 | + var shardIterators = await GetShardIterators(streamArn, stoppingToken); |
| 96 | + |
| 97 | + while (!stoppingToken.IsCancellationRequested) |
| 98 | + { |
| 99 | + var hasRecords = false; |
| 100 | + |
| 101 | + for (int i = 0; i < shardIterators.Count; i++) |
| 102 | + { |
| 103 | + if (stoppingToken.IsCancellationRequested) |
| 104 | + return; |
| 105 | + |
| 106 | + var iterator = shardIterators[i]; |
| 107 | + if (iterator == null) |
| 108 | + continue; |
| 109 | + |
| 110 | + var getRecordsResponse = await _streamsClient.GetRecordsAsync(new GetRecordsRequest |
| 111 | + { |
| 112 | + ShardIterator = iterator, |
| 113 | + Limit = _config.BatchSize |
| 114 | + }, stoppingToken); |
| 115 | + |
| 116 | + shardIterators[i] = getRecordsResponse.NextShardIterator; |
| 117 | + |
| 118 | + if (getRecordsResponse.Records.Count == 0) |
| 119 | + continue; |
| 120 | + |
| 121 | + hasRecords = true; |
| 122 | + var lambdaRecords = ConvertToLambdaRecords(getRecordsResponse.Records, streamArn); |
| 123 | + |
| 124 | + var lambdaPayload = new DynamoDBEvent |
| 125 | + { |
| 126 | + Records = lambdaRecords |
| 127 | + }; |
| 128 | + |
| 129 | + var invokeRequest = new InvokeRequest |
| 130 | + { |
| 131 | + InvocationType = InvocationType.RequestResponse, |
| 132 | + FunctionName = _config.FunctionName, |
| 133 | + Payload = JsonSerializer.Serialize(lambdaPayload, _jsonOptions) |
| 134 | + }; |
| 135 | + |
| 136 | + _logger.LogInformation("Invoking Lambda function {functionName} with {recordCount} DynamoDB stream records", |
| 137 | + _config.FunctionName, lambdaRecords.Count); |
| 138 | + |
| 139 | + var lambdaResponse = await _lambdaClient.InvokeAsync(invokeRequest, _config.LambdaRuntimeApi); |
| 140 | + |
| 141 | + if (lambdaResponse.FunctionError != null) |
| 142 | + { |
| 143 | + _logger.LogError("Invoking Lambda {function} with {recordCount} DynamoDB stream records failed with error {errorMessage}", |
| 144 | + _config.FunctionName, lambdaRecords.Count, lambdaResponse.FunctionError); |
| 145 | + } |
| 146 | + } |
| 147 | + |
| 148 | + // Check for new shards periodically |
| 149 | + if (shardIterators.All(s => s == null)) |
| 150 | + { |
| 151 | + shardIterators = await GetShardIterators(streamArn, stoppingToken); |
| 152 | + if (shardIterators.Count == 0) |
| 153 | + { |
| 154 | + await Task.Delay(1000, stoppingToken); |
| 155 | + } |
| 156 | + } |
| 157 | + else if (!hasRecords) |
| 158 | + { |
| 159 | + // No records found, wait before polling again |
| 160 | + await Task.Delay(1000, stoppingToken); |
| 161 | + } |
| 162 | + } |
| 163 | + } |
| 164 | + |
| 165 | + private async Task<List<string?>> GetShardIterators(string streamArn, CancellationToken stoppingToken) |
| 166 | + { |
| 167 | + var describeResponse = await _streamsClient.DescribeStreamAsync(new DescribeStreamRequest |
| 168 | + { |
| 169 | + StreamArn = streamArn |
| 170 | + }, stoppingToken); |
| 171 | + |
| 172 | + var iterators = new List<string?>(); |
| 173 | + |
| 174 | + foreach (var shard in describeResponse.StreamDescription.Shards) |
| 175 | + { |
| 176 | + var iteratorResponse = await _streamsClient.GetShardIteratorAsync(new GetShardIteratorRequest |
| 177 | + { |
| 178 | + StreamArn = streamArn, |
| 179 | + ShardId = shard.ShardId, |
| 180 | + ShardIteratorType = ShardIteratorType.LATEST |
| 181 | + }, stoppingToken); |
| 182 | + |
| 183 | + iterators.Add(iteratorResponse.ShardIterator); |
| 184 | + } |
| 185 | + |
| 186 | + return iterators; |
| 187 | + } |
| 188 | + |
| 189 | + /// <summary> |
| 190 | + /// Convert from the SDK's DynamoDB Streams records to the Lambda event's DynamoDB record type. |
| 191 | + /// </summary> |
| 192 | + internal static IList<DynamoDBEvent.DynamodbStreamRecord> ConvertToLambdaRecords(List<Record> records, string streamArn) |
| 193 | + { |
| 194 | + return records.Select(r => ConvertToLambdaRecord(r, streamArn)).ToList(); |
| 195 | + } |
| 196 | + |
| 197 | + /// <summary> |
| 198 | + /// Convert a single SDK stream record to the Lambda event record type. |
| 199 | + /// </summary> |
| 200 | + internal static DynamoDBEvent.DynamodbStreamRecord ConvertToLambdaRecord(Record record, string streamArn) |
| 201 | + { |
| 202 | + var lambdaRecord = new DynamoDBEvent.DynamodbStreamRecord |
| 203 | + { |
| 204 | + EventID = record.EventID, |
| 205 | + EventName = record.EventName?.Value, |
| 206 | + EventSource = "aws:dynamodb", |
| 207 | + EventSourceArn = streamArn, |
| 208 | + EventVersion = record.EventVersion, |
| 209 | + AwsRegion = Arn.Parse(streamArn).Region |
| 210 | + }; |
| 211 | + |
| 212 | + if (record.Dynamodb != null) |
| 213 | + { |
| 214 | + lambdaRecord.Dynamodb = new DynamoDBEvent.StreamRecord |
| 215 | + { |
| 216 | + ApproximateCreationDateTime = record.Dynamodb.ApproximateCreationDateTime ?? DateTime.MinValue, |
| 217 | + SequenceNumber = record.Dynamodb.SequenceNumber, |
| 218 | + SizeBytes = record.Dynamodb.SizeBytes ?? 0, |
| 219 | + StreamViewType = record.Dynamodb.StreamViewType?.Value |
| 220 | + }; |
| 221 | + |
| 222 | + if (record.Dynamodb.Keys != null) |
| 223 | + { |
| 224 | + lambdaRecord.Dynamodb.Keys = ConvertAttributeMap(record.Dynamodb.Keys); |
| 225 | + } |
| 226 | + |
| 227 | + if (record.Dynamodb.NewImage != null) |
| 228 | + { |
| 229 | + lambdaRecord.Dynamodb.NewImage = ConvertAttributeMap(record.Dynamodb.NewImage); |
| 230 | + } |
| 231 | + |
| 232 | + if (record.Dynamodb.OldImage != null) |
| 233 | + { |
| 234 | + lambdaRecord.Dynamodb.OldImage = ConvertAttributeMap(record.Dynamodb.OldImage); |
| 235 | + } |
| 236 | + } |
| 237 | + |
| 238 | + if (record.UserIdentity != null) |
| 239 | + { |
| 240 | + lambdaRecord.UserIdentity = new DynamoDBEvent.Identity |
| 241 | + { |
| 242 | + PrincipalId = record.UserIdentity.PrincipalId, |
| 243 | + Type = record.UserIdentity.Type |
| 244 | + }; |
| 245 | + } |
| 246 | + |
| 247 | + return lambdaRecord; |
| 248 | + } |
| 249 | + |
| 250 | + /// <summary> |
| 251 | + /// Convert SDK AttributeValue dictionary to Lambda event AttributeValue dictionary. |
| 252 | + /// </summary> |
| 253 | + internal static Dictionary<string, DynamoDBEvent.AttributeValue> ConvertAttributeMap(Dictionary<string, AttributeValue> sdkMap) |
| 254 | + { |
| 255 | + var result = new Dictionary<string, DynamoDBEvent.AttributeValue>(); |
| 256 | + foreach (var kvp in sdkMap) |
| 257 | + { |
| 258 | + result[kvp.Key] = ConvertAttributeValue(kvp.Value); |
| 259 | + } |
| 260 | + return result; |
| 261 | + } |
| 262 | + |
| 263 | + /// <summary> |
| 264 | + /// Convert a single SDK AttributeValue to the Lambda event AttributeValue. |
| 265 | + /// </summary> |
| 266 | + internal static DynamoDBEvent.AttributeValue ConvertAttributeValue(AttributeValue sdkValue) |
| 267 | + { |
| 268 | + var lambdaValue = new DynamoDBEvent.AttributeValue(); |
| 269 | + |
| 270 | + if (sdkValue.S != null) |
| 271 | + lambdaValue.S = sdkValue.S; |
| 272 | + if (sdkValue.N != null) |
| 273 | + lambdaValue.N = sdkValue.N; |
| 274 | + if (sdkValue.B != null) |
| 275 | + lambdaValue.B = sdkValue.B; |
| 276 | + if (sdkValue.BOOL != null) |
| 277 | + lambdaValue.BOOL = sdkValue.BOOL; |
| 278 | + if (sdkValue.NULL != null) |
| 279 | + lambdaValue.NULL = sdkValue.NULL; |
| 280 | + if (sdkValue.SS?.Count > 0) |
| 281 | + lambdaValue.SS = sdkValue.SS; |
| 282 | + if (sdkValue.NS?.Count > 0) |
| 283 | + lambdaValue.NS = sdkValue.NS; |
| 284 | + if (sdkValue.BS?.Count > 0) |
| 285 | + lambdaValue.BS = sdkValue.BS; |
| 286 | + if (sdkValue.L?.Count > 0) |
| 287 | + lambdaValue.L = sdkValue.L.Select(ConvertAttributeValue).ToList(); |
| 288 | + if (sdkValue.M?.Count > 0) |
| 289 | + lambdaValue.M = ConvertAttributeMap(sdkValue.M); |
| 290 | + |
| 291 | + return lambdaValue; |
| 292 | + } |
| 293 | +} |
0 commit comments