aws-annoying 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

aws_annoying/cli/app.py

@@ -1,6 +1,13 @@
 from __future__ import annotations
 
+import importlib.metadata
+import logging
+import logging.config
+from typing import Optional
+
 import typer
+from rich import print  # noqa: A004
+from rich.console import Console
 
 app = typer.Typer(
     pretty_exceptions_short=True,
@@ -8,3 +15,77 @@ app = typer.Typer(
     rich_markup_mode="rich",
     no_args_is_help=True,
 )
+
+
+def show_version(value: Optional[bool]) -> None:
+    """Show the version of the application."""
+    if not value:
+        return
+
+    print(importlib.metadata.version("aws-annoying"))
+    raise typer.Exit(0)
+
+
+@app.callback()
+def main(  # noqa: D103
+    ctx: typer.Context,
+    *,
+    version: Optional[bool] = typer.Option(  # noqa: ARG001
+        None,
+        "--version",
+        is_eager=True,
+        callback=show_version,
+        help="Show the version and exit.",
+    ),
+    quiet: bool = typer.Option(
+        False,  # noqa: FBT003
+        help="Disable outputs.",
+    ),
+    verbose: bool = typer.Option(
+        False,  # noqa: FBT003
+        help="Enable verbose outputs.",
+    ),
+    dry_run: bool = typer.Option(
+        False,  # noqa: FBT003
+        help="Enable dry-run mode. If enabled, certain commands will avoid making changes.",
+    ),
+) -> None:
+    log_level = logging.DEBUG if verbose else logging.INFO
+    console = Console(soft_wrap=True, emoji=False)
+    logging_config: logging.config._DictConfigArgs = {
+        "version": 1,
+        "disable_existing_loggers": False,
+        "formatters": {
+            "rich": {
+                "format": "%(message)s",
+                "datefmt": "[%X]",
+            },
+        },
+        "handlers": {
+            "null": {
+                "class": "logging.NullHandler",
+            },
+            "rich": {
+                "class": "aws_annoying.cli.logging_handler.RichLogHandler",
+                "formatter": "rich",
+                "console": console,
+            },
+        },
+        "root": {
+            "handlers": ["null"],
+        },
+        "loggers": {
+            "aws_annoying": {
+                "level": log_level,
+                "handlers": ["rich"],
+                "propagate": True,
+            },
+        },
+    }
+    if quiet:
+        logging_config["loggers"]["aws_annoying"]["level"] = logging.CRITICAL
+
+    logging.config.dictConfig(logging_config)
+
+    # Global flags
+    ctx.meta["dry_run"] = dry_run

aws_annoying/cli/ecs/__init__.py

@@ -0,0 +1,3 @@
+from . import task_definition_lifecycle, wait_for_deployment
+
+__all__ = ("task_definition_lifecycle", "wait_for_deployment")

aws_annoying/cli/ecs/_app.py

@@ -0,0 +1,9 @@
+import typer
+
+from aws_annoying.cli.app import app
+
+ecs_app = typer.Typer(
+    no_args_is_help=True,
+    help="ECS (Elastic Container Service) utility commands.",
+)
+app.add_typer(ecs_app, name="ecs")

aws_annoying/cli/{ecs_task_definition_lifecycle.py → ecs/task_definition_lifecycle.py}

@@ -1,21 +1,25 @@
 from __future__ import annotations
 
+import logging
 from typing import TYPE_CHECKING
 
 import boto3
 import typer
-from rich import print  # noqa: A004
 
-from .app import app
+from ._app import ecs_app
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
 
+
+logger = logging.getLogger(__name__)
+
 _DELETE_CHUNK_SIZE = 10
 
 
-@app.command()
-def ecs_task_definition_lifecycle(
+@ecs_app.command()
+def task_definition_lifecycle(
+    ctx: typer.Context,
     *,
     family: str = typer.Option(
         ...,
@@ -33,14 +37,11 @@ def ecs_task_definition_lifecycle(
         False,  # noqa: FBT003
         help="Delete the task definition after deregistering it.",
     ),
-    dry_run: bool = typer.Option(
-        False,  # noqa: FBT003
-        help="Do not perform any changes, only show what would be done.",
-    ),
 ) -> None:
     """Execute ECS task definition lifecycle."""
+    dry_run = ctx.meta["dry_run"]
     if dry_run:
-        print("⚠️ Dry run mode enabled. Will not perform any actual changes.")
+        logger.info("Dry run mode enabled. Will not perform any actual changes.")
 
     ecs = boto3.client("ecs")
 
@@ -59,23 +60,27 @@ def ecs_task_definition_lifecycle(
 
     # Keep the latest N task definitions
     expired_taskdef_arns = task_definition_arns[:-keep_latest]
-    print(f"⚠️ Deregistering {len(expired_taskdef_arns)} task definitions...")
+    logger.warning("Deregistering %d task definitions...", len(expired_taskdef_arns))
     for arn in expired_taskdef_arns:
         if not dry_run:
             ecs.deregister_task_definition(taskDefinition=arn)
 
         # ARN like: "arn:aws:ecs:<region>:<account-id>:task-definition/<family>:<revision>"
         _, family_revision = arn.split(":task-definition/")
-        print(f"✅ Deregistered task definition [yellow]{family_revision!r}[/yellow]")
+        logger.warning("Deregistered task definition [yellow]%r[/yellow]", family_revision)
 
     if delete and expired_taskdef_arns:
         # Delete the expired task definitions in chunks due to API limitation
-        print(f"⚠️ Deleting {len(expired_taskdef_arns)} task definitions in chunks of size {_DELETE_CHUNK_SIZE}...")
+        logger.warning(
+            "Deleting %d task definitions in chunks of size %d...",
+            len(expired_taskdef_arns),
+            _DELETE_CHUNK_SIZE,
+        )
         for idx, chunk in enumerate(_chunker(expired_taskdef_arns, _DELETE_CHUNK_SIZE)):
             if not dry_run:
                 ecs.delete_task_definitions(taskDefinitions=chunk)
 
-            print(f"✅ Deleted {len(chunk)} task definitions in {idx}-th batch.")
+            logger.warning("Deleted %d task definitions in %d-th batch.", len(chunk), idx)
 
 
 def _chunker(sequence: list, size: int) -> Iterator[list]:

aws_annoying/cli/ecs/wait_for_deployment.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+import logging
+from datetime import datetime, timezone
+from typing import Optional
+
+import typer
+
+from aws_annoying.ecs import (
+    DeploymentFailedError,
+    ECSServiceRef,
+    ServiceTaskDefinitionAssertionError,
+    check_service_task_definition,
+    wait_for_deployment_complete,
+    wait_for_deployment_start,
+    wait_for_service_stability,
+)
+from aws_annoying.utils.timeout import OperationTimeoutError, Timeout
+
+from ._app import ecs_app
+
+logger = logging.getLogger(__name__)
+
+
+@ecs_app.command()
+def wait_for_deployment(  # noqa: PLR0913
+    *,
+    cluster: str = typer.Option(
+        ...,
+        help="The name of the ECS cluster.",
+        show_default=False,
+    ),
+    service: str = typer.Option(
+        ...,
+        help="The name of the ECS service.",
+        show_default=False,
+    ),
+    expected_task_definition: Optional[str] = typer.Option(
+        None,
+        help=(
+            "The service's task definition expected after deployment."
+            " If provided, it will be used to assert the service's task definition after deployment finished or timed out."  # noqa: E501
+        ),
+        show_default=False,
+    ),
+    polling_interval: int = typer.Option(
+        5,
+        help="The interval between any polling attempts, in seconds.",
+        min=1,
+    ),
+    timeout_seconds: Optional[int] = typer.Option(
+        None,
+        help=(
+            "The maximum time to wait for the deployment to complete, in seconds."
+            " If not provided, it will wait indefinitely."
+        ),
+        min=1,
+    ),
+    wait_for_start: bool = typer.Option(
+        True,  # noqa: FBT003
+        help=(
+            "Whether to wait for the deployment to start."
+            " Because there could be no deployment right after the deploy,"
+            " this option will wait for a new deployment to start if no running deployment is found."
+        ),
+    ),
+    wait_for_stability: bool = typer.Option(
+        False,  # noqa: FBT003
+        help="Whether to wait for the service to be stable after the deployment.",
+    ),
+) -> None:
+    """Wait for ECS deployment to complete."""
+    start = datetime.now(tz=timezone.utc)
+    try:
+        with Timeout(timeout_seconds):
+            _wait_for_deployment(
+                ECSServiceRef(cluster=cluster, service=service),
+                wait_for_start=wait_for_start,
+                polling_interval=polling_interval,
+                wait_for_stability=wait_for_stability,
+                expected_task_definition=expected_task_definition,
+            )
+    except OperationTimeoutError:
+        logger.error(  # noqa: TRY400
+            "Timeout reached after %s seconds. The deployment may not have finished.",
+            timeout_seconds,
+        )
+        raise typer.Exit(1) from None
+    except DeploymentFailedError as err:
+        elapsed = datetime.now(tz=timezone.utc) - start
+        logger.error(  # noqa: TRY400
+            "Deployment failed in [bold]%.2f[/bold] seconds with error: %s",
+            elapsed.total_seconds(),
+            err,
+        )
+        raise typer.Exit(1) from None
+    else:
+        elapsed = datetime.now(tz=timezone.utc) - start
+        logger.info(
+            "Deployment completed in [bold]%.2f[/bold] seconds.",
+            elapsed.total_seconds(),
+        )
+
+
+def _wait_for_deployment(
+    service_ref: ECSServiceRef,
+    *,
+    wait_for_start: bool,
+    polling_interval: int = 5,
+    wait_for_stability: bool,
+    expected_task_definition: str | None = None,
+) -> None:
+    # Find current deployment for the service
+    logger.info(
+        "Looking up running deployment for service %s",
+        service_ref.service,
+    )
+    latest_deployment_arn = wait_for_deployment_start(
+        service_ref,
+        wait_for_start=wait_for_start,
+        polling_interval=polling_interval,
+    )
+
+    # Polling for the deployment to finish (successfully or unsuccessfully)
+    logger.info(
+        "Start waiting for deployment %s to finish.",
+        latest_deployment_arn,
+    )
+    ok, status = wait_for_deployment_complete(latest_deployment_arn, polling_interval=polling_interval)
+    if ok:
+        logger.info(
+            "Deployment succeeded with status %s",
+            status,
+        )
+    else:
+        msg = f"Deployment failed with status: {status}"
+        raise DeploymentFailedError(msg)
+
+    # Wait for the service to be stable
+    if wait_for_stability:
+        logger.info(
+            "Start waiting for service %s to be stable.",
+            service_ref.service,
+        )
+        wait_for_service_stability(service_ref, polling_interval=polling_interval)
+
+    # Check if the service task definition matches the expected one
+    if expected_task_definition:
+        logger.info(
+            "Checking if the service task definition is the expected one: %s",
+            expected_task_definition,
+        )
+        ok, actual = check_service_task_definition(service_ref, expect=expected_task_definition)
+        if not ok:
+            msg = f"The service task definition is not the expected one; got: {actual!r}"
+            raise ServiceTaskDefinitionAssertionError(msg)
+
+        logger.info("The service task definition matches the expected one.")

aws_annoying/cli/load_variables.py

@@ -1,18 +1,22 @@
 # flake8: noqa: B008
 from __future__ import annotations
 
+import logging
 import os
 import subprocess
+from io import StringIO
 from typing import NoReturn, Optional
 
 import typer
 from rich.console import Console
 from rich.table import Table
 
-from aws_annoying.variables import VariableLoader
+from aws_annoying.variable_loader import VariableLoader
 
 from .app import app
 
+logger = logging.getLogger(__name__)
+
 
 @app.command(
     context_settings={
@@ -21,9 +25,9 @@ from .app import app
         "ignore_unknown_options": True,
     },
 )
-def load_variables(  # noqa: PLR0913
-    *,
+def load_variables(
     ctx: typer.Context,
+    *,
     arns: list[str] = typer.Option(
         [],
         metavar="ARN",
@@ -42,14 +46,6 @@ def load_variables(  # noqa: PLR0913
         False,  # noqa: FBT003
         help="Overwrite the existing environment variables with the same name.",
     ),
-    quiet: bool = typer.Option(
-        False,  # noqa: FBT003
-        help="Suppress all outputs from this command.",
-    ),
-    dry_run: bool = typer.Option(
-        False,  # noqa: FBT003
-        help="Print the progress only. Neither load variables nor run the command.",
-    ),
     replace: bool = typer.Option(
         True,  # noqa: FBT003
         help=(
@@ -82,21 +78,19 @@ def load_variables(  # noqa: PLR0913
     The variables are loaded in the order of option provided, overwriting the variables with the same name in the order of the ARNs.
     Existing environment variables are preserved by default, unless `--overwrite-env` is provided.
     """  # noqa: E501
-    console = Console(quiet=quiet, emoji=False)
-
     command = ctx.args
     if not command:
-        console.print("⚠️ No command provided. Exiting...")
+        logger.warning("No command provided. Exiting...")
         raise typer.Exit(0)
 
     # Mapping of the ARNs by index (index used for ordering)
     map_arns_by_index = {str(idx): arn for idx, arn in enumerate(arns)}
     if env_prefix:
-        console.print(f"🔍 Loading ARNs from environment variables with prefix: {env_prefix!r}")
+        logger.info("Loading ARNs from environment variables with prefix: %r", env_prefix)
         arns_env = {
             key.removeprefix(env_prefix): value for key, value in os.environ.items() if key.startswith(env_prefix)
         }
-        console.print(f"🔍 Found {len(arns_env)} sources from environment variables.")
+        logger.info("Found %d sources from environment variables.", len(arns_env))
         map_arns_by_index = arns_env | map_arns_by_index
 
     # Briefly show the ARNs
@@ -104,21 +98,22 @@ def load_variables(  # noqa: PLR0913
     for idx, arn in sorted(map_arns_by_index.items()):
         table.add_row(idx, arn)
 
-    console.print(table)
+    # Workaround: The logger cannot directly handle the rich table output.
+    with StringIO() as file:
+        Console(file=file, emoji=False).print(table)
+        table_str = file.getvalue().rstrip()
+        logger.info("Summary:\n%s", table_str)
 
     # Retrieve the variables
-    loader = VariableLoader(dry_run=dry_run)
-    console.print("🔍 Retrieving variables from AWS resources...")
-    if dry_run:
-        console.print("⚠️ Dry run mode enabled. Variables won't be loaded from AWS.")
-
+    loader = VariableLoader()
+    logger.info("Retrieving variables from AWS resources...")
     try:
         variables, load_stats = loader.load(map_arns_by_index)
     except Exception as exc:  # noqa: BLE001
-        console.print(f"❌ Failed to load the variables: {exc!s}")
+        logger.error("Failed to load the variables: %s", exc)  # noqa: TRY400
         raise typer.Exit(1) from None
 
-    console.print(f"✅ Retrieved {load_stats['secrets']} secrets and {load_stats['parameters']} parameters.")
+    logger.info("Retrieved %d secrets and %d parameters.", load_stats["secrets"], load_stats["parameters"])
 
     # Prepare the environment variables
     env = os.environ.copy()
@@ -130,7 +125,7 @@ def load_variables(  # noqa: PLR0913
             env.setdefault(key, str(value))
 
     # Run the command with the variables injected as environment variables, replacing current process
-    console.print(f"🚀 Running the command: [bold orchid]{' '.join(command)}[/bold orchid]")
+    logger.info("Running the command: [bold orchid]%s[/bold orchid]", " ".join(command))
     if replace:  # pragma: no cover (not coverable)
         os.execvpe(command[0], command, env=env)  # noqa: S606
         # The above line should never return

aws_annoying/cli/logging_handler.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import logging
+import logging.config
+from typing import TYPE_CHECKING, Any, Final
+
+from typing_extensions import override
+
+if TYPE_CHECKING:
+    from rich.console import Console
+
+
+class RichLogHandler(logging.Handler):
+    """Custom logging handler to use Rich Console."""
+
+    _level_emojis: Final[dict[str, str]] = {
+        "DEBUG": "🔍",
+        "INFO": "ℹ️",  # noqa: RUF001
+        "WARNING": "⚠️",
+        "ERROR": "❗",
+        "CRITICAL": "🚨",
+    }
+
+    def __init__(self, console: Console, *args: Any, **kwargs: Any) -> None:
+        """Initialize the log handler.
+
+        Args:
+            console: Rich console instance.
+            *args: Additional arguments for the logging handler.
+            **kwargs: Additional keyword arguments for the logging handler.
+        """
+        super().__init__(*args, **kwargs)
+        self.console = console
+
+    @override
+    def emit(self, record: logging.LogRecord) -> None:
+        msg = self.format(record)
+        self.console.print(msg)
+
+    @override
+    def format(self, record: logging.LogRecord) -> str:
+        """Format the log record.
+
+        Args:
+            record: The log record to format.
+
+        Returns:
+            The formatted log message.
+        """
+        msg = super().format(record)
+        emoji = self._level_emojis.get(record.levelname)
+        return f"{emoji} {msg}" if emoji else msg

aws_annoying/cli/main.py

@@ -1,7 +1,7 @@
 # flake8: noqa: F401
 from __future__ import annotations
 
-import aws_annoying.cli.ecs_task_definition_lifecycle
+import aws_annoying.cli.ecs
 import aws_annoying.cli.load_variables
 import aws_annoying.cli.mfa
 import aws_annoying.cli.session_manager

aws_annoying/cli/mfa/configure.py

@@ -1,18 +1,18 @@
 from __future__ import annotations
 
+import logging
 from pathlib import Path  # noqa: TC003
 from typing import Optional
 
 import boto3
 import typer
-from rich import print  # noqa: A004
 from rich.prompt import Prompt
 
-from aws_annoying.mfa import MfaConfig, update_credentials
+from aws_annoying.mfa_config import MfaConfig, update_credentials
 
 from ._app import mfa_app
 
-_CONFIG_INI_SECTION = "aws-annoying:mfa"
+logger = logging.getLogger(__name__)
 
 
 @mfa_app.command()
@@ -44,6 +44,10 @@ def configure(  # noqa: PLR0913
         "~/.aws/config",
         help="The path to the AWS config file. Used to persist the MFA configuration.",
     ),
+    aws_config_section: str = typer.Option(
+        "aws-annoying:mfa",
+        help="The section in the AWS config file to persist the MFA configuration.",
+    ),
     persist: bool = typer.Option(
         True,  # noqa: FBT003
         help="Persist the MFA configuration.",
@@ -55,9 +59,9 @@ def configure(  # noqa: PLR0913
     aws_config = aws_config.expanduser()
 
     # Load configuration
-    mfa_config, exists = MfaConfig.from_ini_file(aws_config, _CONFIG_INI_SECTION)
+    mfa_config, exists = MfaConfig.from_ini_file(aws_config, aws_config_section)
     if exists:
-        print(f"⚙️ Loaded MFA configuration from AWS config ({aws_config}).")
+        logger.info("Loaded MFA configuration from AWS config (%s).", aws_config)
 
     mfa_profile = (
         mfa_profile
@@ -83,7 +87,7 @@ def configure(  # noqa: PLR0913
     )
 
     # Get credentials
-    print(f"💬 Retrieving MFA credentials using profile [bold]{mfa_source_profile}[/bold]")
+    logger.info("Retrieving MFA credentials using profile [bold]%s[/bold]", mfa_source_profile)
     session = boto3.session.Session(profile_name=mfa_source_profile)
     sts = session.client("sts")
     response = sts.get_session_token(
@@ -93,7 +97,11 @@ def configure(  # noqa: PLR0913
     credentials = response["Credentials"]
 
     # Update MFA profile in AWS credentials
-    print(f"✅ Updating MFA profile ([bold]{mfa_profile}[/bold]) to AWS credentials ({aws_credentials})")
+    logger.warning(
+        "Updating MFA profile ([bold]%s[/bold]) to AWS credentials ([bold]%s[/bold])",
+        mfa_profile,
+        aws_credentials,
+    )
     update_credentials(
         aws_credentials,
         mfa_profile,  # type: ignore[arg-type]
@@ -104,13 +112,14 @@ def configure(  # noqa: PLR0913
 
     # Persist MFA configuration
     if persist:
-        print(
-            f"✅ Persisting MFA configuration in AWS config ({aws_config}),"
-            f" in [bold]{_CONFIG_INI_SECTION}[/bold] section.",
+        logger.info(
+            "Persisting MFA configuration in AWS config (%s), in [bold]%s[/bold] section.",
+            aws_config,
+            aws_config_section,
         )
         mfa_config.mfa_profile = mfa_profile
         mfa_config.mfa_source_profile = mfa_source_profile
         mfa_config.mfa_serial_number = mfa_serial_number
-        mfa_config.save_ini_file(aws_config, _CONFIG_INI_SECTION)
+        mfa_config.save_ini_file(aws_config, aws_config_section)
     else:
-        print("⚠️ MFA configuration not persisted.")
+        logger.warning("MFA configuration not persisted.")

aws_annoying/cli/session_manager/_common.py

@@ -1,10 +1,7 @@
-# TODO(lasuillard): Using this file until split CLI from library codebase
 from __future__ import annotations
 
-import re
 from typing import Any
 
-import boto3
 import typer
 from rich.prompt import Confirm
 
@@ -17,38 +14,10 @@ class SessionManager(_SessionManager):
         if self._confirm:
             return
 
-        confirm = Confirm.ask(f"⚠️ Will run the following command: [bold red]{' '.join(command)}[/bold red]. Proceed?")
+        confirm = Confirm.ask(f"Will run the following command: [bold red]{' '.join(command)}[/bold red]. Proceed?")
         if not confirm:
             raise typer.Abort
 
     def install(self, *args: Any, confirm: bool = False, **kwargs: Any) -> None:
         self._confirm = confirm
         return super().install(*args, **kwargs)
-
-
-def get_instance_id_by_name(name_or_id: str) -> str | None:
-    """Get the EC2 instance ID by name or ID.
-
-    Be aware that this function will only return the first instance found
-    with the given name, no matter how many instances are found.
-
-    Args:
-        name_or_id: The name or ID of the EC2 instance.
-
-    Returns:
-        The instance ID if found, otherwise `None`.
-    """
-    if re.match(r"m?i-.+", name_or_id):
-        return name_or_id
-
-    ec2 = boto3.client("ec2")
-    response = ec2.describe_instances(Filters=[{"Name": "tag:Name", "Values": [name_or_id]}])
-    reservations = response["Reservations"]
-    if not reservations:
-        return None
-
-    instances = reservations[0]["Instances"]
-    if not instances:
-        return None
-
-    return str(instances[0]["InstanceId"])