fixup! Implement Lambda function versions (#5927)

Author: dsotirho-ucsc

scripts/sell_unused_slots.py

@@ -71,7 +71,7 @@ def _list_contribution_lambda_functions(cls) -> list[Lambda]:
         """
         return [
             lambda_
-            for lambda_ in Lambdas().list_lambdas(deployment='ALL')
+            for lambda_ in Lambdas().list_lambdas()
             if lambda_.is_contribution_lambda
         ]
 

src/azul/lambdas.py

@@ -13,7 +13,6 @@
 )
 
 from azul import (
-    JSON,
     R,
     cache,
     config,
@@ -38,7 +37,6 @@ class Lambda:
     name: str
     role: str
     slot_location: Optional[str]
-    version: str
 
     @property
     def is_contribution_lambda(self) -> bool:
@@ -84,15 +82,13 @@ def has_notification_queue(handler) -> bool:
     def from_response(cls, response: 'FunctionConfigurationTypeDef') -> Self:
         name = response['FunctionName']
         role = response['Role']
-        version = response['Version']
         try:
             slot_location = response['Environment']['Variables']['AZUL_TDR_SOURCE_LOCATION']
         except KeyError:
             slot_location = None
         return cls(name=name,
                    role=role,
-                   slot_location=slot_location,
-                   version=version)
+                   slot_location=slot_location)
 
     def __attrs_post_init__(self):
         if self.slot_location is None:
@@ -109,65 +105,27 @@ class Lambdas:
     def _lambda(self):
         return aws.lambda_
 
-    def list_lambdas(self,
-                     deployment: str
-                     ) -> list[Lambda]:
-        """
-        Return a list of AWS Lambda functions. Only the largest numbered version
-        of each function will be included in the list.
-
-        :param deployment: Limit output to the specified deployment stage. If
-                           'ALL', functions from all deployments will be
-                           returned.
-        """
-        paginator = self._lambda.get_paginator('list_functions')
-        lambda_prefixes = None if deployment == 'ALL' else [
-            config.qualified_resource_name(lambda_name, stage=deployment)
-            for lambda_name in config.lambda_names()
-        ]
-        functions: dict[str, JSON] = dict()
-        for response in paginator.paginate(FunctionVersion='ALL'):
-            for function in response['Functions']:
-                version = function['Version']
-                version = None if version == '$LATEST' else int(version)
-                if version and (lambda_prefixes is None or any(
-                    function['FunctionName'].startswith(prefix)
-                    for prefix in lambda_prefixes
-                )):
-                    name = function['FunctionName']
-                    previous_function = functions.get(name)
-                    if previous_function is None or version > int(previous_function['Version']):
-                        functions[name] = function
+    def list_lambdas(self) -> list[Lambda]:
         return [
             Lambda.from_response(function)
-            for function in functions.values()
+            for response in self._lambda.get_paginator('list_functions').paginate()
+            for function in response['Functions']
         ]
 
-    def get_function(self, function_name: str) -> JSON:
-        """
-        Return the Lambda client `get_function()` response for the largest
-        numbered version of the Lambda function.
-        """
-        paginator = self._lambda.get_paginator('list_versions_by_function')
-        params = {'FunctionName': function_name}
-        version = max([
-            int(function['Version'])
-            for response in paginator.paginate(**params)
-            for function in response['Versions']
-            if function['Version'] != '$LATEST'
-        ])
-        return self._lambda.get_function(FunctionName=function_name,
-                                         Qualifier=str(version))
-
     def manage_lambdas(self, enabled: bool):
-        for function in self.list_lambdas(deployment=config.deployment_stage):
-            self.manage_lambda(function.name, enabled)
+        lambda_prefixes = [
+            config.qualified_resource_name(lambda_infix)
+            for lambda_infix in config.lambda_names()
+        ]
+        assert all(lambda_prefixes)
+        for lambda_ in self.list_lambdas():
+            if any(lambda_.name.startswith(prefix) for prefix in lambda_prefixes):
+                self.manage_lambda(lambda_.name, enabled)
 
     def manage_lambda(self, lambda_name: str, enable: bool):
-        lambda_settings = self.get_function(function_name=lambda_name)
+        lambda_settings = self._lambda.get_function(FunctionName=lambda_name)
         lambda_arn = lambda_settings['Configuration']['FunctionArn']
         # Lambda does not support adding tags to function aliases or versions
-        lambda_arn, _, _ = lambda_arn.rpartition(':')
         lambda_tags = self._lambda.list_tags(Resource=lambda_arn)['Tags']
         lambda_name = lambda_settings['Configuration']['FunctionName']
         if enable:
@@ -211,7 +169,7 @@ def reset_lambda_roles(self):
         client = self._lambda
         lambda_names = set(config.lambda_names())
 
-        for lambda_ in self.list_lambdas(deployment='ALL'):
+        for lambda_ in self.list_lambdas():
             for lambda_name in lambda_names:
                 if lambda_.name.startswith(config.qualified_resource_name(lambda_name)):
                     other_lambda_name = one(lambda_names - {lambda_name})

src/azul/queues.py

@@ -504,11 +504,11 @@ def _wait_for_queue_empty(self, queue: 'Queue'):
 
     def _manage_sqs_push(self,
                          function_name: str,
-                         function_version: str,
+                         version_or_alias: str,
                          queue: 'Queue',
                          enable: bool):
         lambda_ = aws.lambda_
-        partial_arn = f'{function_name}:{function_version}'
+        partial_arn = f'{function_name}:{version_or_alias}'
         response = lambda_.list_event_source_mappings(FunctionName=partial_arn,
                                                       EventSourceArn=queue.attributes['QueueArn'])
         mapping_uuid = one(response['EventSourceMappings'])['UUID']
@@ -559,11 +559,6 @@ def manage_lambdas(self, queues: Mapping[str, 'Queue'], enable: bool):
         Enable or disable the readers and writers of the given queues.
         """
         functions_by_queue = self.functions_by_queue()
-        versions_by_function = {
-            f.name: f.version
-            for f in self._lambdas.list_lambdas(deployment=config.deployment_stage)
-        }
-        assert set(functions_by_queue.values()) <= set(versions_by_function)
 
         with ThreadPoolExecutor(max_workers=len(queues)) as tpe:
             futures = []
@@ -577,11 +572,10 @@ def submit(f, *args, **kwargs):
                 except KeyError:
                     assert queue_name in config.fail_queue_names
                 else:
-                    version = versions_by_function[function]
                     if queue_name == config.notifications_queue.name:
                         # Prevent new notifications from being added
                         submit(self._manage_lambda, config.indexer_name, enable)
-                    submit(self._manage_sqs_push, function, version, queue, enable)
+                    submit(self._manage_sqs_push, function, 'active', queue, enable)
             self._handle_futures(futures)
             futures = [tpe.submit(self._wait_for_queue_idle, queue) for queue in queues.values()]
             self._handle_futures(futures)

src/azul/terraform.py

@@ -756,8 +756,11 @@ def tf_config(self, app_name):
                 '${aws_vpc_endpoint.%s.id}' % app_name
             ]
 
+        assert 'aws_lambda_alias' not in resources
+        resources['aws_lambda_alias'] = dict()
+
         functions = json_item_dicts(json_dict(resources['aws_lambda_function']))
-        for _, resource in functions:
+        for resource_name, resource in functions:
             assert 'layers' not in resource
             resource['layers'] = ['${aws_lambda_layer_version.dependencies.arn}']
             # Publishing a new Lambda function version each time lets us perform
@@ -781,6 +784,15 @@ def tf_config(self, app_name):
             resource['source_code_hash'] = '${filebase64sha256("%s")}' % package_zip
             resource['filename'] = package_zip
 
+            # Inject a Lambda function alias resource that will point to the
+            # Lambda function version published when this and the
+            # aws_lambda_function resources are deployed.
+            json_dict(resources['aws_lambda_alias'])[resource_name] = {
+                'name': 'active',
+                'function_name': '${aws_lambda_function.%s.arn}' % resource_name,
+                'function_version': '${aws_lambda_function.%s.version}' % resource_name
+            }
+
         assert 'aws_cloudwatch_log_group' not in resources
         functions = json_item_dicts(resources['aws_lambda_function'])
         resources['aws_cloudwatch_log_group'] = {
@@ -927,36 +939,37 @@ def tf_config(self, app_name):
                 resource['event_source_arn'] = f'${{aws_sqs_queue.{sqs_name}.arn}}'
 
         # Replace the references to unqualified Lambda function ARNs emitted by
-        # Chalice with qualified ARNs. Note: `aws_lambda_permission` resources
-        # doesn't support the qualified ARN syntax, instead it has a `qualifier`
-        # argument.
+        # Chalice with references to the Lambda function alias.
         #
-        locals[app_name] = re.sub(r'(aws_lambda_function\.[^\.\s]+)\.invoke_arn',
-                                  r'\1.qualified_invoke_arn',
+        locals[app_name] = re.sub(r'\$\{aws_lambda_function\.([^\.\s]+)\.invoke_arn\}',
+                                  r'${aws_lambda_alias.\1.invoke_arn}',
                                   json_str(locals[app_name]))
         if app_name == 'indexer':
             resource_arguments = [
                 ('aws_cloudwatch_event_target', 'arn'),
+                ('aws_lambda_permission', 'function_name'),
                 ('aws_lambda_event_source_mapping', 'function_name'),
             ]
         elif app_name == 'service':
             resource_arguments = [
                 ('aws_cloudwatch_event_target', 'arn'),
+                ('aws_lambda_permission', 'function_name'),
             ]
         else:
             assert False, app_name
         for resource_type, argument in resource_arguments:
             for _, resource in json_item_dicts(resources[resource_type]):
-                value = json_str(resource[argument])
-                assert value.endswith('.arn}'), resource
-                resource[argument] = value.replace('.arn', '.qualified_arn')
+                resource[argument] = re.sub(r'\$\{aws_lambda_function\.([^\.\s]+)\.arn\}',
+                                            r'${aws_lambda_alias.\1.arn}',
+                                            json_str(resource[argument]))
 
-        # Add a qualifier argument to `aws_lambda_permission` resources
+        # Ensure that the Lambda permissions for the previous Lambda aliases
+        # aren't deleted until after the permissions for the new Lambda aliases
+        # have been created.
         #
-        for _, resource in json_item_dicts(resources['aws_lambda_permission']):
-            assert 'qualifier' not in resource, resource
-            lambda_arn = resource['function_name']
-            resource['qualifier'] = lambda_arn.replace('.arn', '.version')
+        for name, resource in json_item_dicts(resources['aws_lambda_permission']):
+            assert 'lifecycle' not in resource, resource
+            resource['lifecycle'] = {'create_before_destroy': True}
 
         return {
             'resource': resources,