Merge branch 'morosi-asset' into 'master'

Add an AWS CDK Asset equivalent

See merge request it/e3-aws!62

Author: adanaja

src/e3/aws/cfn/__init__.py

@@ -482,7 +482,7 @@ def __contains__(self, key: str) -> bool:
         return key in self.resources
 
     def create_data_dir(self, root_dir: str) -> None:
-        """Populate directory that will be exported into a S3 bucket for the stack.
+        """Populate data directory that will be exported into a S3 bucket for the stack.
 
         :param root_dir: temporary local directory
         """

src/e3/aws/cfn/main.py

@@ -13,6 +13,7 @@
 from e3.os.process import PIPE
 from e3.vcs.git import GitRepository
 from e3.aws import AWSEnv, Session
+from e3.aws.s3 import bucket
 from e3.aws.cfn import Stack
 from e3.env import Env
 from e3.fs import find, sync_tree
@@ -26,6 +27,7 @@ def __init__(
         self,
         regions: list[str],
         default_profile: str | None = None,
+        assets_dir: str | None = None,
         data_dir: str | None = None,
         s3_bucket: str | None = None,
         s3_key: str = "",
@@ -37,6 +39,7 @@ def __init__(
 
         :param regions: list of regions on which we can operate
         :param default_profile: default AWS profile to use to create the stack
+        :param assets_dir: directory containing assets of the stack
         :param data_dir: directory containing files used by cfn-init
         :param s3_bucket: if defined S3 will be used as a proxy for resources.
             Template body will be uploaded to S3 before calling operation on
@@ -140,8 +143,11 @@ def __init__(
 
         self.regions = regions
 
+        self.assets_dir = assets_dir
         self.data_dir = data_dir
         self.s3_bucket = s3_bucket
+        self.s3_assets_key = None
+        self.s3_assets_url = None
         self.s3_data_key = None
         self.s3_data_url = None
         self.s3_template_key = None
@@ -150,23 +156,28 @@ 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")
 
         if s3_bucket is not None:
-            s3_root_key = (
-                "/".join([s3_key.rstrip("/"), self.timestamp]).strip("/") + "/"
-            )
-            self.s3_data_key = s3_root_key + "data/"
-            self.s3_data_url = "https://%s.s3.amazonaws.com/%s" % (
-                self.s3_bucket,
-                self.s3_data_key,
-            )
-            self.s3_template_key = s3_root_key + "template"
-            self.s3_template_url = "https://%s.s3.amazonaws.com/%s" % (
-                self.s3_bucket,
-                self.s3_template_key,
+            s3_root_key = f"{s3_key.strip('/')}/"
+            s3_root_url = f"https://{self.s3_bucket}.s3.amazonaws.com/"
+
+            # Assets use a static key
+            self.s3_assets_key = f"{s3_root_key}assets/"
+            self.s3_assets_url = f"{s3_root_url}{self.s3_assets_key}"
+
+            # Data and template use a dynamic key based on the timestamp
+            s3_timestamp_key = (
+                "/".join([s3_root_key.rstrip("/"), self.timestamp]).strip("/") + "/"
             )
+            self.s3_data_key = f"{s3_timestamp_key}data/"
+            self.s3_data_url = f"{s3_root_url}{self.s3_data_key}"
+
+            self.s3_template_key = f"{s3_timestamp_key}template"
+            self.s3_template_url = f"{s3_root_url}{self.s3_template_key}"
 
     @property
     def dry_run(self) -> bool:
@@ -200,6 +211,116 @@ def _prompt_yes(self, msg: str) -> bool:
         ask = input(f"{msg} (y/N): ")
         return ask[0] in "Yy"
 
+    def _upload_dir(
+        self,
+        root_dir: str,
+        s3_bucket: str,
+        s3_key: str,
+        s3_client: botocore.client.S3 | None = None,
+        check_exists: bool = False,
+    ) -> None:
+        """Upload directory to S3 bucket.
+
+        :param root_dir: directory
+        :param s3_bucket: bucket where to upload files
+        :param s3_key: key prefix for uploaded files
+        :param s3_client: a client for the S3 API
+        :param check_exists: check if an S3 object exists before uploading it
+        """
+        assert self.args is not None
+
+        for f in find(root_dir):
+            subkey = os.path.relpath(f, root_dir).replace("\\", "/")
+
+            logging.info(
+                "Upload %s to %s:%s%s",
+                subkey,
+                s3_bucket,
+                s3_key,
+                subkey,
+            )
+
+            if s3_client is None:
+                continue
+
+            with bucket(
+                s3_bucket, client=s3_client, auto_create=False
+            ) as upload_bucket:
+                # Check already existing S3 objects.
+                # Ignore the potential 403 error as CFN roles often only have the
+                # s3:GetObject permission on the bucket
+                s3_object_key = f"{s3_key}{subkey}"
+                if check_exists and upload_bucket.object_exists(
+                    s3_object_key, ignore_error_403=True
+                ):
+                    logging.info(
+                        "Skip already existing %s",
+                        subkey,
+                    )
+                    continue
+
+                if not self.args.dry_run:
+                    with open(f, "rb") as fd:
+                        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.
+
+        :param stack: the stack to upload
+        """
+        # Nothing to upload if there is no S3 bucket set
+        if self.s3_bucket is None:
+            return
+
+        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,
+            )
+
+        with tempfile.TemporaryDirectory() as tempd:
+            # Push data associated with CFNMain and then all data
+            # related to the stack
+            self.create_data_dir(root_dir=tempd)
+            stack.create_data_dir(root_dir=tempd)
+
+            if self.s3_data_key is not None:
+                # synchronize data to the bucket before creating the stack
+                self._upload_dir(
+                    root_dir=tempd,
+                    s3_bucket=self.s3_bucket,
+                    s3_key=self.s3_data_key,
+                    s3_client=s3,
+                )
+
+        if self.s3_template_key is not None:
+            logging.info(
+                "Upload template to %s:%s",
+                self.s3_bucket,
+                self.s3_template_key,
+            )
+            if s3 is not None and not self.args.dry_run:
+                s3.put_object(
+                    Bucket=self.s3_bucket,
+                    Body=stack.body.encode("utf-8"),
+                    ServerSideEncryption="AES256",
+                    Key=self.s3_template_key,
+                )
+
     def _push_stack_changeset(self, stack: Stack, s3_template_url: str | None) -> int:
         """Push the changeset of a stack from an already uploaded S3 template.
 
@@ -280,47 +401,7 @@ def execute_for_stack(self, stack: Stack, aws_env: Session | None = None) -> int
 
             if self.args.command in ("push", "update"):
                 # Synchronize resources to the S3 bucket
-                if not self.args.dry_run:
-                    assert self.aws_env
-                    s3 = self.aws_env.client("s3")
-
-                with tempfile.TemporaryDirectory() as tempd:
-                    # Push data associated with CFNMain and then all data
-                    # related to the stack
-                    self.create_data_dir(root_dir=tempd)
-                    stack.create_data_dir(root_dir=tempd)
-
-                    if self.s3_data_key is not None:
-                        # synchronize data to the bucket before creating the stack
-                        for f in find(tempd):
-                            with open(f, "rb") as fd:
-                                subkey = os.path.relpath(f, tempd).replace("\\", "/")
-                                logging.info(
-                                    "Upload %s to %s:%s%s",
-                                    subkey,
-                                    self.s3_bucket,
-                                    self.s3_data_key,
-                                    subkey,
-                                )
-                                if not self.args.dry_run:
-                                    s3.put_object(
-                                        Bucket=self.s3_bucket,
-                                        Body=fd,
-                                        ServerSideEncryption="AES256",
-                                        Key=self.s3_data_key + subkey,
-                                    )
-
-                if self.s3_template_key is not None:
-                    logging.info(
-                        "Upload template to %s:%s", self.s3_bucket, self.s3_template_key
-                    )
-                    if not self.args.dry_run:
-                        s3.put_object(
-                            Bucket=self.s3_bucket,
-                            Body=stack.body.encode("utf-8"),
-                            ServerSideEncryption="AES256",
-                            Key=self.s3_template_key,
-                        )
+                self._upload_stack(stack)
 
                 logging.info("Validate template for stack %s" % stack.name)
                 if not self.args.dry_run:
@@ -334,6 +415,7 @@ def execute_for_stack(self, stack: Stack, aws_env: Session | None = None) -> int
                     return self._push_stack_changeset(
                         stack=stack, s3_template_url=self.s3_template_url
                     )
+
             elif self.args.command == "show":
                 print(stack.body)
             elif self.args.command == "protect":
@@ -428,19 +510,31 @@ 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.
@@ -449,6 +543,10 @@ 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/mock/__init__.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+from unittest.mock import patch
+from contextlib import contextmanager
+
+if TYPE_CHECKING:
+    from typing import Any
+    from collections.abc import Iterator
+
+    from e3.aws.troposphere.awslambda import PyFunctionAsset
+
+
+@contextmanager
+def mock_pyfunctionasset() -> Iterator[None]:
+    """Mock PyFunctionAsset.
+
+    PyFunctionAsset does a pip install and packaging of source files that
+    may be necessary to disable in some tests. With this mock, the checksum
+    "dummychecksum" is assigned to assets instead.
+    """
+
+    def mock_create_assets_dir(self: PyFunctionAsset, *args: Any, **kargs: Any) -> Any:
+        """Disable create_assets_dir and assign a dummy checksum."""
+        self.checksum = "dummychecksum"
+
+    with patch(
+        "e3.aws.troposphere.awslambda.PyFunctionAsset.create_assets_dir",
+        mock_create_assets_dir,
+    ):
+        yield

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

@@ -8,7 +8,7 @@
 import boto3
 
 if TYPE_CHECKING:
-    from typing import Any
+    from typing import Any, BinaryIO
     from collections.abc import Iterable, Iterator
 
 logger = logging.getLogger("e3.aws.s3")
@@ -99,7 +99,9 @@ def delete_bucket(self) -> None:
         self.clear_bucket()
         self.client.delete_bucket(Bucket=self.bucket)
 
-    def push(self, key: str, content: bytes, exist_ok: bool | None = None) -> None:
+    def push(
+        self, key: str, content: bytes | BinaryIO, exist_ok: bool | None = None
+    ) -> None:
         """Push content to S3.
 
         You can set exist_ok to false to prevent the object from being
@@ -179,6 +181,28 @@ def bucket_exists(self) -> bool:
                 return False
             raise
 
+    def object_exists(self, key: str, /, ignore_error_403: bool = False) -> bool:
+        """Check if an object exists.
+
+        :param key: object key
+        :param ignore_error_403: boto3.head_object returns a 403 error when the
+            object doesn't exist and the IAM role doesn't have the s3:ListBucket
+            permission. Setting ignore_error_403=True makes the function return
+            False instead of raising a ClientError
+        :return: if the object exists
+        :raises ClientError: in case of error, or in case of permission issue
+            when the object doesn't exist, the IAM role doesn't have the
+            s3:ListBucket permission, and ignore_error_403 is False
+        """
+        try:
+            self.client.head_object(Bucket=self.bucket, Key=key)
+            return True
+        except ClientError as e:
+            error_code = e.response["Error"]["Code"]
+            if error_code == "404" or (error_code == "403" and ignore_error_403):
+                return False
+            raise
+
     @property
     def key_count(self) -> int:
         """Return the number of keys from S3."""