Configure Kinesis subscription with StartingPosition

Jude · Jude · commit babf5fa369b3 · 2017-03-28T18:39:02.000+02:00
diff --git a/lambda_uploader/config.py b/lambda_uploader/config.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import json
+from datetime import datetime
 from os import path
 
 # Python 2/3 compatability
@@ -26,7 +27,8 @@
                    u'role': basestring, u'timeout': int, u'memory': int}
 REQUIRED_VPC_PARAMS = {u'subnets': list, u'security_groups': list}
 REQUIRED_KINESIS_SUBSCRIPTION_PARAMS = {u'stream': basestring,
-                                        u'batch_size': int}
+                                        u'batch_size': int,
+                                        u'starting_position': basestring}
 
 DEFAULT_PARAMS = {u'requirements': [], u'publish': False,
                   u'alias': None, u'alias_description': None,
@@ -136,6 +138,20 @@ def validate_kinesis():
                     raise TypeError("Batch size in Kinesis subscription must"
                                     " be greater than 0")
 
+                valid_starting_pos = ['TRIM_HORIZON', 'LATEST', 'AT_TIMESTAMP']
+                if ksub['starting_position'] not in valid_starting_pos:
+                    raise TypeError("Starting position in Kinesis"
+                                    " must be one of %s" % valid_starting_pos)
+
+                if ksub['starting_position'] == 'AT_TIMESTAMP':
+                    ts = ksub.get('starting_position_timestamp')
+                    try:
+                        datetime.strptime(ts, '%Y-%m-%dT%H:%M:%SZ')
+                    except:
+                        raise TypeError("Starting position timestamp"
+                                        " must have format "
+                                        " YYYY-mm-ddTHH:MM:SSZ")
+
         if 'kinesis' in self._config['subscription']:
             validate_kinesis()
 
diff --git a/lambda_uploader/subscribers.py b/lambda_uploader/subscribers.py
@@ -15,30 +15,44 @@
 import boto3
 import botocore
 import logging
+from datetime import datetime
 
 LOG = logging.getLogger(__name__)
 
 
 class KinesisSubscriber(object):
     ''' Invokes the lambda function on events from the Kinesis streams '''
     def __init__(self, config, profile_name,
-                 function_name, stream_name, batch_size):
+                 function_name, stream_name, batch_size,
+                 starting_position, starting_position_ts=None):
         self._aws_session = boto3.session.Session(region_name=config.region,
                                                   profile_name=profile_name)
         self._lambda_client = self._aws_session.client('lambda')
         self.function_name = function_name
         self.stream_name = stream_name
         self.batch_size = batch_size
+        self.starting_position = starting_position
+        self.starting_position_ts = starting_position_ts
 
     def subscribe(self):
         ''' Subscribes the lambda to the Kinesis stream '''
         try:
             LOG.debug('Creating Kinesis subscription')
-            self._lambda_client \
-                .create_event_source_mapping(EventSourceArn=self.stream_name,
-                                             FunctionName=self.function_name,
-                                             BatchSize=self.batch_size,
-                                             StartingPosition='TRIM_HORIZON')
+            if self.starting_position_ts:
+                self._lambda_client \
+                    .create_event_source_mapping(
+                            EventSourceArn=self.stream_name,
+                            FunctionName=self.function_name,
+                            BatchSize=self.batch_size,
+                            StartingPosition=self.starting_position,
+                            StartingPositionTimestamp=self.starting_position_ts)
+            else:
+                self._lambda_client \
+                    .create_event_source_mapping(
+                            EventSourceArn=self.stream_name,
+                            FunctionName=self.function_name,
+                            BatchSize=self.batch_size,
+                            StartingPosition=self.starting_position)
             LOG.debug('Subscription created')
         except botocore.exceptions.ClientError as ex:
             response_code = ex.response['Error']['Code']
@@ -56,6 +70,13 @@ def create_subscriptions(config, profile_name):
         function_name = config.name
         stream_name = data['stream']
         batch_size = data['batch_size']
+        starting_position = data['starting_position']
+        starting_position_ts = None
+        if starting_position == 'AT_TIMESTAMP':
+            ts = data.get('starting_position_timestamp')
+            starting_position_ts = datetime.strptime(ts, '%Y-%m-%dT%H:%M:%SZ')
         s = KinesisSubscriber(config, profile_name,
-                              function_name, stream_name, batch_size)
+                              function_name, stream_name, batch_size,
+                              starting_position,
+                              starting_position_ts=starting_position_ts)
         s.subscribe()
diff --git a/tests/configs/lambda-with-subscription.json b/tests/configs/lambda-with-subscription.json
@@ -15,7 +15,8 @@
   "subscription": {
     "kinesis": {
       "stream": "arn:aws:kinesis:eu-west-1:000000000000:stream/services",
-      "batch_size": 10
+      "batch_size": 10,
+      "starting_position": "TRIM_HORIZON"
     }
   },
   "vpc": {
diff --git a/tests/configs/lambda-with-subscription_at_ts.json b/tests/configs/lambda-with-subscription_at_ts.json
@@ -0,0 +1,31 @@
+{
+  "name": "myFunc",
+  "description": "myfunc",
+  "region": "us-east-1",
+  "handler": "function.lambda_handler",
+  "role": "arn:aws:iam::00000000000:role/lambda_basic_execution",
+  "requirements": ["Jinja2==2.8"],
+  "ignore": [
+    "circle.yml",
+    ".git",
+    "/*.pyc"
+  ],
+  "timeout": 30,
+  "memory": 512,
+  "subscription": {
+    "kinesis": {
+      "stream": "arn:aws:kinesis:eu-west-1:000000000000:stream/services",
+      "batch_size": 10,
+      "starting_position": "TRIM_HORIZON",
+      "starting_position_timestamp": "2017-11-01T11:00:00Z"
+    }
+  },
+  "vpc": {
+    "subnets": [
+      "subnet-00000000"
+    ],
+    "security_groups": [
+      "sg-00000000"
+    ]
+  }
+}
diff --git a/tests/configs/lambda-with-subscription_at_ts_invalid_ts.json b/tests/configs/lambda-with-subscription_at_ts_invalid_ts.json
@@ -0,0 +1,31 @@
+{
+  "name": "myFunc",
+  "description": "myfunc",
+  "region": "us-east-1",
+  "handler": "function.lambda_handler",
+  "role": "arn:aws:iam::00000000000:role/lambda_basic_execution",
+  "requirements": ["Jinja2==2.8"],
+  "ignore": [
+    "circle.yml",
+    ".git",
+    "/*.pyc"
+  ],
+  "timeout": 30,
+  "memory": 512,
+  "subscription": {
+    "kinesis": {
+      "stream": "arn:aws:kinesis:eu-west-1:000000000000:stream/services",
+      "batch_size": 10,
+      "starting_position": "AT_TIMESTAMP",
+      "starting_position_timestamp": null
+    }
+  },
+  "vpc": {
+    "subnets": [
+      "subnet-00000000"
+    ],
+    "security_groups": [
+      "sg-00000000"
+    ]
+  }
+}
diff --git a/tests/configs/lambda-with-subscription_at_ts_not_provided.json b/tests/configs/lambda-with-subscription_at_ts_not_provided.json
@@ -0,0 +1,30 @@
+{
+  "name": "myFunc",
+  "description": "myfunc",
+  "region": "us-east-1",
+  "handler": "function.lambda_handler",
+  "role": "arn:aws:iam::00000000000:role/lambda_basic_execution",
+  "requirements": ["Jinja2==2.8"],
+  "ignore": [
+    "circle.yml",
+    ".git",
+    "/*.pyc"
+  ],
+  "timeout": 30,
+  "memory": 512,
+  "subscription": {
+    "kinesis": {
+      "stream": "arn:aws:kinesis:eu-west-1:000000000000:stream/services",
+      "batch_size": 10,
+      "starting_position": "AT_TIMESTAMP"
+    }
+  },
+  "vpc": {
+    "subnets": [
+      "subnet-00000000"
+    ],
+    "security_groups": [
+      "sg-00000000"
+    ]
+  }
+}
diff --git a/tests/test_config.py b/tests/test_config.py
@@ -58,6 +58,28 @@ def test_kinesis_subscription():
     assert cfg.raw['subscription']['kinesis']['batch_size'] == ksub['batch_size']
 
 
+def test_kinesis_subscription_with_starting_position_at_timestamp():
+    ksub = {
+        'stream': 'arn:aws:kinesis:eu-west-1:000000000000:stream/services',
+        'batch_size': 10,
+        'starting_position_timestamp': '2017-11-01T11:00:00Z'
+    }
+    cfg = config.Config(EX_CONFIG, EX_CONFIG + '/lambda-with-subscription_at_ts.json')
+    assert cfg.raw['subscription']['kinesis']['stream'] == ksub['stream']
+    assert cfg.raw['subscription']['kinesis']['batch_size'] == ksub['batch_size']
+    assert cfg.raw['subscription']['kinesis']['starting_position_timestamp'] == ksub['starting_position_timestamp']
+
+
+def test_kinesis_subscription_with_starting_position_at_timestamp_fails_when_timestamp_is_invalid():
+    with pytest.raises(Exception):
+        cfg = config.Config(EX_CONFIG, EX_CONFIG + '/lambda-with-subscription_at_ts_invalid_ts.json')
+
+
+def test_kinesis_subscription_with_starting_position_at_timestamp_fails_when_timestamp_not_provided():
+    with pytest.raises(Exception):
+        cfg = config.Config(EX_CONFIG, EX_CONFIG + '/lambda-with-subscription_at_ts_not_provided.json')
+
+
 def test_set_publish():
     cfg = config.Config(EX_CONFIG)
     # Check that we default to false
