Permalink
Browse files

Add Puppet Master MongoDB integration files

  • Loading branch information...
1 parent 1b97e54 commit 7601265c1e07539124e0a02c4a019a3d9d641e79 @ripienaar ripienaar committed Jul 27, 2010
@@ -0,0 +1,84 @@
+require 'rubygems'
+require 'mongo'
+
+class Puppet::Util::MongoQuery
+ include Singleton
+
+ def initialize
+ @connected = false
+
+ @host = nil
+ @db = nil
+ @collection = nil
+ end
+
+ # Returns just fqdn matching the query
+ def find_node_names(query)
+ query(query, {:fields => ["fqdn"]}).map{|result| result["fqdn"]}
+ end
+
+ # Finds all about nodes matching query
+ def find_nodes(query)
+ query(query).to_a
+ end
+
+ def setup(host, db, collection)
+ unless (host == @host && db == @db && collection == @collection && connected?)
+ @host = host
+ @db = db
+ @collection = collection
+
+ connect
+ end
+ end
+
+ private
+ # Connect to mongodb
+ def connect
+ begin
+ # Only reconnect if we see a new set of connection params and we arent
+ # already connected to those
+ unless (@host.nil? || @db.nil? || @collection.nil?)
+ Puppet.notice("Creating new mongo db handle: host=#{@host} db=#{@db} collection=#{@collection}")
+ @dbh = Mongo::Connection.new(@host).db(@db)
+ @coll =@dbh.collection(@collection)
+
+ @connected = true
+ else
+ raise Puppet::ParseError, "Cannot connect to mongo db until setup was called"
+ end
+ rescue Exception => e
+ raise Puppet::ParseError, "Failed to connect to mongo db host=#{@host} db=#{@db} collection=#{@collection}: #{e}"
+ end
+ end
+
+ # Checks if we're connected, reconnect if needed
+ def query(q, opts={})
+ tries = 0
+
+ begin
+ connect unless connected?
+
+ @coll.find(q, opts)
+ rescue
+ retries += 1
+
+ if retries == 5
+ raise Puppet::ParseError, "Failed to query Mongo after 5 attempts"
+ end
+
+ connect
+ retry
+ end
+ end
+
+ # Are we connected?
+ def connected?
+ begin
+ return true if @connected && @dbh.connection.connected?
+ rescue Exception => e
+ Puppet.notice("Not connected to mongo db: #{e}")
+ return false
+ end
+ end
+end
@@ -0,0 +1,8 @@
+module Puppet::Parser::Functions
+ newfunction(:search_nodes, :type => :rvalue) do |args|
+ require 'puppet/util/mongoquery'
+
+ searchspec = eval(args[0])
+ Puppet::Util::MongoQuery.instance.find_node_names(searchspec)
+ end
+end
@@ -0,0 +1,15 @@
+module Puppet::Parser::Functions
+ newfunction(:search_setup) do |args|
+ require 'puppet/util/mongoquery'
+
+ unless args.size == 3
+ raise Puppet::ParseError, "search_setup requires host, db and collection arguments"
+ end
+
+ host = args[0]
+ db = args[1]
+ collection = args[2]
+
+ Puppet::Util::MongoQuery.instance.setup(host, db, collection)
+ end
+end

0 comments on commit 7601265

Please sign in to comment.