From 44bab6c94ee50fb30f488cb1383eed1a24a9b3d4 Mon Sep 17 00:00:00 2001
From: Ryan Lim <rlim@chanzuckerberg.com>
Date: Thu, 21 Sep 2023 16:16:54 -0700
Subject: [PATCH] add options to turn off step notifications

---
 main.tf                                       |  1 +
 .../sns_notification/sns_notification.py      | 43 ++++++++++---------
 terraform/modules/swipe-sfn/main.tf           |  2 +-
 terraform/modules/swipe-sfn/variables.tf      |  7 +++
 test/terraform/moto/main.tf                   |  4 +-
 variables.tf                                  |  6 +++
 6 files changed, 39 insertions(+), 24 deletions(-)

diff --git a/main.tf b/main.tf
index f9e8687f..a114f275 100644
--- 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
index 37f26d63..da0e9e91 100644
--- 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
index 1b84557f..7c2f291c 100644
--- 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
index c137aea1..8f58e206 100644
--- 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
index 6b974dd0..d6a6a4be 100644
--- 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
index 2e08455e..8b4d1856 100644
--- 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
+}