Added support for aliasing function uploads

lambda_uploader/config.py

@@ -19,7 +19,8 @@
                    'region': basestring, 'handler': basestring,
                    'role': basestring, 'timeout': int, 'memory': int}
 
-DEFAULT_PARAMS = {u'requirements': [], u'publish': False}
+DEFAULT_PARAMS = {u'requirements': [], u'publish': False,
+                  u'alias': None, u'alias_description': None}
 
 
 class Config(object):
@@ -32,21 +33,44 @@ def __init__(self, pth):
         for param, clss in REQUIRED_PARAMS.iteritems():
             self._validate(param, cls=clss)
 
+    '''
+    Return raw config
+    '''
     @property
     def raw(self):
         if not self._config:
             self._load_config()
 
         return self._config
 
+    '''
+    Return an alias description if set otherwise return an the function
+    description
+    '''
+    @property
+    def alias_description(self):
+        if self._config['alias_description'] is None:
+            return self._config['description']
+        else:
+            return self._config['alias_description']
+
+    '''Set the publish attr to true'''
     def set_publish(self):
         self._config['publish'] = True
 
+    '''Set the alias and description'''
+    def set_alias(self, alias, description=None):
+        self._config['alias'] = alias
+        self._config['alias_description'] = description
+        self._config['publish'] = True
+
+    '''Set all defaults after loading the config'''
     def _set_defaults(self):
         for param, val in DEFAULT_PARAMS.iteritems():
             if self._config.get(param) is None:
                 self._config[param] = val
 
+    '''Validate the configuration file'''
     def _validate(self, key, cls=None):
         if key not in self._config:
             raise ValueError("Config %s must have %s set"
@@ -56,6 +80,7 @@ def _validate(self, key, cls=None):
                     raise TypeError("Config value '%s' should be %s not %s"
                                     % (key, cls, type(self._config[key])))
 
+    '''Load config ... called by init()'''
     def _load_config(self):
         config_file = path.join(self._path, 'lambda.json')
         if not path.isfile(config_file):

lambda_uploader/shell.py

@@ -48,9 +48,6 @@ def _execute(args):
     pth = path.abspath(args.function_dir)
 
     cfg = config.Config(pth)
-    # Set publish if flagged to do so
-    if args.publish:
-        cfg.set_publish()
 
     _print('Building Package')
     pkg = package.build_package(pth, cfg.requirements)
@@ -59,8 +56,23 @@ def _execute(args):
         pkg.clean_workspace()
 
     if not args.no_upload:
+        # Set publish if flagged to do so
+        if args.publish:
+            cfg.set_publish()
+
+        create_alias = False
+        # Set alias if the arg is passed
+        if args.alias is not None:
+            cfg.set_alias(args.alias, args.alias_description)
+            create_alias = True
+
         _print('Uploading Package')
-        uploader.upload_package(pkg, cfg, args.profile)
+        upldr = uploader.PackageUploader(cfg, args.profile)
+        upldr.upload(pkg)
+        # If the alias was set create it
+        if create_alias:
+            upldr.alias()
+
         pkg.clean_zipfile()
 
     _print('Fin')
@@ -88,6 +100,11 @@ def main(arv=None):
                         const=True)
     parser.add_argument('--profile', dest='profile',
                         help='specify AWS cli profile')
+    alias_help = 'alias for published version (WILL SET THE PUBLISH FLAG)'
+    parser.add_argument('--alias', '-a', dest='alias',
+                        default=None, help=alias_help)
+    parser.add_argument('--alias-description', '-m', dest='alias_description',
+                        default=None, help='alias description')
     parser.add_argument('function_dir', default=getcwd(), nargs='?',
                         help='lambda function directory')
 

lambda_uploader/uploader.py

@@ -18,52 +18,144 @@
 LOG = logging.getLogger(__name__)
 
 
-def upload_package(pkg, config, profile_name):
-    with open(pkg.zip_file, "rb") as fil:
-        zip_file = fil.read()
-
-    session = boto3.session.Session(region_name=config.region,
-                                    profile_name=profile_name)
-    client = session.client('lambda')
-    # Assume the function already exists in AWS
-    existing_function = True
-
-    try:
-        get_resp = client.get_function_configuration(FunctionName=config.name)
-        LOG.debug("AWS get_function_configuration response: %s" % get_resp)
-    except:
-        existing_function = False
-        LOG.debug("function not found creating new function")
-
-    if existing_function:
+class PackageUploader(object):
+    def __init__(self, config, profile_name):
+        self._config = config
+        session = boto3.session.Session(region_name=config.region,
+                                        profile_name=profile_name)
+        self._client = session.client('lambda')
+        self.version = None
+
+    '''
+    Calls the AWS methods to upload an existing package and update
+    the function configuration
+
+    returns the package version
+    '''
+    def upload_existing(self, pkg):
+        with open(pkg.zip_file, "rb") as fil:
+            zip_file = fil.read()
+
         LOG.debug('running update_function_code')
-        response = client.update_function_code(
-            FunctionName=config.name,
+        conf_update_resp = self._client.update_function_code(
+            FunctionName=self._config.name,
             ZipFile=zip_file,
-            Publish=config.publish,
+            Publish=False,
         )
-        LOG.debug("AWS update_function_code response: %s" % response)
+        LOG.debug("AWS update_function_code response: %s"
+                  % conf_update_resp)
         LOG.debug('running update_function_configuration')
-        response = client.update_function_configuration(
-            FunctionName=config.name,
-            Handler=config.handler,
-            Role=config.role,
-            Description=config.description,
-            Timeout=config.timeout,
-            MemorySize=config.memory,
+        response = self._client.update_function_configuration(
+            FunctionName=self._config.name,
+            Handler=self._config.handler,
+            Role=self._config.role,
+            Description=self._config.description,
+            Timeout=self._config.timeout,
+            MemorySize=self._config.memory,
         )
-        LOG.debug("AWS update_function_configuration response: %s" % response)
-    else:
+        LOG.debug("AWS update_function_configuration response: %s"
+                  % response)
+
+        version = response.get('Version')
+        # Publish the version after upload and config update if needed
+        if self._config.publish:
+            resp = self._client.publish_version(
+                    FunctionName=self._config.name,
+                    )
+            LOG.debug("AWS publish_version response: %s" % resp)
+            version = resp.get('Version')
+
+        return version
+
+    '''
+    Creates and uploads a new lambda function
+
+    returns the package version
+    '''
+    def upload_new(self, pkg):
+        with open(pkg.zip_file, "rb") as fil:
+            zip_file = fil.read()
+
         LOG.debug('running create_function_code')
-        response = client.create_function(
-            FunctionName=config.name,
+        response = self._client.create_function(
+            FunctionName=self._config.name,
             Runtime='python2.7',
-            Handler=config.handler,
-            Role=config.role,
+            Handler=self._config.handler,
+            Role=self._config.role,
             Code={'ZipFile': zip_file},
-            Description=config.description,
-            Timeout=config.timeout,
-            MemorySize=config.memory,
-            Publish=config.publish,
+            Description=self._config.description,
+            Timeout=self._config.timeout,
+            MemorySize=self._config.memory,
+            Publish=self._config.publish,
         )
         LOG.debug("AWS create_function response: %s" % response)
+
+        return response.get('Version')
+
+    '''
+    Auto determines whether the function exists or not and calls
+    the appropriate method (upload_existing or upload_new).
+    '''
+    def upload(self, pkg):
+        existing_function = True
+        try:
+            get_resp = self._client.get_function_configuration(
+                    FunctionName=self._config.name)
+            LOG.debug("AWS get_function_configuration response: %s" % get_resp)
+        except:
+            existing_function = False
+            LOG.debug("function not found creating new function")
+
+        if existing_function:
+            self.version = self.upload_existing(pkg)
+        else:
+            self.version = self.upload_new(pkg)
+
+    '''
+    Create/update an alias to point to the package. Raises an
+    exception if the package has not been uploaded.
+    '''
+    def alias(self):
+        # if self.version is still None raise exception
+        if self.version is None:
+            raise Exception('Must upload package before applying alias')
+
+        if self._alias_exists():
+            self._update_alias()
+        else:
+            self._create_alias()
+
+    '''
+    Pulls down the current list of aliases and checks to see if
+    an alias exists.
+    '''
+    def _alias_exists(self):
+        resp = self._client.list_aliases(
+                FunctionName=self._config.name)
+
+        for alias in resp.get('Aliases'):
+            if alias.get('Name') == self._config.alias:
+                return True
+        return False
+
+    '''Creates alias'''
+    def _create_alias(self):
+        LOG.debug("Creating new alias %s" % self._config.alias)
+        resp = self._client.create_alias(
+                FunctionName=self._config.name,
+                Name=self._config.alias,
+                FunctionVersion=self.version,
+                Description=self._config.alias_description,
+                )
+        LOG.debug("AWS create_alias response: %s" % resp)
+
+    '''Update alias'''
+    def _update_alias(self):
+        LOG.debug("Updating alias %s" % self._config.alias)
+        resp = self._client.update_alias(
+                FunctionName=self._config.name,
+                Name=self._config.alias,
+                FunctionVersion=self.version,
+                Description=self._config.alias_description,
+                )
+        LOG.debug("AWS update_alias response: %s" % resp)