lib/capistrano/consul.rb in capistrano3-consul-0.3.5 vs lib/capistrano/consul.rb in capistrano3-consul-0.3.6

@@ -4,13 +4,15 @@
 
 module Capistrano
   module Consul
     def self.setup
       return if @url
+
       @url = fetch(:consul_url)
       @token = fetch(:consul_token)
       return false unless @url
+
       @ssh_gateway = fetch(:consul_ssh_gateway)
       if @ssh_gateway
         @gateway = Net::SSH::Gateway.new(@ssh_gateway[:host], @ssh_gateway[:username] || @ssh_gateway[:user], @ssh_gateway[:options] || {})
         @gateway.open('127.0.0.1', @ssh_gateway[:port], @ssh_gateway[:port])
       end
@@ -26,19 +28,27 @@
 
   module DSL
     def consul_all_nodes(properties = {})
       Consul.setup
       Diplomat::Node.get_all.each_with_index do |node, index|
-        extra_properties = index.zero? ? { primary: true } : {}
-        server(node['Address'], properties.merge(extra_properties))
+        if block_given?
+          yield(node)
+        else
+          extra_properties = index.zero? ? { primary: true } : {}
+          server(node['Address'], properties.merge(extra_properties))
+        end
       end
     end
 
     def consul_service(service_name, properties = {})
       Consul.setup
       Diplomat::Service.get(service_name, :all).each_with_index do |node, index|
-        extra_properties = index.zero? ? { primary: true } : {}
-        server(node['Address'], properties.merge(extra_properties))
+        if block_given?
+          yield(node)
+        else
+          extra_properties = index.zero? ? { primary: true } : {}
+          server(node['Address'], properties.merge(extra_properties))
+        end
       end
     end
   end
 end