Merge branch 'morosi-asset' into 'master'

Use parameters for the S3 keys of assets

See merge request it/e3-aws!77

Author: adanaja

src/e3/aws/cfn/main.py

@@ -156,8 +156,6 @@ def __init__(
         self.assume_role = assume_role
         self.aws_env: Session | AWSEnv | None = None
         self.deploy_branch = deploy_branch
-        # A temporary dir will be assigned when generating assets
-        self.gen_assets_dir: str | None = None
 
         self.timestamp = datetime.utcnow().strftime("%Y-%m-%d/%H:%M:%S.%f")
 
@@ -264,7 +262,7 @@ def _upload_dir(
                         upload_bucket.push(key=s3_object_key, content=fd, exist_ok=True)
 
     def _upload_stack(self, stack: Stack) -> None:
-        """Upload stack assets, data, and template to S3.
+        """Upload stack data and template to S3.
 
         :param stack: the stack to upload
         """
@@ -274,23 +272,7 @@ def _upload_stack(self, stack: Stack) -> None:
 
         assert self.args is not None
 
-        if self.aws_env:
-            s3 = self.aws_env.client("s3")
-        else:
-            s3 = None
-            logging.warning(
-                "no aws session, won't be able to check if assets exist in the bucket"
-            )
-
-        # Synchronize assets to the bucket before creating the stack
-        if self.gen_assets_dir is not None and self.s3_assets_key is not None:
-            self._upload_dir(
-                root_dir=self.gen_assets_dir,
-                s3_bucket=self.s3_bucket,
-                s3_key=self.s3_assets_key,
-                s3_client=s3,
-                check_exists=True,
-            )
+        s3 = self.aws_env.client("s3") if self.aws_env else None
 
         with tempfile.TemporaryDirectory() as tempd:
             # Push data associated with CFNMain and then all data
@@ -510,31 +492,19 @@ def execute(
                 return 1
 
         return_val = 0
+        stacks = self.create_stack()
+
+        if isinstance(stacks, list):
+            for stack in stacks:
+                return_val = self.execute_for_stack(stack, aws_env=aws_env)
+                # Stop at first failure
+                if return_val:
+                    return return_val
+        else:
+            return_val = self.execute_for_stack(stacks, aws_env=aws_env)
 
-        # Create a temporary assets dir here as assets need to be generated at the
-        # time of create_stack
-        with tempfile.TemporaryDirectory() as temp_assets_dir:
-            self.gen_assets_dir = temp_assets_dir
-            self.pre_create_stack()
-            stacks = self.create_stack()
-            self.post_create_stack()
-
-            if isinstance(stacks, list):
-                for stack in stacks:
-                    return_val = self.execute_for_stack(stack, aws_env=aws_env)
-                    # Stop at first failure
-                    if return_val:
-                        return return_val
-            else:
-                return_val = self.execute_for_stack(stacks, aws_env=aws_env)
-
-        self.gen_assets_dir = None
         return return_val
 
-    def pre_create_stack(self) -> None:
-        """Before create_stack."""
-        pass
-
     @abc.abstractmethod
     def create_stack(self) -> Stack | list[Stack]:
         """Create a stack.
@@ -543,10 +513,6 @@ def create_stack(self) -> Stack | list[Stack]:
         """
         pass
 
-    def post_create_stack(self) -> None:
-        """After create_stack."""
-        pass
-
     @property
     def stack_policy_body(self) -> str | None:
         """Stack Policy that can be set by calling the command ``protect``.

src/e3/aws/troposphere/__init__.py

@@ -1,15 +1,19 @@
 from __future__ import annotations
 from abc import ABC, abstractmethod
+from tempfile import TemporaryDirectory
 from itertools import chain
-from troposphere import AWSObject, Output, Template
+from troposphere import AWSObject, Output, Template, Parameter
+from collections import deque
+import logging
 
 from e3.aws import cfn, name_to_id, Session
 from e3.aws.cfn.main import CFNMain
 from e3.aws.troposphere.iam.policy_document import PolicyDocument
 from typing import TYPE_CHECKING
 
+
 if TYPE_CHECKING:  # all: no cover
-    from typing import Union
+    from typing import Union, Any
     from collections.abc import Iterable
     from troposphere import And, Condition, Equals, If, Not, Or
 
@@ -40,20 +44,78 @@ def cfn_policy_document(self, stack: Stack) -> PolicyDocument:
         """
         return PolicyDocument([])
 
-    def create_assets_dir(self, root_dir: str) -> None:  # noqa: B027
-        """Put assets in root_dir before export to S3 bucket referenced by the stack.
+    def create_data_dir(self, root_dir: str) -> None:  # noqa: B027
+        """Put data in root_dir before export to S3 bucket referenced by the stack.
 
-        :param root_dir: local directory in which assets should be stored. Assets will
+        :param root_dir: local directory in which data should be stored. Data will
             be then uploaded to an S3 bucket accessible from the template. The
             target location is the one received by resources method. Note that
             the same root_dir is shared by all resources in your stack.
         """
         pass
 
-    def create_data_dir(self, root_dir: str) -> None:  # noqa: B027
-        """Put data in root_dir before export to S3 bucket referenced by the stack.
 
-        :param root_dir: local directory in which data should be stored. Data will
+class Asset(Construct):
+    """Generic asset.
+
+    Assets are local files or directories that are uploaded to S3 and that can be
+    referenced by other resources. For example, an asset might be a directory
+    that contains the handler code for an AWS Lambda function. Assets can represent
+    any artifact that the app needs to operate.
+
+    Each asset insert an additional parameter to the CloudFormation template, that
+    can be used by other resources, with the intrinsic function Fn::Sub, to
+    reference the S3 key of the asset.
+    """
+
+    def __init__(self, name: str) -> None:
+        """Initialize Asset.
+
+        :param name: the logical name for CloudFormation
+        """
+        self.name = name
+        self.s3_key_parameter_name = f"{self.name}S3Key"
+
+    @property
+    @abstractmethod
+    def s3_key(self) -> str | None:
+        """Return the S3 key of this asset.
+
+        It may return None if the S3 key is not yet known.
+        """
+        ...
+
+    @property
+    def s3_key_parameter(self) -> str:
+        """Return the parameter that stores the S3 key.
+
+        The Default value is omitted if the S3 key is not yet known.
+        """
+        params: dict[str, Any] = {}
+        s3_key = self.s3_key
+        if s3_key is not None:
+            params["Default"] = s3_key
+
+        return Parameter(
+            self.s3_key_parameter_name,
+            Type="String",
+            Description=f"S3 key of asset {self.name}",
+            **params,
+        )
+
+    def resources(self, stack: Stack) -> list[AWSObject | Construct]:
+        """Return no resources."""
+        # Add the parameter during template creation even if the S3 key may not
+        # be known yet. This is useful if creating a stack from code, so that
+        # the exported template contains the parameter
+        stack.add_parameter(self.s3_key_parameter)
+        return []
+
+    @abstractmethod
+    def create_assets_dir(self, root_dir: str) -> None:  # noqa: B027
+        """Put assets in root_dir before export to S3 bucket referenced by the stack.
+
+        :param root_dir: local directory in which assets should be stored. Assets will
             be then uploaded to an S3 bucket accessible from the template. The
             target location is the one received by resources method. Note that
             the same root_dir is shared by all resources in your stack.
@@ -75,7 +137,6 @@ def __init__(
         s3_key: str | None = None,
         s3_assets_key: str | None = None,
         version: str | None = None,
-        gen_assets_dir: str | None = None,
     ) -> None:
         """Initialize Stack attributes.
 
@@ -89,7 +150,6 @@ def __init__(
         :param s3_key: s3 prefix in s3_bucket in which data is stored
         :param s3_assets_key: s3 prefix in s3_bucket in which assets are stored
         :param version: template format version
-        :param gen_assets_dir: directory where to generate stack assets to upload to S3
         """
         super().__init__(
             stack_name,
@@ -99,28 +159,22 @@ def __init__(
             s3_key=s3_key,
         )
         self.constructs: list[Construct | AWSObject] = []
+        self.assets: dict[str, Asset] = {}
 
         self.deploy_session = deploy_session
         self.dry_run = dry_run
         self.version = version
         self.s3_assets_key = s3_assets_key
-        self.gen_assets_dir = gen_assets_dir
         self.template = Template()
 
     def construct_to_objects(self, construct: Construct | AWSObject) -> list[AWSObject]:
         """Return list of AWS objects resources from a construct.
 
-        The function create_assets_dir is called at the same time, on each construct,
-        as some AWS objects may have properties that depend on assets.
-
         :param construct: construct to list resources from
         """
         if isinstance(construct, AWSObject):
             return [construct]
         else:
-            if self.gen_assets_dir is not None:
-                construct.create_assets_dir(self.gen_assets_dir)
-
             return list(
                 chain.from_iterable(
                     [
@@ -145,10 +199,23 @@ def add(self, element: AWSObject | Construct | Stack) -> Stack:
         # Add the new constructs (non expanded)
         self.constructs += constructs
 
-        # Update the template
+        # Update the template with AWSObjects.
+        # Convert Constructs to AWSObjects | Constructs recursively
         resources = []
-        for construct in constructs:
-            resources.extend(self.construct_to_objects(construct))
+        constructs_to_objects = deque(constructs)
+        while constructs_to_objects:
+            construct = constructs_to_objects.pop()
+            if isinstance(construct, AWSObject):
+                resources.append(construct)
+            else:
+                # Special case to keep track of Assets and generate parameters
+                # for the S3 keys
+                if isinstance(construct, Asset):
+                    self.add_parameter(construct.s3_key_parameter)
+                    self.assets[construct.name] = construct
+
+                constructs_to_objects.extend(construct.resources(stack=self))
+
         self.template.add_resource(resources)
 
         return self
@@ -163,6 +230,20 @@ def extend(self, elements: Iterable[AWSObject | Construct | Stack]) -> Stack:
 
         return self
 
+    def add_parameter(self, parameter: Parameter | list[Parameter]) -> None:
+        """Add parameters to stack template.
+
+        :param parameter: parameter to add to the template
+        """
+        if not isinstance(parameter, list):
+            parameter = [parameter]
+
+        for param in parameter:
+            if param.title in self.template.parameters:
+                self.template.parameters[param.title] = param
+            else:
+                self.template.add_parameter(param)
+
     def add_output(self, output: Output | list[Output]) -> None:
         """Add outputs to stack template.
 
@@ -266,14 +347,7 @@ def __init__(
             s3_key=self.s3_data_key,
             s3_assets_key=self.s3_assets_key,
         )
-
-    def pre_create_stack(self) -> None:
-        """Assign the temporary assets directory to the stack."""
-        self.stack.gen_assets_dir = self.gen_assets_dir
-
-    def post_create_stack(self) -> None:
-        """Unassign the temporary assets directory from the stack."""
-        self.stack.gen_assets_dir = None
+        self.gen_assets_dir: str | None = None
 
     def add(self, element: AWSObject | Construct | Stack) -> Stack:
         """Add resource to project's stack.
@@ -288,3 +362,59 @@ def extend(self, elements: list[AWSObject | Construct | Stack]) -> Stack:
         :param elements: resources to add to the stack.
         """
         return self.stack.extend(elements)
+
+    def _upload_stack(self, stack: cfn.Stack) -> None:
+        """See CFNMain."""
+        # Nothing to upload if there is no S3 bucket or S3 assets key set
+        if self.s3_bucket is not None and self.s3_assets_key is not None:
+            # Upload assets to S3 first
+            if self.aws_env:
+                s3 = self.aws_env.client("s3")
+            else:
+                s3 = None
+                logging.warning(
+                    "no aws session, won't be able to check if assets exist "
+                    "in the bucket"
+                )
+
+            assert self.gen_assets_dir is not None
+            self._upload_dir(
+                root_dir=self.gen_assets_dir,
+                s3_bucket=self.s3_bucket,
+                s3_key=self.s3_assets_key,
+                s3_client=s3,
+                check_exists=True,
+            )
+
+        # Upload the rest
+        super()._upload_stack(stack)
+
+    def execute_for_stack(
+        self, stack: cfn.Stack, aws_env: Session | None = None
+    ) -> int:
+        """See CFNMain."""
+        # Set the directory where to generate assets
+        with TemporaryDirectory() as tmpd:
+            self.gen_assets_dir = tmpd
+
+            # Create the assets directory
+            assert self.args is not None
+            if isinstance(stack, Stack) and self.args.command in [
+                "show",
+                "push",
+                "update",
+            ]:
+                for asset in stack.assets.values():
+                    # Populate the assets directory
+                    asset.create_assets_dir(root_dir=tmpd)
+
+                    # Add a parameter to the stack with the known S3 key of the
+                    # asset
+                    stack.add_parameter(asset.s3_key_parameter)
+
+            try:
+                # Execute the command for the stack
+                return super().execute_for_stack(stack=stack, aws_env=aws_env)
+            finally:
+                # Unset the assets directory
+                self.gen_assets_dir = None