Merge pull request #237 from adanaja/morosi-feat

Improve RestApi for stages and aliases

Author: adanaja

src/e3/aws/troposphere/apigateway/__init__.py

@@ -170,19 +170,32 @@ def __init__(
         path: str,
         method_list: list[Method],
         resource_list: list[Resource] | None = None,
+        integration_uri: str | Ref | Sub | None = None,
         lambda_arn: str | GetAtt | Ref | None = None,
+        lambda_arn_permission: str
+        | GetAtt
+        | Ref
+        | dict[str, str | GetAtt | Ref]
+        | None = None,
     ) -> None:
         """Initialize a REST API resource.
 
         :param path: the last path segment for this resource
         :param method_list: a list of methods accepted on this resource
         :param resource_list: a list of child resources
+        :param integration_uri: URI of a lambda function for this resource
         :param lambda_arn: arn of the lambda executed for this resource
+        :param lambda_arn_permission: lambda arn for which to add InvokeFunction
+            permission (can be different from the lambda arn executed
+            by the REST API). A mapping from stage names to lambda arns can
+            also be passed
         """
         self.path = path
         self.method_list = method_list
         self.resource_list = resource_list
+        self.integration_uri = integration_uri
         self.lambda_arn = lambda_arn
+        self.lambda_arn_permission = lambda_arn_permission
 
 
 class Api(Construct):
@@ -224,7 +237,6 @@ def __init__(
         :param hosted_zone_id: id of the hosted zone that contains domain_name.
             This parameter is required if domain_name is not None
         :param stages_config: configurations of the different stages
-        :param integration_uri: URI of a Lambda function
         """
         self.name = name
         self.description = description
@@ -895,35 +907,39 @@ def declare_stage(
     def _declare_method(
         self,
         method: Method,
-        resource: Resource,
         resource_id_prefix: str,
         resource_path: str,
+        resource_integration_uri: str | Ref | Sub | None = None,
+        resource_lambda_arn: str | GetAtt | Ref | None = None,
+        resource_lambda_arn_permission: str
+        | GetAtt
+        | Ref
+        | dict[str, str | GetAtt | Ref]
+        | None = None,
     ) -> list[AWSObject]:
         """Declare a method.
 
         :param method: the method definition
-        :param resource: resource associated with the method
         :param resource_id_prefix: resource_id without trailing Resource
         :param resource_path: absolute path to the resource
+        :param resource_integration_uri: integration URI for the resource
+        :param resource_lambda_arn: arn of lambda for the resource
+        :param resource_lambda_arn_permission: lambda arn permission for the resource
         :return: a list of AWSObjects to be added to the stack
         """
         result = []
         id_prefix = name_to_id(f"{resource_id_prefix}-{method.method}")
 
-        # Take the global lambda_arn or the one configured for the resource
-        lambda_arn = (
-            self.lambda_arn if resource.lambda_arn is None else resource.lambda_arn
-        )
-
-        # Integration URI for the resource
+        # Take the global integration uri or the one configured for the resource
         integration_uri = (
             self.integration_uri
-            if self.integration_uri is not None
-            else Sub(
-                "arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31"
-                "/functions/${lambdaArn}/invocations",
-                dict_values={"lambdaArn": lambda_arn},
-            )
+            if resource_integration_uri is None
+            else resource_integration_uri
+        )
+
+        # Take the global lambda arn or the one configured for the resource
+        lambda_arn = (
+            self.lambda_arn if resource_lambda_arn is None else resource_lambda_arn
         )
 
         integration = apigateway.Integration(
@@ -934,7 +950,13 @@ def _declare_method(
             IntegrationHttpMethod="POST",
             PassthroughBehavior="NEVER",
             Type="AWS_PROXY",
-            Uri=integration_uri,
+            Uri=integration_uri
+            if integration_uri is not None
+            else Sub(
+                "arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31"
+                "/functions/${lambdaArn}/invocations",
+                dict_values={"lambdaArn": lambda_arn},
+            ),
         )
 
         method_params = {
@@ -952,6 +974,16 @@ def _declare_method(
         result.append(apigateway.Method(f"{id_prefix}Method", **method_params))
 
         for config in self.stages_config:
+            if resource_lambda_arn_permission is not None:
+                # Use the lambda_arn_permission configured for resource
+                if isinstance(resource_lambda_arn_permission, dict):
+                    assert (
+                        config.name in resource_lambda_arn_permission
+                    ), f"missing lambda arn permission for stage {config.name}"
+                    lambda_arn = resource_lambda_arn_permission[config.name]
+                else:
+                    lambda_arn = resource_lambda_arn_permission
+
             result.append(
                 awslambda.Permission(
                     name_to_id(f"{id_prefix}-{config.name}LambdaPermission"),
@@ -1019,13 +1051,25 @@ def _declare_resources(
         resource_list: list[Resource],
         parent_id_prefix: str | None = None,
         parent_path: str | None = None,
+        parent_integration_uri: str | Ref | Sub | None = None,
+        parent_lambda_arn: str | GetAtt | Ref | None = None,
+        parent_lambda_arn_permission: str
+        | GetAtt
+        | Ref
+        | dict[str, str | GetAtt | Ref]
+        | None = None,
     ) -> list[AWSObject]:
         """Create API resources and methods recursively.
 
         Each resource can define its own methods and have child resources.
 
         :param resource_list: list of resources
         :param parent_id_prefix: id of the parent resource without trailing Resource
+        :param parent_path: absolute path to the parent resource
+        :param parent_integration_uri: integration URI of the parent resource
+        :param parent_lambda_arn: lambda arn of the parent resource
+        :param parent_lambda_arn_permission: lambda arn permission of the
+            parent resource
         :return: a list of AWSObjects to be added to the stack
         """
         result: list[AWSObject] = []
@@ -1059,13 +1103,36 @@ def _declare_resources(
 
             result.append(resource)
 
+            # Get the integration URI of this resource.
+            # It must be forwarded to children so that they recursively use the
+            # same URI
+            resource_integration_uri = (
+                r.integration_uri
+                if r.integration_uri is not None
+                else parent_integration_uri
+            )
+
+            # Same for the lambda arn
+            resource_lambda_arn = (
+                r.lambda_arn if r.lambda_arn is not None else parent_lambda_arn
+            )
+
+            # Same fo the lambda arn permission
+            resource_lambda_arn_permission = (
+                r.lambda_arn_permission
+                if r.lambda_arn_permission is not None
+                else parent_lambda_arn_permission
+            )
+
             # Declare the methods of this resource
             for method in r.method_list:
                 result += self._declare_method(
                     method=method,
-                    resource=r,
                     resource_id_prefix=resource_id_prefix,
                     resource_path=resource_path,
+                    resource_integration_uri=resource_integration_uri,
+                    resource_lambda_arn=resource_lambda_arn,
+                    resource_lambda_arn_permission=resource_lambda_arn_permission,
                 )
 
             # Declare the children of this resource
@@ -1074,6 +1141,9 @@ def _declare_resources(
                     resource_list=r.resource_list,
                     parent_id_prefix=resource_id_prefix,
                     parent_path=resource_path,
+                    parent_integration_uri=resource_integration_uri,
+                    parent_lambda_arn=resource_lambda_arn,
+                    parent_lambda_arn_permission=resource_lambda_arn_permission,
                 )
 
         return result

src/e3/aws/troposphere/awslambda/__init__.py

@@ -509,22 +509,29 @@ def __init__(
         description: str,
         lambda_arn: str | GetAtt | Ref,
         lambda_version: str,
+        alias_name: str | None = None,
         provisioned_concurrency_config: awslambda.ProvisionedConcurrencyConfiguration
         | None = None,
         routing_config: awslambda.AliasRoutingConfiguration | None = None,
     ):
         """Initialize an AWS lambda alias.
 
-        :param name: function name
-        :param description: a description of the function
+        :param name: name of the resource
+        :param description: a description of the alias
         :param lambda_arn: the name of the Lambda function
         :param lambda_version: the function version that the alias invokes
+        :param alias_name: name of the alias. By default the parameter
+            name will be used as both the name of the resource and the name
+            of the alias, so this allows for a different alias name. For
+            example if you have multiple Lambda functions using the same
+            alias names
         :param provisioned_concurrency_config: specifies a provisioned
             concurrency configuration for a function's alias
         :param routing_config: the routing configuration of the alias
         """
         self.name = name
         self.description = description
+        self.alias_name = alias_name
         self.lambda_arn = lambda_arn
         self.lambda_version = lambda_version
         self.provisioned_concurrency_config = provisioned_concurrency_config
@@ -537,7 +544,7 @@ def ref(self) -> Ref:
     def resources(self, stack: Stack) -> list[AWSObject]:
         """Return list of AWSObject associated with the construct."""
         params = {
-            "Name": self.name,
+            "Name": self.alias_name if self.alias_name is not None else self.name,
             "Description": self.description,
             "FunctionName": self.lambda_arn,
             "FunctionVersion": self.lambda_version,
@@ -769,9 +776,11 @@ def create_alias(
             :param default_name: default alias name if none is specified
             """
             name = config.name if config.name is not None else default_name
+            id = name_to_id(f"{self.lambda_name}-{name}-alias")
             return Alias(
-                name=name_to_id(f"{self.lambda_name}-{name}-alias"),
+                name=id,
                 description=f"{name} alias for {self.lambda_name} lambda",
+                alias_name=config.name if config.name is not None else id,
                 lambda_arn=self.lambda_arn,
                 lambda_version=config.version,
                 provisioned_concurrency_config=config.provisioned_concurrency_config,

tests/tests_e3_aws/troposphere/apigateway/apigateway_test.py

@@ -714,29 +714,15 @@ def test_rest_api_nested_resources(stack: Stack, lambda_fun: PyFunction) -> None
     stack.s3_bucket = "cfn_bucket"
     stack.s3_key = "templates/"
 
-    # Lambda for the products resource
-    products_lambda = PyFunction(
-        name="productslambda",
-        description="this is a test",
-        role="somearn",
-        code_dir="my_code_dir",
-        handler="app.main",
-        runtime="python3.8",
-        logs_retention_in_days=None,
-    )
-
     rest_api = RestApi(
         name="testapi",
         description="this is a test",
         lambda_arn=lambda_fun.ref,
         resource_list=[
-            Resource(path="accounts", method_list=[Method("ANY")]),
             Resource(
-                path="products",
-                # Specific lambda for this resource
-                lambda_arn=products_lambda.ref,
+                path="foo",
                 method_list=[Method("ANY")],
-                resource_list=[Resource(path="abcd", method_list=[Method("GET")])],
+                resource_list=[Resource(path="bar", method_list=[Method("GET")])],
             ),
         ],
     )
@@ -751,3 +737,97 @@ def test_rest_api_nested_resources(stack: Stack, lambda_fun: PyFunction) -> None
 
     print(stack.export()["Resources"])
     assert stack.export()["Resources"] == expected
+
+
+def test_rest_api_multi_lambdas_stages(stack: Stack) -> None:
+    """Test REST API with multiple lambdas and stages."""
+    stack.s3_bucket = "cfn_bucket"
+    stack.s3_key = "templates/"
+
+    # Create two lambdas for two different methods
+    accounts_lambda, products_lambda = [
+        PyFunction(
+            name=f"{name}lambda",
+            description="this is a test",
+            role="somearn",
+            code_dir="my_code_dir",
+            handler="app.main",
+            runtime="python3.8",
+            logs_retention_in_days=None,
+        )
+        for name in ("accounts", "products")
+    ]
+
+    # Create lambda versions
+    accounts_lambda_versions, products_lambda_versions = [
+        AutoVersion(2, lambda_function=lambda_fun)
+        for lambda_fun in (accounts_lambda, products_lambda)
+    ]
+
+    # Create lambda aliases.
+    # Share the same alias names as it will make it easier to setup the stage
+    # variable for using the right alias depending on the stage
+    accounts_lambda_aliases, products_lambda_aliases = [
+        BlueGreenAliases(
+            blue_config=BlueGreenAliasConfiguration(
+                name="Blue", version=lambda_versions.previous.version
+            ),
+            green_config=BlueGreenAliasConfiguration(
+                name="Green", version=lambda_versions.latest.version
+            ),
+            lambda_function=lambda_fun,
+        )
+        for lambda_versions, lambda_fun in (
+            (accounts_lambda_versions, accounts_lambda),
+            (products_lambda_versions, products_lambda),
+        )
+    ]
+
+    # Create the REST API
+    rest_api = RestApi(
+        name="testapi",
+        description="this is a test",
+        # Not important as it's overriden in resources
+        lambda_arn=accounts_lambda.ref,
+        # Declare prod and beta stages redirecting to correct aliases
+        stages_config=[
+            StageConfiguration("default", variables={"lambdaAlias": "Blue"}),
+            StageConfiguration("beta", variables={"lambdaAlias": "Green"}),
+        ],
+        # Declare two resources pointing to two different lambdas
+        resource_list=[
+            Resource(
+                path=path,
+                # Action to invoke the lambda with correct alias
+                integration_uri="arn:aws:apigateway:us-east-1:lambda:path/2015-03-31/"
+                "functions/arn:aws:lambda:eu-west-1:123456789012:function:"
+                f"{lambda_fun.name}:${{stageVariables.lambdaAlias}}/invocations",
+                # Lambda ARNs for InvokeFunction permissions depending on the stage
+                lambda_arn_permission={
+                    "default": lambda_aliases.blue.ref,
+                    "beta": lambda_aliases.green.ref,
+                },
+                method_list=[Method("ANY")],
+            )
+            for path, lambda_fun, lambda_aliases in (
+                ("accounts", accounts_lambda, accounts_lambda_aliases),
+                ("products", products_lambda, products_lambda_aliases),
+            )
+        ],
+    )
+
+    stack.add(accounts_lambda)
+    stack.add(products_lambda)
+    stack.add(accounts_lambda_versions)
+    stack.add(products_lambda_versions)
+    stack.add(accounts_lambda_aliases)
+    stack.add(products_lambda_aliases)
+    stack.add(rest_api)
+
+    with open(
+        os.path.join(TEST_DIR, "apigatewayv1_test_multi_lambdas_stages.json"),
+    ) as fd:
+        expected = json.load(fd)
+
+    print(stack.export()["Resources"])
+    assert stack.export()["Resources"] == expected