Report progress (#64)

Author: jaymccon

src/cloudformation_cli_python_lib/callback.py

@@ -0,0 +1,45 @@
+import json
+import logging
+from typing import Optional
+from uuid import uuid4
+
+# boto3 doesn't have stub files
+from boto3 import Session  # type: ignore
+
+from .interface import BaseResourceModel, HandlerErrorCode, OperationStatus
+from .utils import KitchenSinkEncoder
+
+LOG = logging.getLogger(__name__)
+
+
+def report_progress(  # pylint: disable=too-many-arguments
+    session: Session,
+    bearer_token: str,
+    error_code: Optional[HandlerErrorCode],
+    operation_status: OperationStatus,
+    current_operation_status: Optional[OperationStatus],
+    resource_model: Optional[BaseResourceModel],
+    status_message: str,
+) -> None:
+    client = session.client("cloudformation")
+    request = {
+        "BearerToken": bearer_token,
+        "OperationStatus": operation_status.name,
+        "StatusMessage": status_message,
+        "ClientRequestToken": str(uuid4()),
+    }
+    if resource_model:
+        request["ResourceModel"] = json.dumps(
+            resource_model._serialize(),  # pylint: disable=protected-access
+            cls=KitchenSinkEncoder,
+        )
+    if error_code:
+        request["ErrorCode"] = error_code.name
+    if current_operation_status:
+        request["CurrentOperationStatus"] = current_operation_status.name
+    response = client.record_handler_progress(**request)
+    LOG.info(
+        "Record Handler Progress with Request Id %s and Request: {%s}",
+        response["ResponseMetadata"]["RequestId"],
+        request,
+    )

src/cloudformation_cli_python_lib/interface.py

@@ -75,7 +75,7 @@ class ProgressEvent:
     status: OperationStatus
     errorCode: Optional[HandlerErrorCode] = None
     message: str = ""
-    callbackContext: Optional[Mapping[str, Any]] = None
+    callbackContext: Optional[MutableMapping[str, Any]] = None
     callbackDelaySeconds: int = 0
     resourceModel: Optional[BaseResourceModel] = None
     resourceModels: Optional[List[BaseResourceModel]] = None
@@ -90,9 +90,21 @@ def _serialize(
         # mutate to what's expected in the response
         if to_response:
             ser["bearerToken"] = bearer_token
-            ser["operationStatus"] = ser.pop("status")
-            if ser["callbackDelaySeconds"] == 0:
-                del ser["callbackDelaySeconds"]
+            ser["operationStatus"] = ser.pop("status").name
+            if self.resourceModel:
+                # pylint: disable=protected-access
+                ser["resourceModel"] = self.resourceModel._serialize()
+            if self.resourceModels:
+                ser["resourceModels"] = [
+                    # pylint: disable=protected-access
+                    model._serialize()
+                    for model in self.resourceModels
+                ]
+            del ser["callbackDelaySeconds"]
+            if "callbackContext" in ser:
+                del ser["callbackContext"]
+            if self.errorCode:
+                ser["errorCode"] = self.errorCode.name
         return ser
 
     @classmethod

src/cloudformation_cli_python_lib/log_delivery.py

@@ -1,6 +1,6 @@
 import logging
 import time
-from typing import Any, Mapping
+from typing import Any, Mapping, Optional
 
 # boto3 doesn't have stub files
 import boto3  # type: ignore
@@ -28,6 +28,13 @@ def __init__(
         self.client = boto3.client("logs", **creds)
         self.sequence_token = ""
 
+    @classmethod
+    def _get_existing_logger(cls) -> Optional["ProviderLogHandler"]:
+        for handler in logging.getLogger().handlers:
+            if isinstance(handler, cls):
+                return handler
+        return None
+
     @classmethod
     def setup(cls, event_data: Mapping[str, Any]) -> None:
         try:
@@ -46,12 +53,23 @@ def setup(cls, event_data: Mapping[str, Any]) -> None:
         except KeyError:
             stream_name = f'{event_data["awsAccountId"]}-{event_data["region"]}'
 
-        # filter provider messages from platform
-        ProviderFilter.PROVIDER = event_data["resourceType"].replace("::", "_").lower()
-        logging.getLogger().handlers[0].addFilter(ProviderFilter())
-
-        # add log handler to root, so that provider gets plugin logs too
+        log_handler = cls._get_existing_logger()
         if log_creds and log_group:
+            if log_handler:
+                # This is a re-used lambda container, log handler is already setup, so
+                # we just refresh the client with new creds
+                log_handler.client = boto3.client(
+                    "logs",
+                    aws_access_key_id=log_creds["accessKeyId"],
+                    aws_secret_access_key=log_creds["secretAccessKey"],
+                    aws_session_token=log_creds["sessionToken"],
+                )
+                return
+            # filter provider messages from platform
+            ProviderFilter.PROVIDER = (
+                event_data["resourceType"].replace("::", "_").lower()
+            )
+            logging.getLogger().handlers[0].addFilter(ProviderFilter())
             log_handler = cls(
                 group=log_group,
                 stream=stream_name,
@@ -61,6 +79,7 @@ def setup(cls, event_data: Mapping[str, Any]) -> None:
                     "aws_session_token": log_creds["sessionToken"],
                 },
             )
+            # add log handler to root, so that provider gets plugin logs too
             logging.getLogger().addHandler(log_handler)
 
     def _create_log_group(self) -> None:

src/cloudformation_cli_python_lib/resource.py

@@ -8,6 +8,7 @@
 import boto3  # type: ignore
 
 from .boto3_proxy import SessionProxy, _get_boto_session
+from .callback import report_progress
 from .exceptions import InternalFailure, InvalidRequest, _HandlerError
 from .interface import (
     Action,
@@ -211,9 +212,40 @@ def __call__(
             ProviderLogHandler.setup(event_data)
             parsed = self._parse_request(event_data)
             caller_sess, platform_sess, request, action, callback, event = parsed
+            # Acknowledge the task for first time invocation
+            if not event.requestContext:
+                report_progress(
+                    platform_sess,
+                    event.bearerToken,
+                    None,
+                    OperationStatus.IN_PROGRESS,
+                    OperationStatus.PENDING,
+                    None,
+                    "",
+                )
+            else:
+                # If this invocation was triggered by a 're-invoke' CloudWatch Event,
+                # clean it up
+                CloudWatchScheduler(platform_sess).cleanup_cloudwatch_events(
+                    event.requestContext.get("cloudWatchEventsRuleName", ""),
+                    event.requestContext.get("cloudWatchEventsTargetId", ""),
+                )
             invoke = True
             while invoke:
                 progress = self._invoke_handler(caller_sess, request, action, callback)
+                if progress.callbackContext:
+                    callback = progress.callbackContext
+                    event.requestContext["callbackContext"] = callback
+                if event.action in MUTATING_ACTIONS:
+                    report_progress(
+                        platform_sess,
+                        event.bearerToken,
+                        progress.errorCode,
+                        progress.status,
+                        OperationStatus.IN_PROGRESS,
+                        progress.resourceModel,
+                        progress.message,
+                    )
                 invoke = self.schedule_reinvocation(
                     event, progress, context, platform_sess
                 )

src/cloudformation_cli_python_lib/scheduler.py

@@ -55,4 +55,4 @@ def _min_to_cron(minutes: int) -> str:
         schedule_time = datetime.now() + timedelta(minutes=minutes)
         # add another minute, as per java implementation
         schedule_time = schedule_time + timedelta(minutes=1)
-        return schedule_time.strftime("cron('%M %H %d %m ? %Y')")
+        return schedule_time.strftime("cron(%M %H %d %m ? %Y)")

src/cloudformation_cli_python_lib/utils.py

@@ -43,7 +43,7 @@ class RequestData:
     logicalResourceId: str
     resourceProperties: Mapping[str, Any]
     systemTags: Mapping[str, Any]
-    stackTags: Mapping[str, Any]
+    stackTags: Optional[Mapping[str, Any]] = None
     callerCredentials: Optional[Credentials] = field(default=None)
     providerCredentials: Optional[Credentials] = field(default=None)
     previousResourceProperties: Optional[Mapping[str, Any]] = None

tests/lib/callback_test.py

@@ -0,0 +1,61 @@
+# pylint: disable=redefined-outer-name,protected-access
+from unittest.mock import Mock, patch
+from uuid import uuid4
+
+import boto3
+from cloudformation_cli_python_lib.callback import report_progress
+from cloudformation_cli_python_lib.interface import (
+    BaseResourceModel,
+    HandlerErrorCode,
+    OperationStatus,
+)
+
+
+class MockSession:
+    def __init__(self):
+        self._cfn = Mock(boto3.client("cloudformation"), autospec=True)
+        self._cfn.record_handler_progress.return_value = {
+            "ResponseMetadata": {"RequestId": "mock_request"}
+        }
+
+    def client(self, _name):
+        return self._cfn
+
+
+def test_report_progress_minimal():
+    session = MockSession()
+    uuid = uuid4()
+    with patch("cloudformation_cli_python_lib.callback.uuid4", return_value=uuid):
+        report_progress(
+            session, "123", None, OperationStatus.IN_PROGRESS, None, None, ""
+        )
+    session._cfn.record_handler_progress.assert_called_once_with(
+        BearerToken="123",
+        OperationStatus="IN_PROGRESS",
+        StatusMessage="",
+        ClientRequestToken=str(uuid),
+    )
+
+
+def test_report_progress_full():
+    session = MockSession()
+    uuid = uuid4()
+    with patch("cloudformation_cli_python_lib.callback.uuid4", return_value=uuid):
+        report_progress(
+            session,
+            "123",
+            HandlerErrorCode.InternalFailure,
+            OperationStatus.FAILED,
+            OperationStatus.IN_PROGRESS,
+            BaseResourceModel(),
+            "test message",
+        )
+    session._cfn.record_handler_progress.assert_called_once_with(
+        BearerToken="123",
+        OperationStatus="FAILED",
+        CurrentOperationStatus="IN_PROGRESS",
+        StatusMessage="test message",
+        ResourceModel="{}",
+        ErrorCode="InternalFailure",
+        ClientRequestToken=str(uuid),
+    )

tests/lib/interface_test.py

@@ -1,5 +1,6 @@
-# pylint: disable=protected-access,redefined-outer-name
+# pylint: disable=protected-access,redefined-outer-name,abstract-method
 import json
+from dataclasses import dataclass
 from string import ascii_letters
 
 import boto3
@@ -26,6 +27,12 @@ def client():
     )
 
 
+@dataclass
+class TestModel(BaseResourceModel):
+    somekey: str
+    someotherkey: str
+
+
 def test_base_resource_model__deserialize():
     with pytest.raises(NotImplementedError):
         BaseResourceModel()._deserialize({})
@@ -52,16 +59,48 @@ def test_progress_event_failed_is_json_serializable(error_code, message):
 
 
 @given(s.text(ascii_letters), s.text(ascii_letters))
-def test_progress_event_serialize_to_response(message, bearer_token):
+def test_progress_event_serialize_to_response_with_context(message, bearer_token):
+    event = ProgressEvent(
+        status=OperationStatus.SUCCESS, message=message, callbackContext={"a": "b"}
+    )
+
+    assert event._serialize(to_response=True, bearer_token=bearer_token) == {
+        "operationStatus": OperationStatus.SUCCESS.name,  # pylint: disable=no-member
+        "message": message,
+        "bearerToken": bearer_token,
+    }
+
+
+@given(s.text(ascii_letters), s.text(ascii_letters))
+def test_progress_event_serialize_to_response_with_model(message, bearer_token):
+    model = TestModel("a", "b")
+    event = ProgressEvent(
+        status=OperationStatus.SUCCESS, message=message, resourceModel=model
+    )
+
+    assert event._serialize(to_response=True, bearer_token=bearer_token) == {
+        "operationStatus": OperationStatus.SUCCESS.name,  # pylint: disable=no-member
+        "message": message,
+        "bearerToken": bearer_token,
+        "resourceModel": {"somekey": "a", "someotherkey": "b"},
+    }
+
+
+@given(s.text(ascii_letters), s.text(ascii_letters))
+def test_progress_event_serialize_to_response_with_models(message, bearer_token):
+    models = [TestModel("a", "b"), TestModel("c", "d")]
     event = ProgressEvent(
-        status=OperationStatus.SUCCESS, message=message, callbackDelaySeconds=1
+        status=OperationStatus.SUCCESS, message=message, resourceModels=models
     )
 
     assert event._serialize(to_response=True, bearer_token=bearer_token) == {
-        "operationStatus": OperationStatus.SUCCESS.value,
+        "operationStatus": OperationStatus.SUCCESS.name,  # pylint: disable=no-member
         "message": message,
         "bearerToken": bearer_token,
-        "callbackDelaySeconds": 1,
+        "resourceModels": [
+            {"somekey": "a", "someotherkey": "b"},
+            {"somekey": "c", "someotherkey": "d"},
+        ],
     }