✨AWS-library: Added interface to AWS SSM (#6032)

Author: sanderegg

.coveragerc

@@ -26,5 +26,8 @@ exclude_lines =
     # Don't complain about abstract methods, they aren't run:
     @(abc\.)?abstract(((class|static)?method)|property)
 
+    # Don't complain about type checking
+    if TYPE_CHECKING:
+
 ignore_errors = True
 show_missing = True

Makefile

@@ -124,11 +124,8 @@ __check_defined = \
 .PHONY: help
 
 help: ## help on rule's targets
-ifeq ($(IS_WIN),)
-	@awk --posix 'BEGIN {FS = ":.*?## "} /^[[:alpha:][:space:]_-]+:.*?## / {printf "\033[36m%-20s\033[0m %s\n", $$1, $$2}' $(MAKEFILE_LIST)
-else
-	@awk --posix 'BEGIN {FS = ":.*?## "} /^[[:alpha:][:space:]_-]+:.*?## / {printf "%-20s %s\n", $$1, $$2}' $(MAKEFILE_LIST)
-endif
+	@awk 'BEGIN {FS = ":.*?## "}; /^[^.[:space:]].*?:.*?## / {if ($$1 != "help" && NF == 2) {printf "\033[36m%-20s\033[0m %s\n", $$1, $$2}}' $(MAKEFILE_LIST)
+
 
 test_python_version: ## Check Python version, throw error if compilation would fail with the installed version
 	python ./scripts/test_python_version.py
@@ -459,14 +456,13 @@ push-version: tag-version
 .check-uv-installed:
 		@echo "Checking if 'uv' is installed..."
 		@if ! command -v uv >/dev/null 2>&1; then \
-				printf "\033[31mError: 'uv' is not installed.\033[0m\n"; \
-				printf "To install 'uv', run the following command:\n"; \
-				printf "\033[34mcurl -LsSf https://astral.sh/uv/install.sh | sh\033[0m\n"; \
-				exit 1; \
+				curl -LsSf https://astral.sh/uv/install.sh | sh; \
 		else \
 				printf "\033[32m'uv' is installed. Version: \033[0m"; \
 				uv --version; \
 		fi
+		# upgrading uv
+		-@uv self --quiet update
 
 
 .venv: .check-uv-installed

packages/aws-library/requirements/_base.in

@@ -9,5 +9,5 @@
 aioboto3
 aiocache
 pydantic[email]
-types-aiobotocore[ec2,s3]
+types-aiobotocore[ec2,s3,ssm]
 sh

packages/aws-library/requirements/_base.txt

@@ -184,11 +184,13 @@ typer==0.12.3
     #   -r requirements/../../../packages/service-library/requirements/../../../packages/settings-library/requirements/_base.in
     #   -r requirements/../../../packages/settings-library/requirements/_base.in
     #   faststream
-types-aiobotocore==2.12.3
+types-aiobotocore==2.13.0
     # via -r requirements/_base.in
-types-aiobotocore-ec2==2.12.3
+types-aiobotocore-ec2==2.13.0
     # via types-aiobotocore
-types-aiobotocore-s3==2.12.3
+types-aiobotocore-s3==2.13.0
+    # via types-aiobotocore
+types-aiobotocore-ssm==2.13.0
     # via types-aiobotocore
 types-awscrt==0.20.9
     # via botocore-stubs
@@ -205,6 +207,7 @@ typing-extensions==4.11.0
     #   types-aiobotocore
     #   types-aiobotocore-ec2
     #   types-aiobotocore-s3
+    #   types-aiobotocore-ssm
 urllib3==2.2.1
     # via
     #   -c requirements/../../../packages/models-library/requirements/../../../requirements/constraints.txt

packages/aws-library/src/aws_library/ec2/client.py

@@ -1,6 +1,6 @@
 import contextlib
 import logging
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
 from typing import cast
 
@@ -22,6 +22,7 @@
     EC2TooManyInstancesError,
 )
 from .models import (
+    AWSTagKey,
     EC2InstanceConfig,
     EC2InstanceData,
     EC2InstanceType,
@@ -274,6 +275,24 @@ async def get_instances(
         )
         return all_instances
 
+    async def stop_instances(self, instance_datas: Iterable[EC2InstanceData]) -> None:
+        try:
+            with log_context(
+                _logger,
+                logging.INFO,
+                msg=f"stopping instances {[i.id for i in instance_datas]}",
+            ):
+                await self.client.stop_instances(
+                    InstanceIds=[i.id for i in instance_datas]
+                )
+        except botocore.exceptions.ClientError as exc:
+            if (
+                exc.response.get("Error", {}).get("Code", "")
+                == "InvalidInstanceID.NotFound"
+            ):
+                raise EC2InstanceNotFoundError from exc
+            raise  # pragma: no cover
+
     async def terminate_instances(
         self, instance_datas: Iterable[EC2InstanceData]
     ) -> None:
@@ -295,7 +314,7 @@ async def terminate_instances(
             raise  # pragma: no cover
 
     async def set_instances_tags(
-        self, instances: list[EC2InstanceData], *, tags: EC2Tags
+        self, instances: Sequence[EC2InstanceData], *, tags: EC2Tags
     ) -> None:
         try:
             with log_context(
@@ -314,3 +333,21 @@ async def set_instances_tags(
             if exc.response.get("Error", {}).get("Code", "") == "InvalidID":
                 raise EC2InstanceNotFoundError from exc
             raise  # pragma: no cover
+
+    async def remove_instances_tags(
+        self, instances: Sequence[EC2InstanceData], *, tag_keys: Iterable[AWSTagKey]
+    ) -> None:
+        try:
+            with log_context(
+                _logger,
+                logging.DEBUG,
+                msg=f"removing {tag_keys=} of instances '[{[i.id for i in instances]}]'",
+            ):
+                await self.client.delete_tags(
+                    Resources=[i.id for i in instances],
+                    Tags=[{"Key": tag_key} for tag_key in tag_keys],
+                )
+        except botocore.exceptions.ClientError as exc:
+            if exc.response.get("Error", {}).get("Code", "") == "InvalidID":
+                raise EC2InstanceNotFoundError from exc
+            raise  # pragma: no cover

packages/aws-library/src/aws_library/ec2/models.py

@@ -10,7 +10,6 @@
     BaseModel,
     ByteSize,
     ConstrainedStr,
-    Extra,
     Field,
     NonNegativeFloat,
     validator,
@@ -142,7 +141,6 @@ class EC2InstanceBootSpecific(BaseModel):
     )
 
     class Config:
-        extra = Extra.forbid
         schema_extra: ClassVar[dict[str, Any]] = {
             "examples": [
                 {

packages/aws-library/src/aws_library/ssm/__init__.py

@@ -0,0 +1,21 @@
+from ._client import SimcoreSSMAPI
+from ._errors import (
+    SSMAccessError,
+    SSMCommandExecutionError,
+    SSMInvalidCommandError,
+    SSMNotConnectedError,
+    SSMRuntimeError,
+    SSMSendCommandInstancesNotReadyError,
+)
+
+__all__: tuple[str, ...] = (
+    "SimcoreSSMAPI",
+    "SSMAccessError",
+    "SSMNotConnectedError",
+    "SSMRuntimeError",
+    "SSMSendCommandInstancesNotReadyError",
+    "SSMInvalidCommandError",
+    "SSMCommandExecutionError",
+)
+
+# nopycln: file

packages/aws-library/src/aws_library/ssm/_client.py

@@ -0,0 +1,159 @@
+import contextlib
+import logging
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Final, cast
+
+import aioboto3
+import botocore
+import botocore.exceptions
+from aiobotocore.session import ClientCreatorContext
+from servicelib.logging_utils import log_decorator
+from settings_library.ssm import SSMSettings
+from types_aiobotocore_ssm import SSMClient
+from types_aiobotocore_ssm.literals import CommandStatusType
+
+from ._error_handler import ssm_exception_handler
+from ._errors import SSMCommandExecutionError
+
+_logger = logging.getLogger(__name__)
+
+_AWS_WAIT_MAX_DELAY: Final[int] = 5
+_AWS_WAIT_NUM_RETRIES: Final[int] = 3
+
+
+@dataclass(frozen=True)
+class SSMCommand:
+    name: str
+    command_id: str
+    instance_ids: Sequence[str]
+    status: CommandStatusType
+    message: str | None = None
+
+
+@dataclass(frozen=True)
+class SimcoreSSMAPI:
+    _client: SSMClient
+    _session: aioboto3.Session
+    _exit_stack: contextlib.AsyncExitStack
+
+    @classmethod
+    async def create(cls, settings: SSMSettings) -> "SimcoreSSMAPI":
+        session = aioboto3.Session()
+        session_client = session.client(
+            "ssm",
+            endpoint_url=f"{settings.SSM_ENDPOINT}",
+            aws_access_key_id=settings.SSM_ACCESS_KEY_ID.get_secret_value(),
+            aws_secret_access_key=settings.SSM_SECRET_ACCESS_KEY.get_secret_value(),
+            region_name=settings.SSM_REGION_NAME,
+        )
+        assert isinstance(session_client, ClientCreatorContext)  # nosec
+        exit_stack = contextlib.AsyncExitStack()
+        ec2_client = cast(
+            SSMClient, await exit_stack.enter_async_context(session_client)
+        )
+        return cls(ec2_client, session, exit_stack)
+
+    async def close(self) -> None:
+        await self._exit_stack.aclose()
+
+    async def ping(self) -> bool:
+        try:
+            await self._client.list_commands(MaxResults=1)
+            return True
+        except Exception:  # pylint: disable=broad-except
+            return False
+
+    # a function to send a command via ssm
+    @log_decorator(_logger, logging.DEBUG)
+    @ssm_exception_handler(_logger)
+    async def send_command(
+        self, instance_ids: Sequence[str], *, command: str, command_name: str
+    ) -> SSMCommand:
+        # NOTE: using Targets instead of instances as this is limited to 50 instances
+        # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ssm.html#SSM.Client.send_command
+        response = await self._client.send_command(
+            Targets=[{"Key": "InstanceIds", "Values": instance_ids}],
+            DocumentName="AWS-RunShellScript",
+            Comment=command_name,
+            Parameters={"commands": [command]},
+        )
+        assert response["Command"]  # nosec
+        assert "Comment" in response["Command"]  # nosec
+        assert "CommandId" in response["Command"]  # nosec
+        assert "Status" in response["Command"]  # nosec
+
+        return SSMCommand(
+            name=response["Command"]["Comment"],
+            command_id=response["Command"]["CommandId"],
+            status=response["Command"]["Status"],
+            instance_ids=instance_ids,
+        )
+
+    @log_decorator(_logger, logging.DEBUG)
+    @ssm_exception_handler(_logger)
+    async def get_command(self, instance_id: str, *, command_id: str) -> SSMCommand:
+
+        response = await self._client.get_command_invocation(
+            CommandId=command_id, InstanceId=instance_id
+        )
+
+        return SSMCommand(
+            name=response["Comment"],
+            command_id=response["CommandId"],
+            instance_ids=[response["InstanceId"]],
+            status=response["Status"] if response["Status"] != "Delayed" else "Pending",
+            message=response["StatusDetails"],
+        )
+
+    @log_decorator(_logger, logging.DEBUG)
+    @ssm_exception_handler(_logger)
+    async def is_instance_connected_to_ssm_server(self, instance_id: str) -> bool:
+        response = await self._client.describe_instance_information(
+            InstanceInformationFilterList=[
+                {
+                    "key": "InstanceIds",
+                    "valueSet": [
+                        instance_id,
+                    ],
+                }
+            ],
+        )
+        assert response["InstanceInformationList"]  # nosec
+        assert len(response["InstanceInformationList"]) == 1  # nosec
+        assert "PingStatus" in response["InstanceInformationList"][0]  # nosec
+        return bool(response["InstanceInformationList"][0]["PingStatus"] == "Online")
+
+    @log_decorator(_logger, logging.DEBUG)
+    @ssm_exception_handler(_logger)
+    async def wait_for_has_instance_completed_cloud_init(
+        self, instance_id: str
+    ) -> bool:
+        cloud_init_status_command = await self.send_command(
+            (instance_id,),
+            command="cloud-init status",
+            command_name="cloud-init status",
+        )
+        # wait for command to complete
+        waiter = self._client.get_waiter(  # pylint: disable=assignment-from-no-return
+            "command_executed"
+        )
+        try:
+            await waiter.wait(
+                CommandId=cloud_init_status_command.command_id,
+                InstanceId=instance_id,
+                WaiterConfig={
+                    "Delay": _AWS_WAIT_MAX_DELAY,
+                    "MaxAttempts": _AWS_WAIT_NUM_RETRIES,
+                },
+            )
+        except botocore.exceptions.WaiterError as exc:
+            msg = f"Timed-out waiting for {instance_id} to complete cloud-init"
+            raise SSMCommandExecutionError(details=msg) from exc
+        response = await self._client.get_command_invocation(
+            CommandId=cloud_init_status_command.command_id, InstanceId=instance_id
+        )
+        if response["Status"] != "Success":
+            raise SSMCommandExecutionError(details=response["StatusDetails"])
+        # check if cloud-init is done
+        return bool("status: done" in response["StandardOutputContent"])