From df4381ba6a346dca306378c7bb9a1f1d7cf3b2c1 Mon Sep 17 00:00:00 2001
From: Onur Yelekci <onur.yelekci@digital.homeoffice.gov.uk>
Date: Thu, 26 Dec 2019 11:14:52 +0000
Subject: [PATCH] Add db cluster support

---
 lib/rds_helper.rb    | 26 ++++++++++++++++++++++++++
 lib/rds_scheduler.rb | 14 ++++++++++----
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/lib/rds_helper.rb b/lib/rds_helper.rb
index 82e3bba..8c2e182 100755
--- a/lib/rds_helper.rb
+++ b/lib/rds_helper.rb
@@ -14,6 +14,10 @@ def db_instances
     @rds_client.describe_db_instances.db_instances
   end
 
+  def db_clusters
+    @rds_client.describe_db_clusters.db_clusters
+  end
+
   def db_tags(db_arn)
     @rds_client.list_tags_for_resource(
       resource_name: db_arn
@@ -41,4 +45,26 @@ def stop_db_instance(db_name, db_status, db_schedule, schedule_type)
       @logger.warn "DB Instance '#{db_name}' is not in a running state, not taking action (status: #{db_status}, #{schedule_type} schedule: '#{db_schedule}')"
     end
   end
+
+  def start_db_cluster(db_cluster_identifier, db_status, db_schedule, schedule_type)
+    if db_status.eql? 'stopped'
+      @logger.info "Starting DB Cluster '#{db_cluster_identifier}' (#{schedule_type} schedule: '#{db_schedule}', dry_run: #{@dry_run})"
+      @rds_client.start_db_cluster(db_cluster_identifier: db_cluster_identifier) unless @dry_run
+    elsif db_status.eql? 'available'
+      @logger.info "DB Cluster '#{db_cluster_identifier}' is currently available (#{schedule_type} schedule: '#{db_schedule}')"
+    else
+      @logger.warn "DB Cluster '#{db_cluster_identifier}' is not in a stopped state, not taking action (status: #{db_status}, #{schedule_type} schedule: '#{db_schedule}')"
+    end
+  end
+
+  def stop_db_cluster(db_cluster_identifier, db_status, db_schedule, schedule_type)
+    if %w[stopping stopped].include?(db_status)
+      @logger.info "DB Cluster '#{db_cluster_identifier}' is currently stopped (status: #{db_status}, #{schedule_type} schedule: '#{db_schedule}')"
+    elsif db_status.eql? 'available'
+      @logger.info "Stopping DB Cluster '#{db_cluster_identifier}' (#{schedule_type} schedule: '#{db_schedule}', dry_run: #{@dry_run})"
+      @rds_client.stop_db_cluster(db_cluster_identifier: db_cluster_identifier) unless @dry_run
+    else
+      @logger.warn "DB Cluster '#{db_cluster_identifier}' is not in a running state, not taking action (status: #{db_status}, #{schedule_type} schedule: '#{db_schedule}')"
+    end
+  end
 end
diff --git a/lib/rds_scheduler.rb b/lib/rds_scheduler.rb
index 56f2b1a..7ac3559 100755
--- a/lib/rds_scheduler.rb
+++ b/lib/rds_scheduler.rb
@@ -38,7 +38,7 @@ def execute
           break
         end
 
-        process_schedule(rds.db_instance_identifier, rds.db_instance_status, db_schedule, downtime_schedule)
+        process_schedule(rds.db_cluster_identifier, rds.db_instance_identifier, rds.db_instance_status, db_schedule, downtime_schedule)
       end
 
       break if @run_once
@@ -48,7 +48,7 @@ def execute
     end
   end
 
-  def process_schedule(db_name, db_status, db_schedule, downtime_schedule)
+  def process_schedule(db_cluster_identifier, db_name, db_status, db_schedule, downtime_schedule)
     if db_schedule
       begin
         parsed_schedule = @time_parser.parse_schedule(db_schedule)
@@ -58,9 +58,15 @@ def process_schedule(db_name, db_status, db_schedule, downtime_schedule)
         schedule_type = downtime_schedule ? 'downtime' : 'uptime'
         begin
           if (@time_parser.schedule_active?(parsed_schedule) && !downtime_schedule) || (!@time_parser.schedule_active?(parsed_schedule) && downtime_schedule)
-            @rds_client.start_db_instance(db_name, db_status, db_schedule, schedule_type)
-          else
+            if db_cluster_identifier.nil?
+              @rds_client.start_db_instance(db_name, db_status, db_schedule, schedule_type)
+            else
+              @rds_client.start_db_cluster(db_cluster_identifier, db_status, db_schedule, schedule_type)
+            end
+          elsif db_cluster_identifier.nil?
             @rds_client.stop_db_instance(db_name, db_status, db_schedule, schedule_type)
+          else
+            @rds_client.stop_db_cluster(db_cluster_identifier, db_status, db_schedule, schedule_type)
           end
         rescue TimeScheduleParser::TimezoneInvalid => e
           @logger.error "Error processing Time Schedule for DB Instance '#{db_name}': #{e.message}"