@@ -2,14 +2,15 @@ package lock
2
2
3
3
import (
4
4
"fmt"
5
+ "time"
6
+
5
7
"github.com/aws/aws-sdk-go/aws"
6
8
"github.com/aws/aws-sdk-go/aws/awserr"
7
9
"github.com/aws/aws-sdk-go/aws/session"
8
10
"github.com/aws/aws-sdk-go/service/dynamodb"
9
11
"github.com/gruntwork-io/go-commons/errors"
10
12
"github.com/gruntwork-io/go-commons/retry"
11
13
"github.com/sirupsen/logrus"
12
- "time"
13
14
)
14
15
15
16
const (
@@ -40,6 +41,10 @@ type Options struct {
40
41
SleepBetweenRetries time.Duration
41
42
// The logger to use for the lock
42
43
Logger * logrus.Logger
44
+
45
+ // Custom session to use to authenticate to AWS in the SDK. If nil, constructs the session based on the default
46
+ // authentication chain in the SDK.
47
+ AwsSession * session.Session
43
48
}
44
49
45
50
type TimeoutExceeded struct {
@@ -79,23 +84,13 @@ func NewAuthenticatedSession(awsRegion string) (*session.Session, error) {
79
84
return sess , nil
80
85
}
81
86
82
- // NewDynamoDb returns an authenticated client object for accessing DynamoDb
83
- func NewDynamoDb (awsRegion string ) (* dynamodb.DynamoDB , error ) {
84
- sess , err := NewAuthenticatedSession (awsRegion )
85
- if err != nil {
86
- return nil , err
87
- }
88
- dynamodbSvc := dynamodb .New (sess )
89
- return dynamodbSvc , nil
90
- }
91
-
92
87
// AcquireLock will attempt to acquire a lock in DynamoDB table while taking the configuration options into account.
93
88
// We are using DynamoDB to create a table to help us track the lock status of different resources.
94
89
// The acquiring of a lock attempts to create a table. The intention is that we have 1 table per resource in a single region.
95
90
// This would allow the locking mechanism to flexibly decide if a resource is locked or not. For test cases where the AWS resource
96
91
// is multi-region, or global, the configuration of which regions to use should reflect that.
97
92
func AcquireLock (options * Options ) error {
98
- client , err := NewDynamoDb (options . AwsRegion )
93
+ client , err := getDynamoDBClient (options )
99
94
if err != nil {
100
95
options .Logger .Errorf ("Error authenticating to AWS: %s\n " , err )
101
96
return err
@@ -156,7 +151,7 @@ func acquireLock(options *Options, client *dynamodb.DynamoDB) error {
156
151
// ReleaseLock will attempt to release the lock defined by the provided lock string in the configured lock table for the
157
152
// configured region
158
153
func ReleaseLock (options * Options ) error {
159
- client , err := NewDynamoDb (options . AwsRegion )
154
+ client , err := getDynamoDBClient (options )
160
155
if err != nil {
161
156
options .Logger .Errorf ("Error authenticating to AWS: %s\n " , err )
162
157
return err
@@ -297,7 +292,7 @@ func lockTableExistsAndIsActive(tableName string, client *dynamodb.DynamoDB) (bo
297
292
// GetLockStatus attempts to acquire the lock and check if the expected item is there
298
293
// If there's the expected Item with the correct `LockString` value - then the status is `locked`, if the item is not there - then the status is `not locked`
299
294
func GetLockStatus (options * Options ) (* dynamodb.GetItemOutput , error ) {
300
- client , err := NewDynamoDb (options . AwsRegion )
295
+ client , err := getDynamoDBClient (options )
301
296
if err != nil {
302
297
options .Logger .Errorf ("Error authenticating to AWS: %s\n " , err )
303
298
return nil , err
@@ -318,3 +313,20 @@ func GetLockStatus(options *Options) (*dynamodb.GetItemOutput, error) {
318
313
319
314
return item , nil
320
315
}
316
+
317
+ func getDynamoDBClient (options * Options ) (* dynamodb.DynamoDB , error ) {
318
+ if options .AwsSession == nil {
319
+ return NewDynamoDb (options .AwsRegion )
320
+ }
321
+ return dynamodb .New (options .AwsSession ), nil
322
+ }
323
+
324
+ // NewDynamoDb returns an authenticated client object for accessing DynamoDb
325
+ func NewDynamoDb (awsRegion string ) (* dynamodb.DynamoDB , error ) {
326
+ sess , err := NewAuthenticatedSession (awsRegion )
327
+ if err != nil {
328
+ return nil , err
329
+ }
330
+ dynamodbSvc := dynamodb .New (sess )
331
+ return dynamodbSvc , nil
332
+ }
0 commit comments