dbt-platform-helper 13.1.0__py3-none-any.whl → 15.16.0__py3-none-any.whl

dbt_platform_helper/providers/ecs.py

@@ -1,9 +1,32 @@
 import random
 import string
-import time
-from typing import List
+import subprocess
+from typing import Any
+from typing import Optional
+
+from botocore.exceptions import ClientError
 
 from dbt_platform_helper.platform_exception import PlatformException
+from dbt_platform_helper.platform_exception import ValidationException
+from dbt_platform_helper.providers.vpc import Vpc
+from dbt_platform_helper.utilities.decorators import retry
+from dbt_platform_helper.utilities.decorators import wait_until
+
+
+class ECSException(PlatformException):
+    pass
+
+
+class ECSAgentNotRunningException(ECSException):
+    def __init__(self):
+        super().__init__("""ECS exec agent never reached "RUNNING" status""")
+
+
+class NoClusterException(ECSException):
+    def __init__(self, application_name: str, environment: str):
+        super().__init__(
+            f"""No ECS cluster found for "{application_name}" in "{environment}" environment."""
+        )
 
 
 class ECS:
@@ -13,7 +36,49 @@ class ECS:
         self.application_name = application_name
         self.env = env
 
-    def get_cluster_arn(self) -> str:
+    def start_ecs_task(
+        self,
+        cluster_name: str,
+        container_name: str,
+        task_def_arn: str,
+        vpc_config: Vpc,
+        env_vars: list[dict] = None,
+    ):
+        container_override = {"name": container_name}
+        if env_vars:
+            container_override["environment"] = env_vars
+
+        response = self.ecs_client.run_task(
+            taskDefinition=task_def_arn,
+            cluster=cluster_name,
+            capacityProviderStrategy=[
+                {"capacityProvider": "FARGATE", "weight": 1, "base": 0},
+            ],
+            enableExecuteCommand=True,
+            networkConfiguration={
+                "awsvpcConfiguration": {
+                    "subnets": vpc_config.public_subnets,
+                    "securityGroups": vpc_config.security_groups,
+                    "assignPublicIp": "ENABLED",
+                }
+            },
+            overrides={"containerOverrides": [container_override]},
+        )
+
+        return response.get("tasks", [{}])[0].get("taskArn")
+
+    def get_cluster_arn_by_name(self, cluster_name: str) -> str:
+        clusters = self.ecs_client.describe_clusters(
+            clusters=[
+                cluster_name,
+            ],
+        )["clusters"]
+        if len(clusters) == 1 and "clusterArn" in clusters[0]:
+            return clusters[0]["clusterArn"]
+
+        raise NoClusterException(self.application_name, self.env)
+
+    def get_cluster_arn_by_copilot_tag(self) -> str:
         """Returns the ARN of the ECS cluster for the given application and
         environment."""
         for cluster_arn in self.ecs_client.list_clusters()["clusterArns"]:
@@ -45,58 +110,185 @@ class ECS:
             random_id = "".join(random.choices(string.ascii_lowercase + string.digits, k=12))
             return f"conduit-{self.application_name}-{self.env}-{addon_name}-{random_id}"
 
-    def get_ecs_task_arns(self, cluster_arn: str, task_name: str):
-        """Gets the ECS task ARNs for a given task name and cluster ARN."""
-        tasks = self.ecs_client.list_tasks(
-            cluster=cluster_arn,
-            desiredStatus="RUNNING",
-            family=f"copilot-{task_name}",
-        )
+    def get_ecs_task_arns(
+        self,
+        cluster: str,
+        max_results: int = 100,
+        desired_status: str = "RUNNING",
+        service_name: Optional[str] = None,
+        started_by: Optional[str] = None,
+        task_def_family: Optional[str] = None,
+    ) -> list[str]:
+        """Returns the ECS task ARNs based on the parameters provided."""
+
+        params = {
+            "cluster": cluster,
+            "maxResults": max_results,
+            "desiredStatus": desired_status,
+        }
+
+        if service_name:
+            params["serviceName"] = service_name
+        if started_by:
+            params["startedBy"] = started_by
+        if task_def_family:
+            params["family"] = task_def_family
+
+        tasks = self.ecs_client.list_tasks(**params)
 
         if not tasks["taskArns"]:
             return []
 
         return tasks["taskArns"]
 
-    def ecs_exec_is_available(self, cluster_arn: str, task_arns: List[str]):
+    @retry()
+    def exec_task(self, cluster_arn: str, task_arn: str, subprocess_call=subprocess.call):
+        result = subprocess_call(
+            f"aws ecs execute-command --cluster {cluster_arn} "
+            f"--task {task_arn} "
+            f"--interactive --command bash ",
+            shell=True,
+        )
+        if result != 0:
+            raise PlatformException("Failed to exec into ECS task.")
+        return result
+
+    @wait_until(
+        max_attempts=25,
+        exceptions_to_catch=(ECSException,),
+        message_on_false="ECS Agent Not running",
+    )
+    def ecs_exec_is_available(self, cluster_arn: str, task_arns: list[str]) -> bool:
         """
         Checks if the ExecuteCommandAgent is running on the specified ECS task.
 
         Waits for up to 25 attempts, then raises ECSAgentNotRunning if still not
         running.
         """
-        current_attempts = 0
-        execute_command_agent_status = ""
+        if not task_arns:
+            raise ValidationException("No task ARNs provided")
+        task_details = self.ecs_client.describe_tasks(cluster=cluster_arn, tasks=task_arns)
 
-        while execute_command_agent_status != "RUNNING" and current_attempts < 25:
-            current_attempts += 1
+        if not task_details["tasks"]:
+            raise ECSException("No ECS tasks returned.")
+        container_details = task_details["tasks"][0]["containers"][0]
+        if container_details.get("managedAgents", None):
+            managed_agents = container_details["managedAgents"]
+        else:
+            raise ECSException("No managed agent on ecs task.")
 
-            task_details = self.ecs_client.describe_tasks(cluster=cluster_arn, tasks=task_arns)
+        execute_command_agent = [
+            agent for agent in managed_agents if agent["name"] == "ExecuteCommandAgent"
+        ]
+        if not execute_command_agent:
+            raise ECSException("No ExecuteCommandAgent on ecs task.")
+        return execute_command_agent[0]["lastStatus"] == "RUNNING"
 
-            managed_agents = task_details["tasks"][0]["containers"][0]["managedAgents"]
-            execute_command_agent_status = [
-                agent["lastStatus"]
-                for agent in managed_agents
-                if agent["name"] == "ExecuteCommandAgent"
-            ][0]
-            if execute_command_agent_status != "RUNNING":
-                time.sleep(1)
+    @wait_until(
+        max_attempts=20,
+        message_on_false="ECS task did not register in time",
+    )
+    def wait_for_task_to_register(self, cluster_arn: str, task_family: str) -> list[str]:
+        task_arns = self.get_ecs_task_arns(cluster=cluster_arn, task_def_family=task_family)
+        if task_arns:
+            return task_arns
+        return False
 
-        if execute_command_agent_status != "RUNNING":
-            raise ECSAgentNotRunningException
+    def get_service_deployment_state(
+        self, cluster_name: str, service_name: str, start_time: float
+    ) -> tuple[Optional[str], Optional[str]]:
+        """
+        Returns status & statusReason for the deployment of an ECS service.
 
+        statusReason can be:
+            PENDING | SUCCESSFUL | STOPPED | STOP_REQUESTED |
+            IN_PROGRESS | ROLLBACK_REQUESTED | ROLLBACK_IN_PROGRESS |
+            ROLLBACK_SUCCESSFUL | ROLLBACK_FAILED
+        """
+        resp = self.ecs_client.list_service_deployments(
+            cluster=cluster_name, service=service_name, createdAt={"after": start_time}
+        )
+        deployments = resp.get("serviceDeployments", [])
 
-class ECSException(PlatformException):
-    pass
+        if not deployments:
+            return None, f"No deployments found for '{service_name}'"
 
+        return deployments[0].get("status"), deployments[0].get("statusReason")
 
-class ECSAgentNotRunningException(ECSException):
-    def __init__(self):
-        super().__init__("""ECS exec agent never reached "RUNNING" status""")
+    def get_container_names_from_ecs_tasks(
+        self, cluster_name: str, task_ids: list[str]
+    ) -> list[str]:
+        """Retrieve container names from each ECS task provided."""
 
+        response = self.ecs_client.describe_tasks(cluster=cluster_name, tasks=task_ids)
 
-class NoClusterException(ECSException):
-    def __init__(self, application_name: str, environment: str):
-        super().__init__(
-            f"""No ECS cluster found for "{application_name}" in "{environment}" environment."""
-        )
+        names = []
+        for task in response.get("tasks", []):
+            for container in task.get("containers", []):
+                if container["name"] not in names:
+                    names.append(container["name"])
+        return names
+
+    def register_task_definition(
+        self,
+        service: str,
+        task_definition: dict,
+        image_tag: Optional[str] = None,
+    ) -> str:
+        """Register a new task definition revision using provided model and
+        containerDefinitions."""
+
+        for container in task_definition["containerDefinitions"]:
+            if container["name"] == service:
+                container["image"] = f"{container['image']}:{image_tag}"
+                break
+
+        try:
+            task_definition_response = self.ecs_client.register_task_definition(**task_definition)
+            return task_definition_response["taskDefinition"]["taskDefinitionArn"]
+        except ClientError as err:
+            raise PlatformException(f"Error registering task definition: {err}")
+
+    def update_service(
+        self,
+        service: str,
+        task_def_arn: str,
+        environment: str,
+        application: str,
+        desired_count: int,
+    ) -> dict[str, Any]:
+        """Update an ECS service and return the response."""
+
+        try:
+            service_response = self.ecs_client.update_service(
+                cluster=f"{application}-{environment}-cluster",
+                service=f"{application}-{environment}-{service}",
+                taskDefinition=task_def_arn,
+                desiredCount=desired_count,
+            )
+            return service_response["service"]
+        except ClientError as err:
+            raise PlatformException(f"Error updating ECS service: {err}")
+
+    def describe_service(self, service: str, environment: str, application: str) -> dict[str, Any]:
+        """Return information about an ECS service."""
+
+        try:
+            service_response = self.ecs_client.describe_services(
+                cluster=f"{application}-{environment}-cluster",
+                services=[
+                    f"{application}-{environment}-{service}",
+                ],
+            )
+            return service_response["services"][0]
+        except ClientError as err:
+            raise PlatformException(f"Error retrieving ECS service: {err}")
+
+    def describe_tasks(self, cluster_name: str, task_ids: list[str]) -> list[dict[str, Any]]:
+        """Return information about ECS tasks."""
+
+        try:
+            response = self.ecs_client.describe_tasks(cluster=cluster_name, tasks=task_ids)
+            return response["tasks"]
+        except ClientError as err:
+            raise PlatformException(f"Error retrieving ECS tasks: {err}")

dbt_platform_helper/providers/environment_variable.py

@@ -0,0 +1,24 @@
+import os
+from typing import Optional
+
+from dbt_platform_helper.platform_exception import PlatformException
+
+
+class EnvironmentVariableProvider:
+
+    @staticmethod
+    def get(env_var: str, default=None) -> Optional[str]:
+        """Returns the stripped value or None if not set or empty."""
+        value = os.environ.get(env_var)
+        if value and value.strip():
+            return value.strip()
+        return default
+
+    @staticmethod
+    def get_required(env_var: str) -> str:
+        """Returns the stripped value or raises a PlatformException if not set
+        or empty."""
+        value = os.environ.get(env_var)
+        if not value or not value.strip():
+            raise PlatformException(f"Environment variable '{env_var}' is not set or is empty")
+        return value.strip()

dbt_platform_helper/providers/files.py

@@ -17,7 +17,7 @@ class FileProvider:
         file_path.write_text(contents)
 
         action = "overwritten" if file_exists and overwrite else "created"
-        return f"File {file_name} {action}"
+        return f"File {file_path} {action}"
 
     @staticmethod
     def delete_file(base_path: str, file_name: str):

dbt_platform_helper/providers/io.py

@@ -1,20 +1,35 @@
 import click
 
 from dbt_platform_helper.platform_exception import PlatformException
+from dbt_platform_helper.providers.environment_variable import (
+    EnvironmentVariableProvider,
+)
 
 
 class ClickIOProvider:
+    def __init__(self, env_var_provider=EnvironmentVariableProvider()):
+        self.env_var_provider = env_var_provider
+        self.debug_flag = self.env_var_provider.get("DEBUG")
+
     def warn(self, message: str):
         click.secho(message, fg="magenta")
 
+    def debug(self, message: str):
+        if self.debug_flag and self.debug_flag.strip().upper() == "TRUE":
+            click.secho(message, fg="green")
+
     def error(self, message: str):
         click.secho(f"Error: {message}", fg="red")
 
-    def info(self, message: str):
-        click.secho(message)
+    def info(self, message: str, **kwargs):
+        click.secho(message, **kwargs)
 
-    def input(self, message: str) -> str:
-        return click.prompt(message)
+    def input(
+        self, message: str, hide_input=False, confirmation_prompt=False, input_type=str
+    ) -> str:
+        return click.prompt(
+            message, hide_input=hide_input, confirmation_prompt=confirmation_prompt, type=input_type
+        )
 
     def confirm(self, message: str) -> bool:
         try:
@@ -26,6 +41,23 @@ class ClickIOProvider:
         click.secho(f"Error: {message}", err=True, fg="red")
         exit(1)
 
+    def deploy_error(self, message: str):
+        click.secho(message, fg="red")
+
+    # TODO: DBTP-1979: messages will be a ValidationMessages class rather than a free-rein dictionary
+    def process_messages(self, messages: dict):
+        if not messages:
+            return
+
+        if messages.get("errors"):
+            self.error("\n".join(messages["errors"]))
+
+        if messages.get("warnings"):
+            self.warn("\n".join(messages["warnings"]))
+
+        if messages.get("info"):
+            self.info("\n".join(messages["info"]))
+
 
 class ClickIOProviderException(PlatformException):
     pass

dbt_platform_helper/providers/kms.py

@@ -0,0 +1,22 @@
+import boto3
+
+
+class KMSProvider:
+    """A provider class for interacting with the AWS KMS (Key Management
+    Service)."""
+
+    def __init__(self, kms_client: boto3.client):
+        self.kms_client = kms_client
+
+    def describe_key(self, alias_name: str) -> dict:
+        """
+        Retrieves metadata about a KMS key using its alias.
+
+        Args:
+            alias_name (str): The alias name of the KMS key.
+
+        Returns:
+            dict: A dictionary containing metadata about the specified KMS key.
+        """
+        # The kms client can take an alias name as the KeyId
+        return self.kms_client.describe_key(KeyId=alias_name)