quiqua · quiqua · Jul 6, 2016 · Jul 6, 2016
diff --git a/aws_lambda/__init__.py b/aws_lambda/__init__.py
@@ -4,7 +4,7 @@
 __email__ = 'nficano@gmail.com'
 __version__ = '0.4.0'
 
-from .aws_lambda import deploy, invoke, init, build
+from .aws_lambda import deploy, invoke, init, build, cleanup_old_versions
 
 # Set default logging handler to avoid "No handler found" warnings.
 import logging

diff --git a/aws_lambda/aws_lambda.py b/aws_lambda/aws_lambda.py
@@ -8,6 +8,7 @@
 from shutil import copy, copyfile
 from tempfile import mkdtemp
 
+import botocore
 import boto3
 import pip
 import yaml
@@ -18,6 +19,49 @@
 log = logging.getLogger(__name__)
 
 
+def cleanup_old_versions(src, keep_last_versions):
+    """Deletes old deployed versions of the function in AWS Lambda.
+
+    Won't delete $Latest and any aliased version
+
+    :param str src:
+        The path to your Lambda ready project (folder must contain a valid
+        config.yaml and handler module (e.g.: service.py).
+    :param int keep_last_versions:
+        The number of recent versions to keep and not delete
+    """
+    if keep_last_versions <= 0:
+        print("Won't delete all versions. Please do this manually")
+    else:
+        path_to_config_file = os.path.join(src, 'config.yaml')
+        cfg = read(path_to_config_file, loader=yaml.load)
+
+        aws_access_key_id = cfg.get('aws_access_key_id')
+        aws_secret_access_key = cfg.get('aws_secret_access_key')
+
+        client = get_client('lambda', aws_access_key_id, aws_secret_access_key,
+                            cfg.get('region'))
+
+        response = client.list_versions_by_function(
+            FunctionName=cfg.get("function_name")
+        )
+        versions = response.get("Versions")
+        if len(response.get("Versions")) < keep_last_versions:
+            print("Nothing to delete. (Too few versions published)")
+        else:
+            version_numbers = [elem.get("Version") for elem in
+                               versions[1:-keep_last_versions]]
+            for version_number in version_numbers:
+                try:
+                    client.delete_function(
+                        FunctionName=cfg.get("function_name"),
+                        Qualifier=version_number
+                    )
+                except botocore.exceptions.ClientError as e:
+                    print("Skipping Version {}: {}".format(version_number,
+                                                              e.message))
+
+
 def deploy(src, local_package=None):
     """Deploys a new function to AWS Lambda.
 

diff --git a/requirements.txt b/requirements.txt
@@ -6,6 +6,5 @@ futures==3.0.5
 jmespath==0.9.0
 pyaml==15.8.2
 python-dateutil==2.5.3
-python-lambda==0.2.0
 PyYAML==3.11
 six==1.10.0
diff --git a/scripts/lambda b/scripts/lambda
@@ -38,9 +38,16 @@ def invoke(event_file, verbose):
 def deploy(local_package):
     aws_lambda.deploy(CURRENT_DIR, local_package)
 
+
+@click.command(help="Delete old versions of your functions")
+@click.option("--keep-last", type=int, prompt="Please enter the number of recent versions to keep")
+def cleanup(keep_last):
+    aws_lambda.cleanup_old_versions(CURRENT_DIR, keep_last)
+
 if __name__ == '__main__':
     cli.add_command(init)
     cli.add_command(invoke)
     cli.add_command(deploy)
     cli.add_command(build)
+    cli.add_command(cleanup)
     cli()