Merge pull request #137 from dsouzajude/configure-kinesis-subscription

Configure and update kinesis subscription

Author: jarosser06

Diff for: 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 for: lambda_uploader/subscribers.py

@@ -15,35 +15,60 @@
 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']
             if response_code == 'ResourceConflictException':
-                LOG.debug('Subscription exists')
+                LOG.debug('Subscription exists. Updating ...')
+                resp = self._lambda_client\
+                           .list_event_source_mappings(
+                                FunctionName=self.function_name,
+                                EventSourceArn=self.stream_name)
+                uuid = resp['EventSourceMappings'][0]['UUID']
+                self._lambda_client \
+                    .update_event_source_mapping(
+                        UUID=uuid,
+                        FunctionName=self.function_name,
+                        Enabled=True,
+                        BatchSize=self.batch_size)
             else:
                 LOG.error('Subscription failed, error=%s' % str(ex))
                 raise ex
@@ -56,6 +81,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 for: 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 for: 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 for: 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 for: 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 for: 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

Diff for: tests/test_subscribers.py

@@ -1,6 +1,7 @@
 from os import path
 from mock import patch, Mock
 import nose.tools as nt
+import botocore
 
 from lambda_uploader import subscribers, config
 
@@ -21,3 +22,21 @@ def test_successfully_adds_kinesis_subscription(self, mocked_session):
                              config_file=path.join(EX_CONFIG, 'lambda-with-subscription.json'))
         subscribers.create_subscriptions(conf, None)
         nt.assert_equals(True, _mocked_lambda.create_event_source_mapping.called)
+
+    @patch('lambda_uploader.subscribers.boto3.session.Session')
+    def test_successfully_updates_kinesis_subscription(self, mocked_session):
+        resonse = {"Error": {"Code": "ResourceConflictException", "Message": ""}}
+        err = botocore.exceptions.ClientError(resonse, "create_event_source_mapping")
+        _mocked_lambda = Mock()
+        _mocked_lambda.create_event_source_mapping.side_effect = err
+        _mocked_lambda.list_event_source_mappings.return_value = {
+            'EventSourceMappings': [{'UUID': 'myuuid'}]
+        }
+        _mocked_session = Mock()
+        _mocked_session.client = Mock()
+        _mocked_session.client.return_value = _mocked_lambda
+        mocked_session.return_value = _mocked_session
+        conf = config.Config(path.dirname(__file__),
+                             config_file=path.join(EX_CONFIG, 'lambda-with-subscription.json'))
+        subscribers.create_subscriptions(conf, None)
+        nt.assert_equals(True, _mocked_lambda.update_event_source_mapping.called)