Repository: Miserlou/Zappa

Issue:
Way to limit concurrent executions?
<!--- Provide a general summary of the issue in the Title above -->
## Context
Recovering from rouge script issue and trying to see if there was a mechanism within Zappa deployment to set concurrent executions on a per function basis.

Python 2.7/3.6 -->
2.7

## Expected Behavior
A setting variable within the settings_file to able to add/limit concurrent execution limit on the deployed functions.

## Actual Behavior
<!--- Tell us what happens instead -->

## Possible Fix
<!--- Not obligatory, but suggest a fix or reason for the bug -->


## Your Environment
<!--- Include as many relevant details about the environment you experienced the bug in -->
* Zappa version used: 0.45.1
* Operating System and Python version: 




Discussion:
You can do it through the aws. Have a look at https://aws.amazon.com/blogs/compute/managing-aws-lambda-function-concurrency/
It would be great to be able to control the concurrency limit via the Zappa settings file as well :D 
+1 I'm also waiting for this great feature!
Related to #1628

Resolving patch:
diff --git a/zappa/cli.py b/zappa/cli.py
--- a/zappa/cli.py
+++ b/zappa/cli.py
@@ -100,6 +100,7 @@ class ZappaCLI(object):
     lambda_arn = None
     lambda_name = None
     lambda_description = None
+    lambda_concurrency = None
     s3_bucket_name = None
     settings_file = None
     zip_path = None
@@ -762,7 +763,8 @@ def deploy(self, source_zip=None):
                 runtime=self.runtime,
                 aws_environment_variables=self.aws_environment_variables,
                 aws_kms_key_arn=self.aws_kms_key_arn,
-                use_alb=self.use_alb
+                use_alb=self.use_alb,
+                concurrency=self.lambda_concurrency,
             )
             if source_zip and source_zip.startswith('s3://'):
                 bucket, key_name = parse_s3_url(source_zip)
@@ -936,7 +938,8 @@ def update(self, source_zip=None, no_upload=False):
         kwargs = dict(
             bucket=self.s3_bucket_name,
             function_name=self.lambda_name,
-            num_revisions=self.num_retained_versions
+            num_revisions=self.num_retained_versions,
+            concurrency=self.lambda_concurrency,
         )
         if source_zip and source_zip.startswith('s3://'):
             bucket, key_name = parse_s3_url(source_zip)
@@ -970,7 +973,7 @@ def update(self, source_zip=None, no_upload=False):
                                                         memory_size=self.memory_size,
                                                         runtime=self.runtime,
                                                         aws_environment_variables=self.aws_environment_variables,
-                                                        aws_kms_key_arn=self.aws_kms_key_arn
+                                                        aws_kms_key_arn=self.aws_kms_key_arn,
                                                     )
 
         # Finally, delete the local copy our zip package
@@ -2087,6 +2090,7 @@ def load_settings(self, settings_file=None, session=None):
         self.iam_authorization = self.stage_config.get('iam_authorization', False)
         self.cors = self.stage_config.get("cors", False)
         self.lambda_description = self.stage_config.get('lambda_description', "Zappa Deployment")
+        self.lambda_concurrency = self.stage_config.get('lambda_concurrency', None)
         self.environment_variables = self.stage_config.get('environment_variables', {})
         self.aws_environment_variables = self.stage_config.get('aws_environment_variables', {})
         self.check_environment(self.environment_variables)
diff --git a/zappa/core.py b/zappa/core.py
--- a/zappa/core.py
+++ b/zappa/core.py
@@ -1052,6 +1052,7 @@ def create_lambda_function( self,
                                 xray_tracing=False,
                                 local_zip=None,
                                 use_alb=False,
+                                concurrency=None,
                             ):
         """
         Given a bucket and key (or a local path) of a valid Lambda-zip, a function name and a handler, register that Lambda function.
@@ -1095,7 +1096,6 @@ def create_lambda_function( self,
             }
 
         response = self.lambda_client.create_function(**kwargs)
-
         resource_arn = response['FunctionArn']
         version = response['Version']
 
@@ -1114,9 +1114,15 @@ def create_lambda_function( self,
         if self.tags:
             self.lambda_client.tag_resource(Resource=resource_arn, Tags=self.tags)
 
+        if concurrency is not None:
+            self.lambda_client.put_function_concurrency(
+                FunctionName=resource_arn,
+                ReservedConcurrentExecutions=concurrency,
+            )
+
         return resource_arn
 
-    def update_lambda_function(self, bucket, function_name, s3_key=None, publish=True, local_zip=None, num_revisions=None):
+    def update_lambda_function(self, bucket, function_name, s3_key=None, publish=True, local_zip=None, num_revisions=None, concurrency=None):
         """
         Given a bucket and key (or a local path) of a valid Lambda-zip, a function name and a handler, update that Lambda function's code.
         Optionally, delete previous versions if they exceed the optional limit.
@@ -1161,6 +1167,16 @@ def update_lambda_function(self, bucket, function_name, s3_key=None, publish=Tru
                 Name=ALB_LAMBDA_ALIAS,
             )
 
+        if concurrency is not None:
+            self.lambda_client.put_function_concurrency(
+                FunctionName=function_name,
+                ReservedConcurrentExecutions=concurrency,
+            )
+        else:
+            self.lambda_client.delete_function_concurrency(
+                FunctionName=function_name
+            )
+
         if num_revisions:
             # Find the existing revision IDs for the given function
             # Related: https://github.com/Miserlou/Zappa/issues/1402


Test patch:
diff --git a/test_settings.json b/test_settings.json
--- a/test_settings.json
+++ b/test_settings.json
@@ -116,5 +116,9 @@
       "delete_local_zip": true,
       "use_precompiled_packages": false,
       "delete_s3_zip": false
-    }
+    },
+    "lambda_concurrency_enabled": {
+        "extends": "ttt888",
+        "lambda_concurrency": 6
+     }
 }
diff --git a/tests/placebo/TestZappa.test_cli_aws/lambda.DeleteFunctionConcurrency_1.json b/tests/placebo/TestZappa.test_cli_aws/lambda.DeleteFunctionConcurrency_1.json
new file mode 100644
--- /dev/null
+++ b/tests/placebo/TestZappa.test_cli_aws/lambda.DeleteFunctionConcurrency_1.json
@@ -0,0 +1,9 @@
+{
+    "status_code": 200, 
+    "data": {
+        "ResponseMetadata": {
+            "HTTPStatusCode": 200, 
+            "RequestId": "aff3a3f9-28f4-11e6-9dbb-5dd116b9ddf1"
+        }
+    }
+}
diff --git a/tests/placebo/TestZappa.test_create_lambda_function_local/lambda.DeleteFunctionConcurrency_1.json b/tests/placebo/TestZappa.test_create_lambda_function_local/lambda.DeleteFunctionConcurrency_1.json
new file mode 100644
--- /dev/null
+++ b/tests/placebo/TestZappa.test_create_lambda_function_local/lambda.DeleteFunctionConcurrency_1.json
@@ -0,0 +1,9 @@
+{
+    "status_code": 200, 
+    "data": {
+        "ResponseMetadata": {
+            "HTTPStatusCode": 200, 
+            "RequestId": "aff3a3f9-28f4-11e6-9dbb-5dd116b9ddf1"
+        }
+    }
+}
diff --git a/tests/placebo/TestZappa.test_create_lambda_function_s3/lambda.DeleteFunctionConcurrency_1.json b/tests/placebo/TestZappa.test_create_lambda_function_s3/lambda.DeleteFunctionConcurrency_1.json
new file mode 100644
--- /dev/null
+++ b/tests/placebo/TestZappa.test_create_lambda_function_s3/lambda.DeleteFunctionConcurrency_1.json
@@ -0,0 +1,9 @@
+{
+    "status_code": 200, 
+    "data": {
+        "ResponseMetadata": {
+            "HTTPStatusCode": 200, 
+            "RequestId": "aff3a3f9-28f4-11e6-9dbb-5dd116b9ddf1"
+        }
+    }
+}
diff --git a/tests/tests.py b/tests/tests.py
--- a/tests/tests.py
+++ b/tests/tests.py
@@ -753,6 +753,12 @@ def test_load_extended_settings(self):
         self.assertTrue(zappa_cli.stage_config['touch'])  # First Extension
         self.assertTrue(zappa_cli.stage_config['delete_local_zip'])  # The base
 
+    def test_load_settings__lambda_concurrency_enabled(self):
+        zappa_cli = ZappaCLI()
+        zappa_cli.api_stage = 'lambda_concurrency_enabled'
+        zappa_cli.load_settings('test_settings.json')
+        self.assertEqual(6, zappa_cli.stage_config['lambda_concurrency'])
+
     def test_load_settings_yml(self):
         zappa_cli = ZappaCLI()
         zappa_cli.api_stage = 'ttt888'
@@ -2171,5 +2177,72 @@ def test_zappa_core_undeploy_lambda_alb(self):
         zappa_core.undeploy_lambda_alb(**kwargs)
 
 
+    @mock.patch('botocore.client')
+    def test_set_lambda_concurrency(self, client):
+        boto_mock = mock.MagicMock()
+        zappa_core = Zappa(
+            boto_session=boto_mock,
+            profile_name="test",
+            aws_region="test",
+            load_credentials=True
+        )
+        zappa_core.lambda_client.create_function.return_value = {
+            "FunctionArn": "abc",
+            "Version": 1,
+        }
+        access_logging_patch = zappa_core.create_lambda_function(
+            concurrency=5,
+        )
+        boto_mock.client().put_function_concurrency.assert_called_with(
+            FunctionName="abc",
+            ReservedConcurrentExecutions=5,
+        )
+
+    @mock.patch('botocore.client')
+    def test_update_lambda_concurrency(self, client):
+        boto_mock = mock.MagicMock()
+        zappa_core = Zappa(
+            boto_session=boto_mock,
+            profile_name="test",
+            aws_region="test",
+            load_credentials=True
+        )
+        zappa_core.lambda_client.create_function.return_value = {
+            "FunctionArn": "abc",
+            "Version": 1,
+        }
+        access_logging_patch = zappa_core.update_lambda_function(
+            bucket="test",
+            function_name="abc",
+            concurrency=5,
+        )
+        boto_mock.client().put_function_concurrency.assert_called_with(
+            FunctionName="abc",
+            ReservedConcurrentExecutions=5,
+        )
+        boto_mock.client().delete_function_concurrency.assert_not_called()
+
+    @mock.patch('botocore.client')
+    def test_delete_lambda_concurrency(self, client):
+        boto_mock = mock.MagicMock()
+        zappa_core = Zappa(
+            boto_session=boto_mock,
+            profile_name="test",
+            aws_region="test",
+            load_credentials=True
+        )
+        zappa_core.lambda_client.create_function.return_value = {
+            "FunctionArn": "abc",
+            "Version": 1,
+        }
+        access_logging_patch = zappa_core.update_lambda_function(
+            bucket="test",
+            function_name="abc",
+        )
+        boto_mock.client().put_function_concurrency.assert_not_called()
+        boto_mock.client().delete_function_concurrency.assert_called_with(
+            FunctionName="abc",
+        )
+
 if __name__ == '__main__':
     unittest.main()