add options to turn off step notifications

chanzuckerberg · Sep 21, 2023 · 44bab6c · 44bab6c
1 parent e9f58e4
commit 44bab6c
Show file tree

Hide file tree

Showing 6 changed files with 39 additions and 24 deletions.
diff --git a/main.tf b/main.tf
@@ -64,6 +64,7 @@ module "sfn" {
   stage_vcpu_defaults           = var.stage_vcpu_defaults
   extra_env_vars                = var.extra_env_vars
   sqs_queues                    = var.sqs_queues
+  step_notifications            = var.step_notifications
   call_cache                    = var.call_cache
   output_status_json_files      = var.output_status_json_files
   tags                          = var.tags

diff --git a/miniwdl-plugins/sns_notification/sns_notification.py b/miniwdl-plugins/sns_notification/sns_notification.py
@@ -44,27 +44,28 @@ def task(cfg, logger, run_id, run_dir, task, **recv):
 
     log.info(_("sending message to sns"))
 
-    message_attributes = {
-        "WorkflowName": {"DataType": "String", "StringValue": run_id[0]},
-        "TaskName": {"DataType": "String", "StringValue": run_id[-1]},
-        "ExecutionId": {
-            "DataType": "String",
-            "StringValue": "execution_id_to_be_passed_in",
-        },
-    }
-
-    outputs = process_outputs(values_to_json(recv["outputs"]))
-    message_body = {
-        "version": "0",
-        "id": "0",
-        "detail-type": "Step Functions Execution Step Notification",
-        "source": "aws.batch",
-        "account": "",
-        "time": datetime.now().strftime("%Y-%m-%dT%H:%M:%SZ"),
-        "resources": [],
-        "detail": outputs,
-    }
-    send_message(message_attributes, json.dumps(message_body))
+    if topic_arn:
+        message_attributes = {
+            "WorkflowName": {"DataType": "String", "StringValue": run_id[0]},
+            "TaskName": {"DataType": "String", "StringValue": run_id[-1]},
+            "ExecutionId": {
+                "DataType": "String",
+                "StringValue": "execution_id_to_be_passed_in",
+            },
+        }
+
+        outputs = process_outputs(values_to_json(recv["outputs"]))
+        message_body = {
+            "version": "0",
+            "id": "0",
+            "detail-type": "Step Functions Execution Step Notification",
+            "source": "aws.batch",
+            "account": "",
+            "time": datetime.now().strftime("%Y-%m-%dT%H:%M:%SZ"),
+            "resources": [],
+            "detail": outputs,
+        }
+        send_message(message_attributes, json.dumps(message_body))
 
     yield recv
 

diff --git a/terraform/modules/swipe-sfn/main.tf b/terraform/modules/swipe-sfn/main.tf
@@ -44,7 +44,7 @@ module "batch_job" {
   docker_network             = var.docker_network
   call_cache                 = var.call_cache
   output_status_json_files   = var.output_status_json_files
-  sfn_notification_topic_arn = aws_sns_topic.sfn_notifications_topic[0].arn
+  sfn_notification_topic_arn = length(var.sqs_queues) > 0 && var.step_notifications ? aws_sns_topic.sfn_notifications_topic[0].arn : ""
   tags                       = var.tags
 }
 

diff --git a/terraform/modules/swipe-sfn/variables.tf b/terraform/modules/swipe-sfn/variables.tf
@@ -116,3 +116,10 @@ variable "metrics_schedule" {
   type        = string
   default     = "rate(1 minute)"
 }
+
+variable "step_notifications" {
+  description = "Boolean to determine whether or not to use send step notifications with SNS"
+  type        = bool
+  default     = false
+}
+
diff --git a/test/terraform/moto/main.tf b/test/terraform/moto/main.tf
@@ -32,7 +32,7 @@ module "swipetest" {
     "Two" : { "spot" : 12800, "on_demand" : 256000 },
   }
 
-  workspace_s3_prefixes = ["swipe-test"]
-
+  workspace_s3_prefixes    = ["swipe-test"]
   output_status_json_files = true
+  step_notifications       = true
 }
diff --git a/variables.tf b/variables.tf
@@ -194,3 +194,9 @@ variable "output_status_json_files" {
   type        = bool
   default     = false
 }
+
+variable "step_notifications" {
+  description = "Boolean to determine whether or not to use send step notifications with SNS"
+  type        = bool
+  default     = false
+}