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
+}