Merge pull request #22 from snowplow-devops/add-merge-upstream

Add merge upstream

Author: colmsnowplow

config.go

@@ -49,6 +49,9 @@ type Config struct {
 	dynamoWriteCapacity int64
 	// Time to wait between attempts to verify tables were created/deleted completely
 	dynamoWaiterDelay time.Duration
+
+	// use ListShards to avoid LimitExceedException from DescribeStream
+	useListShardsForKinesisStreamReady bool
 }
 
 // NewConfig returns a default Config struct
@@ -148,6 +151,12 @@ func (c Config) WithLogger(logger Logger) Config {
 	return c
 }
 
+// WithUseListShardsForKinesisStreamReady returns a config with a modified useListShardsForKinesisStreamReady toggle
+func (c Config) WithUseListShardsForKinesisStreamReady(shouldUse bool) Config {
+	c.useListShardsForKinesisStreamReady = shouldUse
+	return c
+}
+
 // Verify that a config struct has sane and valid values
 func validateConfig(c *Config) error {
 	if c.throttleDelay < 200*time.Millisecond {

config_test.go

@@ -64,7 +64,8 @@ func TestConfigWithMethods(t *testing.T) {
 		WithLeaderActionFrequency(1 * time.Second).
 		WithThrottleDelay(1 * time.Second).
 		WithStats(stats).
-		WithIteratorStartTimestamp(&tstamp)
+		WithIteratorStartTimestamp(&tstamp).
+		WithUseListShardsForKinesisStreamReady(true)
 
 	err := validateConfig(&config)
 	require.NoError(t, err)
@@ -76,4 +77,5 @@ func TestConfigWithMethods(t *testing.T) {
 	require.Equal(t, 1*time.Second, config.leaderActionFrequency)
 	require.Equal(t, stats, config.stats)
 	require.Equal(t, &tstamp, config.iteratorStartTimestamp)
+	require.Equal(t, true, config.useListShardsForKinesisStreamReady)
 }

kinsumer.go

@@ -325,6 +325,15 @@ func (k *Kinsumer) dynamoDeleteTableIfExists(name string) error {
 
 // kinesisStreamReady returns an error if the given stream is not ACTIVE
 func (k *Kinsumer) kinesisStreamReady() error {
+	if k.config.useListShardsForKinesisStreamReady {
+		_, err := k.kinesis.ListShards(&kinesis.ListShardsInput{
+			StreamName: aws.String(k.streamName),
+		})
+		if err != nil {
+			return fmt.Errorf("error listing shards for stream %s: %v", k.streamName, err)
+		}
+		return nil
+	}
 	out, err := k.kinesis.DescribeStream(&kinesis.DescribeStreamInput{
 		StreamName: aws.String(k.streamName),
 	})
@@ -336,7 +345,6 @@ func (k *Kinsumer) kinesisStreamReady() error {
 	if status != "ACTIVE" && status != "UPDATING" {
 		return fmt.Errorf("stream %s exists but state '%s' is not 'ACTIVE' or 'UPDATING'", k.streamName, status)
 	}
-
 	return nil
 }