From d7992a09d8979d3a37f5cdd42cfba06931cc4cc4 Mon Sep 17 00:00:00 2001
From: John Ferlito <johnf@inodes.org>
Date: Tue, 3 May 2011 12:10:50 +1000
Subject: [PATCH] Add support for ROLEFILTER environment variable

ROLEFILTER is similar to HOSTFILTER. It basically takes a list of
roles and converts that to a list of hosts and acts as if you called
HOSTFILTER on those hosts.
---
 lib/capistrano/cli/help.txt             |  3 +++
 lib/capistrano/configuration/servers.rb | 21 ++++++++++++++++++---
 test/configuration/servers_test.rb      |  8 ++++++++
 3 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/lib/capistrano/cli/help.txt b/lib/capistrano/cli/help.txt
index 198d7d15..8cf55fb9 100644
--- a/lib/capistrano/cli/help.txt
+++ b/lib/capistrano/cli/help.txt
@@ -74,5 +74,8 @@ The following options are understood:
   <%= color 'HOSTFILTER', :bold %>
 	Execute tasks against this comma-separated list of host, but only if the host has the proper role for the task.
 
+  <%= color 'HOSTROLEFILTER', :bold %>
+	Execute tasks against the hosts in this comma-separated list of roles, but only if the host has the proper role for the task.
+
   <%= color 'ROLES', :bold %>
 	Execute tasks against this comma-separated list of roles.  Hosts which do not have the right roles will be skipped.
diff --git a/lib/capistrano/configuration/servers.rb b/lib/capistrano/configuration/servers.rb
index b32e0e2d..70883663 100644
--- a/lib/capistrano/configuration/servers.rb
+++ b/lib/capistrano/configuration/servers.rb
@@ -26,6 +26,9 @@ module Capistrano
       # Yet additionally, if the HOSTFILTER environment variable is set, it
       # will limit the result to hosts found in that (comma-separated) list.
       #
+      # If the HOSTROLEFILTER environment variable is set, it will limit the
+      # result to hosts found in that (comma-separated) list of roles 
+      #
       # Usage:
       #
       #   # return all known servers
@@ -69,9 +72,21 @@ module Capistrano
     protected
 
       def filter_server_list(servers)
-        return servers unless ENV['HOSTFILTER']
-        filters = ENV['HOSTFILTER'].split(/,/)
-        servers.select { |server| filters.include?(server.host) }
+        return servers unless ENV['HOSTFILTER'] or ENV['HOSTROLEFILTER']
+        if ENV['HOSTFILTER']
+          filters = ENV['HOSTFILTER'].split(/,/)
+          servers.select { |server| filters.include?(server.host) }
+        elsif ENV['HOSTROLEFILTER']
+          filters = ENV['HOSTROLEFILTER'].split(/,/).map do |role|
+            local_roles = roles[role.to_sym]
+            if local_roles.is_a? Array
+              roles[role.to_sym]
+            else
+              roles[role.to_sym].servers
+            end
+          end.flatten
+          servers.select { |server| filters.include?(server) }
+        end
       end
 
       def server_list_from(hosts)
diff --git a/test/configuration/servers_test.rb b/test/configuration/servers_test.rb
index 6cfc404b..b14f6eaa 100644
--- a/test/configuration/servers_test.rb
+++ b/test/configuration/servers_test.rb
@@ -131,6 +131,14 @@ class ConfigurationServersTest < Test::Unit::TestCase
     ENV.delete('HOSTFILTER')
   end
 
+  def test_task_with_hostrolefilter_environment_variable_should_apply_only_to_those_hosts
+    ENV['HOSTROLEFILTER'] = "web"
+    task = new_task(:testing)
+    assert_equal %w(web1 web2).sort, @config.find_servers_for_task(task).map { |s| s.host }.sort
+  ensure
+    ENV.delete('HOSTROLEFILTER')
+  end
+
   def test_task_with_only_should_apply_only_to_matching_tasks
     task = new_task(:testing, @config, :roles => :app, :only => { :primary => true })
     assert_equal %w(app1), @config.find_servers_for_task(task).map { |s| s.host }