Merge pull request #754 from umccr/fix/mover-payload-limit

fix(data-migrate): add structured output to step functions

lib/workload/stateless/stacks/data-migrate/data_mover/cli.py

@@ -5,6 +5,7 @@
 
 from data_mover.data_mover import DataMover
 
+logging.basicConfig()
 logger = logging.getLogger()
 logger.setLevel(logging.DEBUG)
 
@@ -15,6 +16,7 @@ def main():
         sys.exit(0)
     except Exception as e:
         DataMover.send_failure(str(e))
+        logger.error(str(e))
         sys.exit(1)
 
 
@@ -42,7 +44,7 @@ def move(source, destination):
     data_mover = DataMover(source, destination, logger=logger)
     data_mover.sync()
     data_mover.delete()
-    data_mover.send_output()
+    data_mover.send_output(command="move")
 
 
 @cli.command()
@@ -63,7 +65,7 @@ def copy(source, destination):
     """
     data_mover = DataMover(source, destination, logger=logger)
     data_mover.sync()
-    data_mover.send_output()
+    data_mover.send_output(command="copy")
 
 
 if __name__ == "__main__":

lib/workload/stateless/stacks/data-migrate/data_mover/data_mover.py

@@ -2,6 +2,7 @@
 import logging
 import os
 import subprocess
+from typing import Literal
 
 import boto3
 from mypy_boto3_stepfunctions import SFNClient
@@ -26,7 +27,6 @@ def __init__(
         self.repeat = repeat
         self.timeout = timeout
         self.logger = logger
-        self.output = ""
 
     def sync(self):
         """
@@ -45,9 +45,7 @@ def sync(self):
                 stdout=subprocess.PIPE,
                 stderr=subprocess.STDOUT,
             )
-            self.logger.info(str(out.stdout))
-
-            self.output += str(out.stdout) or ""
+            self.logger.info(out.stdout.decode())
 
         if out.stdout != b"":
             raise Exception("failed to sync - non-empty output")
@@ -65,19 +63,24 @@ def delete(self):
             stdout=subprocess.PIPE,
             stderr=subprocess.STDOUT,
         )
-        self.logger.info(str(out.stdout))
-
-        self.output += str(out.stdout) or ""
+        self.logger.info(out.stdout.decode())
 
-    def send_output(self):
+    def send_output(self, command: Literal["copy", "move"] = "copy"):
         """
         Send successful task response with the output.
         """
         task_token = os.getenv("DM_TASK_TOKEN")
         if task_token is not None:
             client: SFNClient = boto3.client("stepfunctions")
             client.send_task_success(
-                taskToken=task_token, output=json.dumps(self.output)
+                taskToken=task_token,
+                output=json.dumps(
+                    {
+                        "command": command,
+                        "source": self.source,
+                        "destination": self.destination,
+                    }
+                ),
             )
 
     @staticmethod