Add branch specification to email service also

Signed-off-by: Balasankar "Balu" C <balasankar@gitlab.com>

2 files changed, 42 insertions, 9 deletions

diff --git a/app/models/project_services/pipelines_email_service.rb b/app/models/project_services/pipelines_email_service.rb
index ae5d5038099..649c252130a 100644
--- a/app/models/project_services/pipelines_email_service.rb
+++ b/app/models/project_services/pipelines_email_service.rb
@@ -1,12 +1,32 @@
 # frozen_string_literal: true
 
 class PipelinesEmailService < Service
-  prop_accessor :recipients
+  prop_accessor :recipients, :branches_to_be_notified
   boolean_accessor :notify_only_broken_pipelines, :notify_only_default_branch
   validates :recipients, presence: true, if: :valid_recipients?
 
+  BRANCH_CHOICES = [
+    ['All branches', 'all'],
+    ['Default branch', 'default'],
+    ['Protected branches', 'protected'],
+    ['Default branch and protected branches', 'default_and_protected']
+  ].freeze
+
   def initialize_properties
-    self.properties ||= { notify_only_broken_pipelines: true, notify_only_default_branch: false }
+    if properties.nil?
+      self.properties = {}
+      self.notify_only_broken_pipelines = true
+      self.branches_to_be_notified = "default"
+    elsif !self.notify_only_default_branch.nil? && self.branches_to_be_notified.nil?
+      # In older versions, there was only a boolean property named
+      # `notify_only_default_branch`. Now we have a string property named
+      # `branches_to_be_notified`. Instead of doing a background migration, we
+      # opted to set a value for the new property based on the old one, if
+      # users hasn't specified one already. When users edit the service and
+      # selects a value for this new property, it will override everything.
+
+      self.branches_to_be_notified = notify_only_default_branch? ? "default" : "all"
+    end
   end
 
   def title
@@ -55,8 +75,9 @@ class PipelinesEmailService < Service
         required: true },
       { type: 'checkbox',
         name: 'notify_only_broken_pipelines' },
-      { type: 'checkbox',
-        name: 'notify_only_default_branch' }
+      { type: 'select',
+        name: 'branches_to_be_notified',
+        choices: BRANCH_CHOICES }
     ]
   end
 
@@ -73,9 +94,21 @@ class PipelinesEmailService < Service
   end
 
   def notify_for_pipeline_branch?(data)
-    return true unless notify_only_default_branch?
+    ref = if data[:ref]
+            Gitlab::Git.ref_name(data[:ref])
+          else
+            data.dig(:object_attributes, :ref)
+          end
 
-    data[:object_attributes][:ref] == data[:project][:default_branch]
+    if branches_to_be_notified == "all"
+      true
+    elsif %w[default default_and_protected].include?(branches_to_be_notified)
+      ref == project.default_branch
+    elsif %w[protected default_and_protected].include?(branches_to_be_notified)
+      project.protected_branches.exists?(name: ref)
+    else
+      false
+    end
   end
 
   def notify_for_pipeline?(data)
diff --git a/lib/api/helpers/services_helpers.rb b/lib/api/helpers/services_helpers.rb
index 6c56fb9e751..608b599105d 100644
--- a/lib/api/helpers/services_helpers.rb
+++ b/lib/api/helpers/services_helpers.rb
@@ -540,9 +540,9 @@ module API
             },
             {
               required: false,
-              name: :notify_only_default_branch,
-              type: Boolean,
-              desc: 'Send notifications only for the default branch'
+              name: :branches_to_be_notified,
+              type: String,
+              desc: 'Branches for which notifications are to be sent'
             }
           ],
           'pivotaltracker' => [