Add functionality for Kinesis subscriptions

Author: Jude

lambda_uploader/config.py

@@ -25,12 +25,14 @@
                    u'region': basestring, u'handler': basestring,
                    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}
 
 DEFAULT_PARAMS = {u'requirements': [], u'publish': False,
                   u'alias': None, u'alias_description': None,
                   u'ignore': [], u'extra_files': [], u'vpc': None,
                   u's3_bucket': None, u's3_key': None, u'runtime': 'python2.7',
-                  u'variables': {}}
+                  u'variables': {}, u'subscription': {}}
 
 
 class Config(object):
@@ -45,6 +47,8 @@ def __init__(self, pth, config_file=None, role=None, variables=None):
         self._set_defaults()
         if self._config['vpc']:
             self._validate_vpc()
+        if self._config['subscription']:
+            self._validate_subscription()
 
         for param, clss in REQUIRED_PARAMS.items():
             self._validate(param, cls=clss)
@@ -120,6 +124,21 @@ def _validate_vpc(self):
                                     " strings. '%s' contains something else"
                                     % param)
 
+    '''Validate the subscription configuration.
+    All kinds of subscription will be validated here'''
+    def _validate_subscription(self):
+        def validate_kinesis():
+            ksub = self._config['subscription']['kinesis']
+            for param, clss in REQUIRED_KINESIS_SUBSCRIPTION_PARAMS.items():
+                self._compare(param, clss, ksub.get(param))
+
+                if ksub['batch_size'] <= 0:
+                    raise TypeError("Batch size in Kinesis subscription must"
+                                    " be greater than 0")
+
+        if 'kinesis' in self._config['subscription']:
+            validate_kinesis()
+
     '''Compare if a string is a certain type'''
     def _compare(self, key, cls, value):
         if cls:

lambda_uploader/shell.py

@@ -23,7 +23,7 @@
 import lambda_uploader
 
 from os import getcwd, path, getenv
-from lambda_uploader import package, config, uploader
+from lambda_uploader import package, config, uploader, subscribers
 from boto3 import __version__ as boto3_version
 from botocore import __version__ as botocore_version
 
@@ -104,6 +104,10 @@ def _execute(args):
         if create_alias:
             upldr.alias()
 
+        if cfg.subscription:
+            _print('Creating subscription')
+            subscribers.create_subscriptions(cfg, args.profile)
+
         pkg.clean_zipfile()
 
     _print('Fin')

lambda_uploader/subscribers.py

@@ -0,0 +1,65 @@
+# Copyright 2015-2016 Rackspace US, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import time
+import boto3
+import botocore
+import logging
+
+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):
+        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
+
+    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')
+            LOG.debug('Subscription created')
+        except botocore.exceptions.ClientError as ex:
+            response_code = ex.response['Error']['Code']
+            if response_code == 'InvalidParameterValueException':
+                LOG.debug('Retrying subscription')
+                time.sleep(3)
+            elif response_code == 'ResourceConflictException':
+                LOG.debug('Subscription exists')
+            else:
+                LOG.error('Subscription failed, error=%s' % str(ex))
+                raise ex
+
+
+def create_subscriptions(config, profile_name):
+    ''' Adds supported subscriptions '''
+    if 'kinesis' in config.subscription.keys():
+        data = config.subscription['kinesis']
+        function_name = config.name
+        stream_name = data['stream']
+        batch_size = data['batch_size']
+        s = KinesisSubscriber(config, profile_name,
+                              function_name, stream_name, batch_size)
+        s.subscribe()

tests/configs/lambda-with-subscription.json

@@ -0,0 +1,29 @@
+{
+  "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
+    }
+  },
+  "vpc": {
+    "subnets": [
+      "subnet-00000000"
+    ],
+    "security_groups": [
+      "sg-00000000"
+    ]
+  }
+}

tests/test_config.py

@@ -48,6 +48,16 @@ def test_no_vpc():
     assert cfg.raw['vpc'] is None
 
 
+def test_kinesis_subscription():
+    ksub = {
+        'stream': 'arn:aws:kinesis:eu-west-1:000000000000:stream/services',
+        'batch_size': 10
+    }
+    cfg = config.Config(EX_CONFIG, EX_CONFIG + '/lambda-with-subscription.json')
+    assert cfg.raw['subscription']['kinesis']['stream'] == ksub['stream']
+    assert cfg.raw['subscription']['kinesis']['batch_size'] == ksub['batch_size']
+
+
 def test_set_publish():
     cfg = config.Config(EX_CONFIG)
     # Check that we default to false

tests/test_subscribers.py

@@ -0,0 +1,23 @@
+from os import path
+from mock import patch, Mock
+import nose.tools as nt
+
+from lambda_uploader import subscribers, config
+
+
+EX_CONFIG = path.normpath(path.join(path.dirname(__file__),
+                          '../tests/configs'))
+
+class TestKinesisSubscriber(object):
+
+    @patch('lambda_uploader.subscribers.boto3.session.Session')
+    def test_successfully_adds_kinesis_subscription(self, mocked_session):
+        _mocked_lambda = Mock()
+        _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.create_event_source_mapping.called)