state.dynamodb: validate AWS connection

state/aws/dynamodb/dynamodb.go

@@ -19,10 +19,11 @@
 	"encoding/binary"
 	"errors"
 	"fmt"
+	"github.com/google/uuid"
 	"reflect"
 	"strconv"
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/dynamodb"
 	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbattribute"
@@ -74,25 +75,48 @@
 }
 
 // Init does metadata and connection parsing.
-func (d *StateStore) Init(_ context.Context, metadata state.Metadata) error {
+func (d *StateStore) Init(ctx context.Context, metadata state.Metadata) error {
 	meta, err := d.getDynamoDBMetadata(metadata)
 	if err != nil {
 		return err
 	}
 
-	client, err := d.getClient(meta)
-	if err != nil {
-		return err
+	// We have this check because we need to set the client to  a mock in tests
+	if d.client == nil {
+		client, err := d.getClient(meta)
+		if err != nil {
+			return err
+		}
+		d.client = client
 	}
-
-	d.client = client
 	d.table = meta.Table
 	d.ttlAttributeName = meta.TTLAttributeName
 	d.partitionKey = meta.PartitionKey
 
+	if err := d.validateTableAccess(ctx); err != nil {
+		return fmt.Errorf("error validating DynamoDB table '%s' access: %w", d.table, err)
+	}
+
 	return nil
 }
 
+// validateConnection runs a dummy Get operation to validate the connection credentials,
+// as well as validating that the table exists, and we have access to it
+func (d *StateStore) validateTableAccess(ctx context.Context) error {
+	input := &dynamodb.GetItemInput{
+		ConsistentRead: aws.Bool(false),
+		TableName:      aws.String(d.table),
+		Key: map[string]*dynamodb.AttributeValue{
+			d.partitionKey: {
+				S: aws.String(uuid.NewString()),
+			},
+		},
+	}
+
+	_, err := d.client.GetItemWithContext(ctx, input)
+	return err
+}
+
 // Features returns the features available in this state store.
 func (d *StateStore) Features() []state.Feature {
 	// TTLs are enabled only if ttlAttributeName is set

state/aws/dynamodb/dynamodb_test.go

@@ -17,6 +17,7 @@ package dynamodb
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"testing"
 	"time"
@@ -76,6 +77,12 @@ func TestInit(t *testing.T) {
 	m := state.Metadata{}
 	s := &StateStore{
 		partitionKey: defaultPartitionKeyName,
+		client: &mockedDynamoDB{
+			// We're adding this so we can pass the connection check on Init
+			GetItemWithContextFn: func(ctx context.Context, input *dynamodb.GetItemInput, op ...request.Option) (*dynamodb.GetItemOutput, error) {
+				return nil, nil
+			},
+		},
 	}
 
 	t.Run("NewDynamoDBStateStore Default Partition Key", func(t *testing.T) {
@@ -124,6 +131,23 @@ func TestInit(t *testing.T) {
 		require.NoError(t, err)
 		assert.Equal(t, s.partitionKey, pkey)
 	})
+
+	t.Run("Init with bad table name or permissions", func(t *testing.T) {
+		m.Properties = map[string]string{
+			"Table":  "does-not-exist",
+			"Region": "eu-west-1",
+		}
+
+		s.client = &mockedDynamoDB{
+			GetItemWithContextFn: func(ctx context.Context, input *dynamodb.GetItemInput, op ...request.Option) (*dynamodb.GetItemOutput, error) {
+				return nil, errors.New("Requested resource not found")
+			},
+		}
+
+		err := s.Init(context.Background(), m)
+		require.Error(t, err)
+		require.Equal(t, "error validating DynamoDB table 'does-not-exist' access: Requested resource not found", err.Error())
+	})
 }
 
 func TestGet(t *testing.T) {