124 lines
3.5 KiB
Go
124 lines
3.5 KiB
Go
package consumergroup
|
|
|
|
import (
|
|
"github.com/aws/aws-sdk-go/aws"
|
|
"github.com/aws/aws-sdk-go/service/dynamodb"
|
|
"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbattribute"
|
|
|
|
consumer "github.com/harlow/kinesis-consumer"
|
|
)
|
|
|
|
// DynamoDb simple and minimal interface for DynamoDb that helps with testing
|
|
type DynamoDb interface {
|
|
PutItem(*dynamodb.PutItemInput) (*dynamodb.PutItemOutput, error)
|
|
UpdateItem(*dynamodb.UpdateItemInput) (*dynamodb.UpdateItemOutput, error)
|
|
GetItem(*dynamodb.GetItemInput) (*dynamodb.GetItemOutput, error)
|
|
Scan(*dynamodb.ScanInput) (*dynamodb.ScanOutput, error)
|
|
}
|
|
|
|
// DynamoStorage struct that implements the storage interface and uses simplified DynamoDb struct
|
|
type DynamoStorage struct {
|
|
Db DynamoDb
|
|
tableName string
|
|
}
|
|
|
|
// CreateLease - stores the lease in dynamo
|
|
func (dynamoClient DynamoStorage) CreateLease(lease consumer.Lease) error {
|
|
|
|
av, err := dynamodbattribute.MarshalMap(lease)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
//TODO add conditional expression
|
|
input := &dynamodb.PutItemInput{
|
|
Item: av,
|
|
TableName: aws.String(dynamoClient.tableName),
|
|
}
|
|
|
|
if _, err := dynamoClient.Db.PutItem(input); err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// TODO add conditional expressions
|
|
// UpdateLease - updates the lease in dynamo
|
|
func (dynamoClient DynamoStorage) UpdateLease(leaseKey string, leaseUpdate consumer.LeaseUpdate) error {
|
|
|
|
key := mapShardIdToKey(leaseKey)
|
|
update, err := dynamodbattribute.MarshalMap(leaseUpdate)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
input := &dynamodb.UpdateItemInput{
|
|
ExpressionAttributeValues: update,
|
|
Key: key,
|
|
ReturnValues: aws.String("UPDATED_NEW"),
|
|
TableName: aws.String(dynamoClient.tableName),
|
|
UpdateExpression: aws.String("set checkpoint = :cp, leaseCounter= :lc, leaseOwner= :lo, heartbeatID= :hb"),
|
|
}
|
|
|
|
if _, err := dynamoClient.Db.UpdateItem(input); err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// GetLease returns the latest stored records sorted by clockID in descending order
|
|
// It is assumed that we won't be keeping many records per ID otherwise, this may need to be optimized
|
|
// later (possibly to use a map)
|
|
func (dynamoClient DynamoStorage) GetLease(shardID string) (*consumer.Lease, error) {
|
|
|
|
key := mapShardIdToKey(shardID)
|
|
input := &dynamodb.GetItemInput{
|
|
Key: key,
|
|
TableName: aws.String(dynamoClient.tableName),
|
|
ConsistentRead: aws.Bool(true),
|
|
}
|
|
result, err := dynamoClient.Db.GetItem(input)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var lease consumer.Lease
|
|
if err := dynamodbattribute.UnmarshalMap(result.Item, &lease); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &lease, nil
|
|
}
|
|
|
|
func mapShardIdToKey(shardID string) map[string]*dynamodb.AttributeValue {
|
|
return map[string]*dynamodb.AttributeValue{
|
|
"leaseKey": {S: aws.String(shardID)},
|
|
}
|
|
}
|
|
|
|
// GetAllLeases this can be used at start up (or anytime to grab all the leases)
|
|
func (dynamoClient DynamoStorage) GetAllLeases() (map[string]consumer.Lease, error) {
|
|
|
|
// TODO if we have a lot of shards, we might have to worry about limits here
|
|
input := &dynamodb.ScanInput{
|
|
ConsistentRead: aws.Bool(true),
|
|
TableName: aws.String(dynamoClient.tableName),
|
|
}
|
|
result, err := dynamoClient.Db.Scan(input)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
leases := make(map[string]consumer.Lease, len(result.Items))
|
|
for _, item := range result.Items {
|
|
var record consumer.Lease
|
|
if err := dynamodbattribute.UnmarshalMap(item, &record); err != nil {
|
|
return nil, err
|
|
}
|
|
leases[record.LeaseKey] = record
|
|
}
|
|
|
|
return leases, nil
|
|
}
|