dbt-platform-helper 12.2.4__py3-none-any.whl → 12.4.0__py3-none-any.whl

dbt_platform_helper/domain/database_copy.py

@@ -8,13 +8,15 @@ from boto3 import Session
 
 from dbt_platform_helper.constants import PLATFORM_CONFIG_FILE
 from dbt_platform_helper.domain.maintenance_page import MaintenancePageProvider
-from dbt_platform_helper.exceptions import ApplicationNotFoundError
-from dbt_platform_helper.exceptions import AWSException
+from dbt_platform_helper.providers.aws import AWSException
 from dbt_platform_helper.utils.application import Application
+from dbt_platform_helper.utils.application import ApplicationNotFoundException
 from dbt_platform_helper.utils.application import load_application
 from dbt_platform_helper.utils.aws import Vpc
 from dbt_platform_helper.utils.aws import get_connection_string
 from dbt_platform_helper.utils.aws import get_vpc_info_by_name
+from dbt_platform_helper.utils.aws import wait_for_log_group_to_exist
+from dbt_platform_helper.utils.files import apply_environment_defaults
 from dbt_platform_helper.utils.messages import abort_with_error
 from dbt_platform_helper.utils.validation import load_and_validate_platform_config
 
@@ -25,80 +27,80 @@ class DatabaseCopy:
         app: str,
         database: str,
         auto_approve: bool = False,
-        load_application_fn: Callable[[str], Application] = load_application,
-        vpc_config_fn: Callable[[Session, str, str, str], Vpc] = get_vpc_info_by_name,
-        db_connection_string_fn: Callable[
+        load_application: Callable[[str], Application] = load_application,
+        vpc_config: Callable[[Session, str, str, str], Vpc] = get_vpc_info_by_name,
+        db_connection_string: Callable[
             [Session, str, str, str, Callable], str
         ] = get_connection_string,
         maintenance_page_provider: Callable[
             [str, str, list[str], str, str], None
         ] = MaintenancePageProvider(),
-        input_fn: Callable[[str], str] = click.prompt,
-        echo_fn: Callable[[str], str] = click.secho,
-        abort_fn: Callable[[str], None] = abort_with_error,
+        input: Callable[[str], str] = click.prompt,
+        echo: Callable[[str], str] = click.secho,
+        abort: Callable[[str], None] = abort_with_error,
     ):
         self.app = app
         self.database = database
         self.auto_approve = auto_approve
-        self.vpc_config_fn = vpc_config_fn
-        self.db_connection_string_fn = db_connection_string_fn
+        self.vpc_config = vpc_config
+        self.db_connection_string = db_connection_string
         self.maintenance_page_provider = maintenance_page_provider
-        self.input_fn = input_fn
-        self.echo_fn = echo_fn
-        self.abort_fn = abort_fn
+        self.input = input
+        self.echo = echo
+        self.abort = abort
 
         if not self.app:
             if not Path(PLATFORM_CONFIG_FILE).exists():
-                self.abort_fn("You must either be in a deploy repo, or provide the --app option.")
+                self.abort("You must either be in a deploy repo, or provide the --app option.")
 
             config = load_and_validate_platform_config()
             self.app = config["application"]
 
         try:
-            self.application = load_application_fn(self.app)
-        except ApplicationNotFoundError:
-            abort_fn(f"No such application '{app}'.")
+            self.application = load_application(self.app)
+        except ApplicationNotFoundException:
+            abort(f"No such application '{app}'.")
 
-    def _execute_operation(self, is_dump: bool, env: str, vpc_name: str):
+    def _execute_operation(self, is_dump: bool, env: str, vpc_name: str, filename: str):
         vpc_name = self.enrich_vpc_name(env, vpc_name)
 
         environments = self.application.environments
         environment = environments.get(env)
         if not environment:
-            self.abort_fn(
+            self.abort(
                 f"No such environment '{env}'. Available environments are: {', '.join(environments.keys())}"
             )
 
         env_session = environment.session
 
         try:
-            vpc_config = self.vpc_config_fn(env_session, self.app, env, vpc_name)
+            vpc_config = self.vpc_config(env_session, self.app, env, vpc_name)
         except AWSException as ex:
-            self.abort_fn(str(ex))
+            self.abort(str(ex))
 
         database_identifier = f"{self.app}-{env}-{self.database}"
 
         try:
-            db_connection_string = self.db_connection_string_fn(
+            db_connection_string = self.db_connection_string(
                 env_session, self.app, env, database_identifier
             )
         except Exception as exc:
-            self.abort_fn(f"{exc} (Database: {database_identifier})")
+            self.abort(f"{exc} (Database: {database_identifier})")
 
         try:
             task_arn = self.run_database_copy_task(
-                env_session, env, vpc_config, is_dump, db_connection_string
+                env_session, env, vpc_config, is_dump, db_connection_string, filename
             )
         except Exception as exc:
-            self.abort_fn(f"{exc} (Account id: {self.account_id(env)})")
+            self.abort(f"{exc} (Account id: {self.account_id(env)})")
 
         if is_dump:
             message = f"Dumping {self.database} from the {env} environment into S3"
         else:
             message = f"Loading data into {self.database} in the {env} environment from S3"
 
-        self.echo_fn(message, fg="white", bold=True)
-        self.echo_fn(
+        self.echo(message, fg="white", bold=True)
+        self.echo(
             f"Task {task_arn} started. Waiting for it to complete (this may take some time)...",
             fg="white",
         )
@@ -107,11 +109,10 @@ class DatabaseCopy:
     def enrich_vpc_name(self, env, vpc_name):
         if not vpc_name:
             if not Path(PLATFORM_CONFIG_FILE).exists():
-                self.abort_fn(
-                    "You must either be in a deploy repo, or provide the vpc name option."
-                )
+                self.abort("You must either be in a deploy repo, or provide the vpc name option.")
             config = load_and_validate_platform_config()
-            vpc_name = config.get("environments", {}).get(env, {}).get("vpc")
+            env_config = apply_environment_defaults(config)["environments"]
+            vpc_name = env_config.get(env, {}).get("vpc")
         return vpc_name
 
     def run_database_copy_task(
@@ -121,12 +122,15 @@ class DatabaseCopy:
         vpc_config: Vpc,
         is_dump: bool,
         db_connection_string: str,
+        filename: str,
     ) -> str:
         client = session.client("ecs")
         action = "dump" if is_dump else "load"
+        dump_file_name = filename if filename else "data_dump"
         env_vars = [
             {"name": "DATA_COPY_OPERATION", "value": action.upper()},
             {"name": "DB_CONNECTION_STRING", "value": db_connection_string},
+            {"name": "DUMP_FILE_NAME", "value": dump_file_name},
         ]
         if not is_dump:
             env_vars.append({"name": "ECS_CLUSTER", "value": f"{self.app}-{env}"})
@@ -156,12 +160,12 @@ class DatabaseCopy:
 
         return response.get("tasks", [{}])[0].get("taskArn")
 
-    def dump(self, env: str, vpc_name: str):
-        self._execute_operation(True, env, vpc_name)
+    def dump(self, env: str, vpc_name: str, filename: str = None):
+        self._execute_operation(True, env, vpc_name, filename)
 
-    def load(self, env: str, vpc_name: str):
+    def load(self, env: str, vpc_name: str, filename: str = None):
         if self.is_confirmed_ready_to_load(env):
-            self._execute_operation(False, env, vpc_name)
+            self._execute_operation(False, env, vpc_name, filename)
 
     def copy(
         self,
@@ -176,8 +180,8 @@ class DatabaseCopy:
         to_vpc = self.enrich_vpc_name(to_env, to_vpc)
         if not no_maintenance_page:
             self.maintenance_page_provider.activate(self.app, to_env, services, template, to_vpc)
-        self.dump(from_env, from_vpc)
-        self.load(to_env, to_vpc)
+        self.dump(from_env, from_vpc, f"data_dump_{to_env}")
+        self.load(to_env, to_vpc, f"data_dump_{to_env}")
         if not no_maintenance_page:
             self.maintenance_page_provider.deactivate(self.app, to_env)
 
@@ -185,7 +189,7 @@ class DatabaseCopy:
         if self.auto_approve:
             return True
 
-        user_input = self.input_fn(
+        user_input = self.input(
             f"\nWARNING: the load operation is destructive and will delete the {self.database} database in the {env} environment. Continue? (y/n)"
         )
         return user_input.lower().strip() in ["y", "yes"]
@@ -194,9 +198,11 @@ class DatabaseCopy:
         action = "dump" if is_dump else "load"
         log_group_name = f"/ecs/{self.app}-{env}-{self.database}-{action}"
         log_group_arn = f"arn:aws:logs:eu-west-2:{self.account_id(env)}:log-group:{log_group_name}"
-        self.echo_fn(f"Tailing {log_group_name} logs", fg="yellow")
+        self.echo(f"Tailing {log_group_name} logs", fg="yellow")
         session = self.application.environments[env].session
-        response = session.client("logs").start_live_tail(logGroupIdentifiers=[log_group_arn])
+        log_client = session.client("logs")
+        wait_for_log_group_to_exist(log_client, log_group_name)
+        response = log_client.start_live_tail(logGroupIdentifiers=[log_group_arn])
 
         stopped = False
         for data in response["responseStream"]:
@@ -210,9 +216,9 @@ class DatabaseCopy:
                     match = re.match(r"(Stopping|Aborting) data (load|dump).*", message)
                     if match:
                         if match.group(1) == "Aborting":
-                            self.abort_fn("Task aborted abnormally. See logs above for details.")
+                            self.abort("Task aborted abnormally. See logs above for details.")
                         stopped = True
-                    self.echo_fn(message)
+                    self.echo(message)
 
     def account_id(self, env):
         envs = self.application.environments

dbt_platform_helper/domain/maintenance_page.py

@@ -9,9 +9,9 @@ from typing import Union
 import boto3
 import click
 
-from dbt_platform_helper.providers.load_balancers import ListenerNotFoundError
-from dbt_platform_helper.providers.load_balancers import ListenerRuleNotFoundError
-from dbt_platform_helper.providers.load_balancers import LoadBalancerNotFoundError
+from dbt_platform_helper.providers.load_balancers import ListenerNotFoundException
+from dbt_platform_helper.providers.load_balancers import ListenerRuleNotFoundException
+from dbt_platform_helper.providers.load_balancers import LoadBalancerNotFoundException
 from dbt_platform_helper.providers.load_balancers import find_https_listener
 from dbt_platform_helper.utils.application import Environment
 from dbt_platform_helper.utils.application import Service
@@ -75,13 +75,13 @@ class MaintenancePageProvider:
             else:
                 raise click.Abort
 
-        except LoadBalancerNotFoundError:
+        except LoadBalancerNotFoundException:
             click.secho(
                 f"No load balancer found for environment {env} in the application {app}.", fg="red"
             )
             raise click.Abort
 
-        except ListenerNotFoundError:
+        except ListenerNotFoundException:
             click.secho(
                 f"No HTTPS listener found for environment {env} in the application {app}.", fg="red"
             )
@@ -110,13 +110,13 @@ class MaintenancePageProvider:
                 f"Maintenance page removed from environment {env} in application {app}", fg="green"
             )
 
-        except LoadBalancerNotFoundError:
+        except LoadBalancerNotFoundException:
             click.secho(
                 f"No load balancer found for environment {env} in the application {app}.", fg="red"
             )
             raise click.Abort
 
-        except ListenerNotFoundError:
+        except ListenerNotFoundException:
             click.secho(
                 f"No HTTPS listener found for environment {env} in the application {app}.", fg="red"
             )
@@ -180,7 +180,7 @@ def remove_maintenance_page(session: boto3.Session, listener_arn: str):
         deleted = delete_listener_rule(tag_descriptions, name, lb_client)
 
         if name == "MaintenancePage" and not deleted:
-            raise ListenerRuleNotFoundError()
+            raise ListenerRuleNotFoundException()
 
 
 def get_rules_tag_descriptions(rules: list, lb_client):

dbt_platform_helper/platform_exception.py

@@ -0,0 +1,5 @@
+# This exception exists so that we can easily catch exceptions
+# at the command level where we know we can just output the
+# error and abort.
+class PlatformException(Exception):
+    pass

dbt_platform_helper/providers/aws.py

@@ -0,0 +1,32 @@
+from dbt_platform_helper.platform_exception import PlatformException
+
+
+class AWSException(PlatformException):
+    pass
+
+
+class CreateTaskTimeoutException(AWSException):
+    def __init__(self, addon_name: str, application_name: str, environment: str):
+        super().__init__(
+            f"""Client ({addon_name}) ECS task has failed to start for "{application_name}" in "{environment}" environment."""
+        )
+
+
+class ImageNotFoundException(AWSException):
+    def __init__(self, commit: str):
+        super().__init__(
+            f"""The commit hash "{commit}" has not been built into an image, try the `platform-helper codebase build` command first."""
+        )
+
+
+class LogGroupNotFoundException(AWSException):
+    def __init__(self, log_group_name: str):
+        super().__init__(f"""No log group called "{log_group_name}".""")
+
+
+# Todo: This should probably be in the AWS Copilot provider, but was causing circular import when we tried it pre refactoring the utils/aws.py
+class CopilotCodebaseNotFoundException(PlatformException):
+    def __init__(self, codebase: str):
+        super().__init__(
+            f"""The codebase "{codebase}" either does not exist or has not been deployed."""
+        )

dbt_platform_helper/providers/cloudformation.py

@@ -1,105 +1,134 @@
 import json
 
+import botocore
 from cfn_tools import dump_yaml
 from cfn_tools import load_yaml
 
-
-def add_stack_delete_policy_to_task_role(cloudformation_client, iam_client, task_name: str):
-
-    stack_name = f"task-{task_name}"
-    stack_resources = cloudformation_client.list_stack_resources(StackName=stack_name)[
-        "StackResourceSummaries"
-    ]
-
-    for resource in stack_resources:
-        if resource["LogicalResourceId"] == "DefaultTaskRole":
-            task_role_name = resource["PhysicalResourceId"]
-            iam_client.put_role_policy(
-                RoleName=task_role_name,
-                PolicyName="DeleteCloudFormationStack",
-                PolicyDocument=json.dumps(
-                    {
-                        "Version": "2012-10-17",
-                        "Statement": [
-                            {
-                                "Action": ["cloudformation:DeleteStack"],
-                                "Effect": "Allow",
-                                "Resource": f"arn:aws:cloudformation:*:*:stack/{stack_name}/*",
-                            },
-                        ],
-                    },
-                ),
-            )
-
-
-def update_conduit_stack_resources(
-    cloudformation_client,
-    iam_client,
-    ssm_client,
-    application_name: str,
-    env: str,
-    addon_type: str,
-    addon_name: str,
-    task_name: str,
-    parameter_name: str,
-    access: str,
-):
-
-    conduit_stack_name = f"task-{task_name}"
-    template = cloudformation_client.get_template(StackName=conduit_stack_name)
-    template_yml = load_yaml(template["TemplateBody"])
-    template_yml["Resources"]["LogGroup"]["DeletionPolicy"] = "Retain"
-    template_yml["Resources"]["TaskNameParameter"] = load_yaml(
-        f"""
-        Type: AWS::SSM::Parameter
-        Properties:
-          Name: {parameter_name}
-          Type: String
-          Value: {task_name}
-        """
-    )
-
-    log_filter_role_arn = iam_client.get_role(RoleName="CWLtoSubscriptionFilterRole")["Role"]["Arn"]
-
-    destination_log_group_arns = json.loads(
-        ssm_client.get_parameter(Name="/copilot/tools/central_log_groups")["Parameter"]["Value"]
-    )
-
-    destination_arn = destination_log_group_arns["dev"]
-    if env.lower() in ("prod", "production"):
-        destination_arn = destination_log_group_arns["prod"]
-
-    template_yml["Resources"]["SubscriptionFilter"] = load_yaml(
-        f"""
-        Type: AWS::Logs::SubscriptionFilter
-        DeletionPolicy: Retain
-        Properties:
-          RoleArn: {log_filter_role_arn}
-          LogGroupName: /copilot/{task_name}
-          FilterName: /copilot/conduit/{application_name}/{env}/{addon_type}/{addon_name}/{task_name.rsplit("-", 1)[1]}/{access}
-          FilterPattern: ''
-          DestinationArn: {destination_arn}
-        """
-    )
-
-    params = []
-    if "Parameters" in template_yml:
-        for param in template_yml["Parameters"]:
-            # TODO testing missed in codecov, update test to assert on method call below with params including ExistingParameter from cloudformation template.
-            params.append({"ParameterKey": param, "UsePreviousValue": True})
-
-    cloudformation_client.update_stack(
-        StackName=conduit_stack_name,
-        TemplateBody=dump_yaml(template_yml),
-        Parameters=params,
-        Capabilities=["CAPABILITY_IAM"],
-    )
-
-    return conduit_stack_name
-
-
-# TODO Catch errors and raise a more human friendly Exception is the CloudFormation stack goes into a "unhappy" state, e.g. ROLLBACK_IN_PROGRESS. Currently we get things like botocore.exceptions.WaiterError: Waiter StackUpdateComplete failed: Waiter encountered a terminal failure state: For expression "Stacks[].StackStatus" we matched expected path: "UPDATE_ROLLBACK_COMPLETE" at least once
-def wait_for_cloudformation_to_reach_status(cloudformation_client, stack_status, stack_name):
-
-    waiter = cloudformation_client.get_waiter(stack_status)
-    waiter.wait(StackName=stack_name, WaiterConfig={"Delay": 5, "MaxAttempts": 20})
+from dbt_platform_helper.platform_exception import PlatformException
+
+
+class CloudFormation:
+    def __init__(self, cloudformation_client, iam_client, ssm_client):
+        self.cloudformation_client = cloudformation_client
+        self.iam_client = iam_client
+        self.ssm_client = ssm_client
+
+    def add_stack_delete_policy_to_task_role(self, task_name: str):
+        stack_name = f"task-{task_name}"
+        stack_resources = self.cloudformation_client.list_stack_resources(StackName=stack_name)[
+            "StackResourceSummaries"
+        ]
+
+        for resource in stack_resources:
+            if resource["LogicalResourceId"] == "DefaultTaskRole":
+                task_role_name = resource["PhysicalResourceId"]
+                self.iam_client.put_role_policy(
+                    RoleName=task_role_name,
+                    PolicyName="DeleteCloudFormationStack",
+                    PolicyDocument=json.dumps(
+                        {
+                            "Version": "2012-10-17",
+                            "Statement": [
+                                {
+                                    "Action": ["cloudformation:DeleteStack"],
+                                    "Effect": "Allow",
+                                    "Resource": f"arn:aws:cloudformation:*:*:stack/{stack_name}/*",
+                                },
+                            ],
+                        },
+                    ),
+                )
+
+    def update_conduit_stack_resources(
+        self,
+        application_name: str,
+        env: str,
+        addon_type: str,
+        addon_name: str,
+        task_name: str,
+        parameter_name: str,
+        access: str,
+    ):
+        conduit_stack_name = f"task-{task_name}"
+        template = self.cloudformation_client.get_template(StackName=conduit_stack_name)
+        template_yml = load_yaml(template["TemplateBody"])
+
+        template_yml["Resources"]["LogGroup"]["DeletionPolicy"] = "Retain"
+
+        template_yml["Resources"]["TaskNameParameter"] = load_yaml(
+            f"""
+            Type: AWS::SSM::Parameter
+            Properties:
+              Name: {parameter_name}
+              Type: String
+              Value: {task_name}
+            """
+        )
+
+        log_filter_role_arn = self.iam_client.get_role(RoleName="CWLtoSubscriptionFilterRole")[
+            "Role"
+        ]["Arn"]
+
+        destination_log_group_arns = json.loads(
+            self.ssm_client.get_parameter(Name="/copilot/tools/central_log_groups")["Parameter"][
+                "Value"
+            ]
+        )
+
+        destination_arn = destination_log_group_arns["dev"]
+        if env.lower() in ("prod", "production"):
+            destination_arn = destination_log_group_arns["prod"]
+
+        template_yml["Resources"]["SubscriptionFilter"] = load_yaml(
+            f"""
+            Type: AWS::Logs::SubscriptionFilter
+            DeletionPolicy: Retain
+            Properties:
+              RoleArn: {log_filter_role_arn}
+              LogGroupName: /copilot/{task_name}
+              FilterName: /copilot/conduit/{application_name}/{env}/{addon_type}/{addon_name}/{task_name.rsplit("-", 1)[1]}/{access}
+              FilterPattern: ''
+              DestinationArn: {destination_arn}
+            """
+        )
+
+        params = []
+        # TODO Currently not covered by tests - see https://uktrade.atlassian.net/browse/DBTP-1582
+        if "Parameters" in template_yml:
+            for param in template_yml["Parameters"]:
+                params.append({"ParameterKey": param, "UsePreviousValue": True})
+
+        self.cloudformation_client.update_stack(
+            StackName=conduit_stack_name,
+            TemplateBody=dump_yaml(template_yml),
+            Parameters=params,
+            Capabilities=["CAPABILITY_IAM"],
+        )
+
+        return conduit_stack_name
+
+    def wait_for_cloudformation_to_reach_status(self, stack_status, stack_name):
+        waiter = self.cloudformation_client.get_waiter(stack_status)
+
+        try:
+            waiter.wait(StackName=stack_name, WaiterConfig={"Delay": 5, "MaxAttempts": 20})
+        except botocore.exceptions.WaiterError as err:
+            current_status = err.last_response.get("Stacks", [{}])[0].get("StackStatus", "")
+
+            if current_status in [
+                "ROLLBACK_IN_PROGRESS",
+                "UPDATE_ROLLBACK_IN_PROGRESS",
+                "ROLLBACK_FAILED",
+            ]:
+                raise CloudFormationException(stack_name, current_status)
+            else:
+                raise CloudFormationException(
+                    stack_name, f"Error while waiting for stack status: {str(err)}"
+                )
+
+
+class CloudFormationException(PlatformException):
+    def __init__(self, stack_name: str, current_status: str):
+        super().__init__(
+            f"The CloudFormation stack '{stack_name}' is not in a good state: {current_status}"
+        )