FR - Allow to specify directory to write all the templates to in addition to S3.

ecs_composex/__init__.py

@@ -5,3 +5,5 @@
 __version__ = "0.1.3"
 
 XFILE_DEST = "ComposeXFile"
+DIR_DEST = "OutputDirectory"
+FILE_DEST = "OutputFile"

ecs_composex/cli.py

@@ -3,19 +3,21 @@
 
 """Console script for ecs_composex."""
 
-import sys
-import argparse
+import os
 import warnings
-import json
+
+import argparse
+import sys
 from boto3 import session
-from ecs_composex import XFILE_DEST
-from ecs_composex.common import LOG, write_template_to_file
-from ecs_composex.common.aws import BUCKET_NAME, CURATED_AZS
+
+from ecs_composex import XFILE_DEST, DIR_DEST
 from ecs_composex.common import KEYISSET
-from ecs_composex.common.cfn_tools import (
-    build_config_template_file,
-    write_config_template_file,
-)
+from ecs_composex.common import LOG
+from ecs_composex.common.aws import BUCKET_NAME, CURATED_AZS
+from ecs_composex.common.cfn_params import USE_FLEET_T
+from ecs_composex.common.cfn_tools import build_config_template_file
+from ecs_composex.common.templates import FileArtifact
+from ecs_composex.compute.compute_params import CLUSTER_NAME_T
 from ecs_composex.root import generate_full_template
 from ecs_composex.vpc.vpc_params import (
     APP_SUBNETS_T,
@@ -24,8 +26,6 @@
     VPC_ID_T,
     VPC_MAP_ID_T,
 )
-from ecs_composex.common.cfn_params import USE_FLEET_T
-from ecs_composex.compute.compute_params import CLUSTER_NAME_T
 
 
 def validate_vpc_input(args):
@@ -88,10 +88,17 @@ def main():
     parser.add_argument(
         "-o",
         "--output-file",
+        required=False,
+        default=f"{os.path.basename(os.path.dirname(__file__))}.yml",
+        help="Output file. Extension determines the file format",
+    )
+    parser.add_argument(
+        "-d",
+        "--output-dir",
+        required=False,
+        help="Output directory to write all the templates to.",
         type=str,
-        required=True,
-        help="The name and path of the main output file. If you specify extra arguments, it will create a parameters"
-        " file as well for creating your CFN Stack",
+        dest=DIR_DEST,
     )
     parser.add_argument(
         "--cfn-config-file",
@@ -239,15 +246,24 @@ def main():
 
     print("Arguments: " + str(args._))
     templates_and_params = generate_full_template(**vars(args))
-    write_template_to_file(templates_and_params[0], args.output_file)
+
+    template_file = FileArtifact(
+        args.output_file, template=templates_and_params[0], **vars(args)
+    )
+    template_file.create()
     cfn_config = build_config_template_file(templates_and_params[1])
     if KEYISSET("CfnConfigFile", vars(args)):
-        config_file_path = args.CfnConfigFile
+        config_file_name = args.CfnConfigFile
     else:
-        config_file_path = f"{args.output_file.split('.')[0]}.config.json"
-    write_config_template_file(cfn_config, config_file_path)
-    with open(f"{args.output_file.split('.')[0]}.params.json", "w") as params_fd:
-        params_fd.write(json.dumps(templates_and_params[1], indent=4))
+        config_file_name = f"{args.output_file.split('.')[0]}.config.json"
+    config_file = FileArtifact(config_file_name, content=cfn_config, **vars(args))
+    config_file.create()
+    params_file = FileArtifact(
+        f"{args.output_file.split('.')[0]}.params.json",
+        content=templates_and_params[1],
+        **vars(args),
+    )
+    params_file.create()
 
 
 if __name__ == "__main__":

ecs_composex/common/templates.py

@@ -3,11 +3,22 @@
 """
 Functions to manage a template and wheter it should be stored in S3
 """
+import yaml
+
+try:
+    from yaml import CDumper as Dumper
+except ImportError:
+    from yaml import Dumper
+from os import path, mkdir
 
 import boto3
+import json
 from botocore.exceptions import ClientError
+from datetime import datetime as dt
+from troposphere import Template
 
-from ecs_composex.common import DATE_PREFIX
+from ecs_composex import DIR_DEST
+from ecs_composex.common import DATE_PREFIX, KEYISSET
 from ecs_composex.common import LOG
 
 
@@ -100,6 +111,7 @@ def upload_template(
     prefix=None,
     session=None,
     client=None,
+    mime=None,
     **kwargs,
 ):
     """Upload template_body to a file in s3 with given prefix and bucket_name
@@ -125,7 +137,8 @@ def upload_template(
     if session is None:
         session = boto3.session.Session()
     assert check_bucket(bucket_name=bucket_name, session=session)
-
+    if mime is None:
+        mime = "application/json"
     if prefix is None:
         prefix = DATE_PREFIX
 
@@ -138,7 +151,7 @@ def upload_template(
             Key=key,
             Bucket=bucket_name,
             ContentEncoding="utf-8",
-            ContentType="application/json",
+            ContentType=mime,
             ServerSideEncryption="AES256",
             **kwargs,
         )
@@ -154,3 +167,130 @@ def upload_template(
     except Exception as error:
         LOG.debug(error)
         return None
+
+
+def write_file_to_directory(file_content, directory, file_name=None, **kwargs):
+    """
+    Function to write a file to a directory.
+    :param file_content:
+    :param directory:
+    :param file_name:
+    :return:
+    """
+
+
+class FileArtifact(object):
+    """
+    Class for a template file built.
+    """
+
+    url = None
+    body = None
+    template = None
+    file_name = None
+    mime = "text/plain"
+    session = boto3.session.Session()
+    bucket = None
+    uploadable = False
+    validated = False
+    output_dir = f"/tmp/{dt.utcnow().strftime('%s')}"
+    uploaded = False
+
+    def upload(self):
+        if not self.uploadable:
+            LOG.error("BucketName was not specified, not attempting upload")
+        elif not self.validated:
+            LOG.error("The template was not correctly validated. Not uploading")
+        else:
+            self.url = upload_template(
+                self.body, self.bucket, self.file_name, mime=self.mime, validate=False
+            )
+            LOG.info(f"Template {self.file_name} uploaded successfully to {self.url}")
+            self.uploaded = True
+
+    def write(self):
+        try:
+            mkdir(self.output_dir)
+            LOG.info(f"Created directory {self.output_dir} to store files")
+        except FileExistsError:
+            LOG.debug(f"Output directory {self.output_dir} already exists")
+        with open(f"{self.output_dir}/{self.file_name}", "w") as template_fd:
+            template_fd.write(self.body)
+            LOG.info(
+                f"Template {self.file_name} written successfully at {self.output_dir}/{self.file_name}"
+            )
+
+    def validate(self):
+        try:
+            self.session.client("cloudformation").validate_template(
+                TemplateBody=self.template.to_json()
+            )
+            LOG.debug(f"Template {self.file_name} was validated successfully by CFN")
+            self.validated = True
+        except ClientError as error:
+            LOG.error(error)
+
+    def define_file_specs(self):
+        """
+        Function to set the file body from template if self.template is Template
+        :return:
+        """
+        if self.file_name.endswith(".json"):
+            self.mime = "application/json"
+        elif self.file_name.endswith(".yml") or self.file_name.endswith(".yaml"):
+            self.mime = "application/x-yaml"
+        if isinstance(self.template, Template):
+            if self.mime == "application/x-yaml":
+                self.body = self.template.to_yaml()
+            else:
+                self.body = self.template.to_json()
+        elif isinstance(self.content, (list, dict, tuple)):
+            self.validated = True
+            self.uploadable = True
+            if self.mime == "application/x-yaml":
+                self.body = yaml.dump(self.content, Dumper=Dumper)
+            elif self.mime == "application/json":
+                self.body = json.dumps(self.content, indent=4)
+
+    def set_from_kwargs(self, **kwargs):
+        """
+        Function to set internal settings based on kwargs keys
+        :param kwargs:
+        """
+        if KEYISSET(DIR_DEST, kwargs):
+            self.output_dir = path.abspath(kwargs[DIR_DEST])
+        if KEYISSET("BucketName", kwargs):
+            self.bucket = kwargs["BucketName"]
+            self.uploadable = True
+
+    def create(self):
+        """
+        Function to write to file and upload in a single function
+        """
+        self.write()
+        self.upload()
+
+    def __init__(self, file_name, template=None, content=None, session=None, **kwargs):
+        """
+        Init function for our template file object
+        :param file_name: Name of the file. Mandatory
+        :param template: If you are providing a template to generate
+        :param body: raw content to write
+        :param session:
+        :param kwargs:
+        """
+        self.file_name = file_name
+        if template is not None and not isinstance(template, Template):
+            raise TypeError("template must be of type", Template, "got", type(template))
+        elif isinstance(template, Template):
+            self.template = template
+            self.validate()
+        if session is not None:
+            self.session = session
+        self.set_from_kwargs(**kwargs)
+        if content is not None and isinstance(content, (tuple, dict, str, list)):
+            self.content = content
+        self.define_file_specs()
+
+    def __repr__(self):
+        return f"{self.output_dir}/{self.file_name}"

ecs_composex/compute/cli.py

@@ -8,11 +8,11 @@
 
 import argparse
 import sys
-import json
+import os
 
 from boto3 import session
-from ecs_composex import XFILE_DEST
-from ecs_composex.common import write_template_to_file
+from ecs_composex import XFILE_DEST, DIR_DEST
+from ecs_composex.common.templates import FileArtifact
 from ecs_composex.common.aws import CURATED_AZS, BUCKET_NAME
 from ecs_composex.ecs.ecs_params import CLUSTER_NAME_T
 from ecs_composex.vpc.vpc_params import VPC_ID_T, APP_SUBNETS_T
@@ -35,9 +35,18 @@ def root_parser():
     parser.add_argument(
         "-o",
         "--output-file",
-        required=True,
+        required=False,
+        default=f"{os.path.basename(os.path.dirname(__file__))}.yml",
         help="Output file. Extension determines the file format",
     )
+    parser.add_argument(
+        "-d",
+        "--output-dir",
+        required=False,
+        help="Output directory to write all the templates to.",
+        type=str,
+        dest=DIR_DEST,
+    )
     #  AWS SETTINGS
     parser.add_argument(
         "--region",
@@ -90,6 +99,12 @@ def root_parser():
         help="Runs spotfleet for EC2. If used in combination "
         "of --use-fargate, it will create an additional SpotFleet",
     )
+    parser.add_argument(
+        "--no-upload",
+        action="store_true",
+        default=False,
+        help="Do not upload the file to S3.",
+    )
     return parser
 
 
@@ -100,9 +115,19 @@ def main():
     args = parser.parse_args()
 
     template_params = create_compute_stack(**vars(args))
-    write_template_to_file(template_params[0], args.output_file)
-    with open(f"{args.output_file.split('.')[0]}.params.json", "w") as params_fd:
-        params_fd.write(json.dumps(template_params[1], indent=4))
+    template_file = FileArtifact(
+        args.output_file, template=template_params[0], **vars(args)
+    )
+    params_file_name = f"{args.output_file.split('.')[0]}.params.json"
+    params_file = FileArtifact(
+        params_file_name, content=template_params[1], **vars(args)
+    )
+    if args.no_upload:
+        template_file.write()
+        params_file.write()
+    else:
+        template_file.create()
+        params_file.create()
     return 0