From e76d2d50cc9bdc942d3a0a4106918520562c71ad Mon Sep 17 00:00:00 2001
From: RoseSecurity <michael@rosesecurity.dev>
Date: Thu, 2 Apr 2026 09:39:56 -0400
Subject: [PATCH] feat(mwaa): add worker_replacement_strategy variable

Add support for configuring the worker replacement strategy when
updating MWAA environments. This allows users to choose between
FORCED (immediate replacement) and GRACEFUL (wait up to 12 hours
for tasks to complete) strategies.

- Add worker_replacement_strategy variable to root module
- Add worker_replacement_strategy to complete example
- Default to null in root module to preserve existing behavior
---
 examples/complete/fixtures.us-east-2.tfvars |  2 ++
 examples/complete/main.tf                   |  1 +
 examples/complete/variables.tf              | 12 ++++++++++++
 main.tf                                     |  1 +
 variables.tf                                | 11 +++++++++++
 5 files changed, 27 insertions(+)

diff --git a/examples/complete/fixtures.us-east-2.tfvars b/examples/complete/fixtures.us-east-2.tfvars
index 8a56e4d..a7b2c67 100644
--- a/examples/complete/fixtures.us-east-2.tfvars
+++ b/examples/complete/fixtures.us-east-2.tfvars
@@ -22,6 +22,8 @@ min_workers = 1
 
 max_workers = 10
 
+worker_replacement_strategy = "FORCED"
+
 webserver_access_mode = "PRIVATE_ONLY"
 
 dag_processing_logs_enabled = true
diff --git a/examples/complete/main.tf b/examples/complete/main.tf
index 99c5d1a..dcebe35 100644
--- a/examples/complete/main.tf
+++ b/examples/complete/main.tf
@@ -32,6 +32,7 @@ module "mwaa" {
   environment_class             = var.environment_class
   min_workers                   = var.min_workers
   max_workers                   = var.max_workers
+  worker_replacement_strategy   = var.worker_replacement_strategy
   webserver_access_mode         = var.webserver_access_mode
   dag_processing_logs_enabled   = var.dag_processing_logs_enabled
   dag_processing_logs_level     = var.dag_processing_logs_level
diff --git a/examples/complete/variables.tf b/examples/complete/variables.tf
index 42a3a55..f6f8b41 100644
--- a/examples/complete/variables.tf
+++ b/examples/complete/variables.tf
@@ -104,3 +104,15 @@ variable "min_workers" {
   description = "The minimum number of workers that you want to run in your environment."
   default     = 1
 }
+
+variable "worker_replacement_strategy" {
+  type        = string
+  description = "The worker replacement strategy to use when updating the environment. Valid values: `FORCED`, `GRACEFUL`. `FORCED` means Apache Airflow workers will be stopped and replaced without waiting for tasks to complete before an update. `GRACEFUL` means Apache Airflow workers will be able to complete running tasks for up to 12 hours during an update before being stopped and replaced."
+  default     = "FORCED"
+
+  validation {
+    condition     = contains(["FORCED", "GRACEFUL"], var.worker_replacement_strategy)
+    error_message = "Valid values are: `FORCED` or `GRACEFUL`."
+  }
+}
+
diff --git a/main.tf b/main.tf
index f688755..a06be70 100644
--- a/main.tf
+++ b/main.tf
@@ -223,6 +223,7 @@ resource "aws_mwaa_environment" "default" {
   kms_key                          = var.kms_key
   max_workers                      = var.max_workers
   min_workers                      = var.min_workers
+  worker_replacement_strategy      = var.worker_replacement_strategy
   min_webservers                   = var.environment_class == "mw1.micro" ? 1 : var.min_webservers
   max_webservers                   = var.environment_class == "mw1.micro" ? 1 : var.max_webservers
   schedulers                       = var.schedulers
diff --git a/variables.tf b/variables.tf
index ea61f6c..ca5802c 100644
--- a/variables.tf
+++ b/variables.tf
@@ -81,6 +81,17 @@ variable "min_workers" {
   default     = 1
 }
 
+variable "worker_replacement_strategy" {
+  type        = string
+  description = "The worker replacement strategy to use when updating the environment. Valid values: `FORCED`, `GRACEFUL`. `FORCED` means Apache Airflow workers will be stopped and replaced without waiting for tasks to complete before an update. `GRACEFUL` means Apache Airflow workers will be able to complete running tasks for up to 12 hours during an update before being stopped and replaced."
+  default     = null
+
+  validation {
+    condition     = var.worker_replacement_strategy == null || contains(["FORCED", "GRACEFUL"], var.worker_replacement_strategy)
+    error_message = "Valid values are: `FORCED`, `GRACEFUL`, or `null`."
+  }
+}
+
 variable "max_webservers" {
   type        = number
   description = "The maximum number of web servers that you want to run in your environment."