chef.rb in capistrano-chef-0.0.3

- old
+ new
@@ -1,17 +1,61 @@
+require 'capistrano'
 require 'chef/knife'
+require 'chef/data_bag_item'
 require 'chef/search/query'
 
-Capistrano::Configuration.instance.load do
-  Chef::Knife.new.configure_chef
+module Capistrano::Chef
+  # Set up chef configuration
+  def self.configure_chef
+    knife = Chef::Knife.new
+    # If you don't do this it gets thrown into debug mode
+    knife.config = { :verbosity => 1 }
+    knife.configure_chef
+  end
 
-  # Define a role for capistrano, but instead of a list of addresses, use a chef
-  # query to search nodes.
-  def chef_role(name, query = "*:*", options = {})
-    # TODO: This can only get a node's top-level attributes. Make it get nested
-    # ones.
-    attr = options.delete(:attribute) || :ipaddress
-    nodes = Chef::Search::Query.new.search(:node, query)[0].map {|n| n[attr] }
-    role name, *nodes, options
-    nodes
+  # Do a search on the Chef server and return an attary of the requested
+  # matching attributes
+  def self.search_chef_nodes(query = '*:*', arg = :ipaddress, limit = 1000)
+    search_proc = \
+      case arg
+      when Proc
+        arg
+      when Hash
+        iface, family = arg.keys.first.to_s, arg.values.first.to_s
+        Proc.new do |n|
+          addresses = n["network"]["interfaces"][iface]["addresses"]
+          addresses.select{|address, data| data["family"] == family }.keys.first
+        end
+      when Symbol, String
+        Proc.new{|n| n[arg.to_s]}
+      else
+        raise ArgumentError, 'Search arguments must be Proc, Hash, Symbol, String.'
+      end
+    Chef::Search::Query.new.search(:node, query, 'X_CHEF_id_CHEF_X asc', 0, limit)[0].map(&search_proc)
   end
+
+  def self.get_data_bag_item(id, data_bag = :apps)
+    Chef::DataBagItem.load(data_bag, id).raw_data
+  end
+
+  # Load into Capistrano
+  def self.load_into(configuration)
+    self.configure_chef
+    configuration.set :capistrano_chef, self
+    configuration.load do
+      def chef_role(name, query = '*:*', options = {})
+        role name, *(capistrano_chef.search_chef_nodes(query, options.delete(:attribute), options.delete(:limit)) + [options])
+      end
+
+      def set_from_data_bag(data_bag = :apps)
+        raise ':application must be set' if fetch(:application).nil?
+        capistrano_chef.get_data_bag_item(application, data_bag).each do |k, v|
+          set k, v
+        end
+      end
+    end
+  end
+end
+
+if Capistrano::Configuration.instance
+  Capistrano::Chef.load_into(Capistrano::Configuration.instance)
 end