From 5fa1afc73251ba659b0b4d723ad6a31a94637a1a Mon Sep 17 00:00:00 2001
From: Tarun Menon <tarunmenon95@gmail.com>
Date: Mon, 22 Jul 2024 14:17:04 +1000
Subject: [PATCH] add constraint

---
 ecs-task.cfhighlander.rb | 1 +
 ecs-task.cfndsl.rb       | 5 +++--
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/ecs-task.cfhighlander.rb b/ecs-task.cfhighlander.rb
index c50b912..4dcd8a1 100644
--- a/ecs-task.cfhighlander.rb
+++ b/ecs-task.cfhighlander.rb
@@ -7,6 +7,7 @@
     Parameters do
       ComponentParam 'EnvironmentName', 'dev', isGlobal: true
       ComponentParam 'EnvironmentType', 'development', allowedValues: ['development','production'], isGlobal: true
+      ComponentParam 'EbsAZ', ''
 
       task_definition.each do |task_def, task|
         if task.has_key?('tag_param')
diff --git a/ecs-task.cfndsl.rb b/ecs-task.cfndsl.rb
index c99d5b5..1f65763 100644
--- a/ecs-task.cfndsl.rb
+++ b/ecs-task.cfndsl.rb
@@ -126,7 +126,10 @@
         EC2_Volume(ebs_volume['name']) do
           Size 100
           VolumeType "gp3"
+          AvailabilityZone Ref(:EbsAZ)
         end
+        
+        task_constraints << {Type: "memberOf", Expression: "attribute:ecs.availability-zone in #{Ref(:EbsAZ)}"}
         mount_points << { ContainerPath: ebs_volume['container_path'], SourceVolume: Ref(ebs_volume['name']), ReadOnly: false}
         task_volumes << { Name: Ref(ebs_volume['name']), ConfiguredAtLaunch: true }
         task_def.merge!({MountPoints: mount_points })
@@ -312,8 +315,6 @@
       end
     end
 
-    print("Task Volumes: #{task_volumes}")
-
     task_type = external_parameters.fetch(:task_type, 'EC2')
     unless task_definition.empty?