Adding a fact handler, along with an abstract interface for fact stores and a simple yaml fact store, towards the Node Classification work.

git-svn-id: https://reductivelabs.com/svn/puppet/trunk@2457 980ebf18-57e1-0310-9a29-db15c13687c0

6 files changed, 358 insertions, 0 deletions

diff --git a/lib/puppet/configuration.rb b/lib/puppet/configuration.rb
index 86d7dbfdc..ffd147cd1 100644
--- a/lib/puppet/configuration.rb
+++ b/lib/puppet/configuration.rb
@@ -326,6 +326,16 @@ module Puppet
         :smtpserver => ["none",
             "The server through which to send email reports."]
     )
+
+    self.setdefaults(:facts,
+        :factstore => ["yaml",
+            "The backend store to use for client facts."]
+    )
+
+    self.setdefaults(:yamlfacts,
+        :yamlfactdir => ["$vardir/facts",
+            "The directory in which client facts are stored when the yaml fact store is used."]
+    )
 end
 
 # $Id$
diff --git a/lib/puppet/fact_stores/yaml.rb b/lib/puppet/fact_stores/yaml.rb
new file mode 100644
index 000000000..b4b0c0e7c
--- /dev/null
+++ b/lib/puppet/fact_stores/yaml.rb
@@ -0,0 +1,42 @@
+Puppet::Util::FactStore.newstore(:yaml) do
+    desc "Store client facts as flat files, serialized using YAML."
+
+    # Get a client's facts.
+    def get(node)
+        file = path(node)
+
+        return nil unless FileTest.exists?(file)
+
+        begin
+            facts = YAML::load(File.read(file))
+        rescue => detail
+            Puppet.err "Could not load facts for %s: %s" % [node, detail]
+        end
+        facts
+    end
+
+    def initialize
+        Puppet.config.use(:yamlfacts)
+    end
+
+    # Store the facts to disk.
+    def set(node, facts)
+        File.open(path(node), "w", 0600) do |f|
+            begin
+                f.print YAML::dump(facts)
+            rescue => detail
+                Puppet.err "Could not write facts for %s: %s" % [node, detail]
+            end
+        end
+        nil
+    end
+
+    private
+
+    # Return the path to a given node's file.
+    def path(node)
+        File.join(Puppet[:yamlfactdir], node + ".yaml")
+    end
+end
+
+# $Id$
diff --git a/lib/puppet/network/handler/facts.rb b/lib/puppet/network/handler/facts.rb
new file mode 100755
index 000000000..46c94b91a
--- /dev/null
+++ b/lib/puppet/network/handler/facts.rb
@@ -0,0 +1,67 @@
+require 'yaml'
+require 'puppet/util/fact_store'
+
+class Puppet::Network::Handler
+    # Receive logs from remote hosts.
+    class Facts < Handler
+        @interface = XMLRPC::Service::Interface.new("facts") { |iface|
+            iface.add_method("void set(string, string)")
+            iface.add_method("string get(string)")
+            iface.add_method("integer store_date(string)")
+        }
+
+        def initialize(hash = {})
+            super
+
+            backend = Puppet[:factstore]
+
+            unless klass = Puppet::Util::FactStore.store(backend)
+                raise Puppet::Error, "Could not find fact store %s" % backend
+            end
+
+            @backend = klass.new
+        end
+
+        # Get the facts from our back end.
+        def get(node)
+            if facts = @backend.get(node)
+                return strip_internal(facts)
+            else
+                return nil
+            end
+        end
+
+        # Set the facts in the backend.
+        def set(node, facts)
+            @backend.set(node, add_internal(facts))
+            nil
+        end
+
+        # Retrieve a client's storage date.
+        def store_date(node)
+            if facts = get(node)
+                facts[:_puppet_timestamp].to_i
+            else
+                nil
+            end
+        end
+
+        private
+
+        # Add internal data to the facts for storage.
+        def add_internal(facts)
+            facts = facts.dup
+            facts[:_puppet_timestamp] = Time.now
+            facts
+        end
+
+        # Strip out that internal data.
+        def strip_internal(facts)
+            facts = facts.dup
+            facts.find_all { |name, value| name.to_s =~ /^_puppet_/ }.each { |name, value| facts.delete(name) }
+            facts
+        end
+    end
+end
+
+# $Id$
diff --git a/lib/puppet/util/fact_store.rb b/lib/puppet/util/fact_store.rb
new file mode 100644
index 000000000..f872ebb01
--- /dev/null
+++ b/lib/puppet/util/fact_store.rb
@@ -0,0 +1,60 @@
+# Created on 2007-05-02
+# Copyright Luke Kanies
+
+module Puppet::Util
+    # The abstract base class for client fact storage.
+    class FactStore
+        extend Puppet::Util
+        extend Puppet::Util::Docs
+        extend Puppet::Util::ClassGen
+
+        @loader = Puppet::Util::Autoload.new(self, "puppet/fact_stores")
+        @stores = {}
+
+        # Add a new report type.
+        def self.newstore(name, options = {}, &block)
+            klass = genclass(name,
+                :block => block,
+                :prefix => "FactStore",
+                :hash => @stores,
+                :attributes => options
+            )
+        end
+
+        # Remove a store; really only used for testing.
+        def self.rmstore(name)
+            rmclass(name, :hash => @stores)
+        end
+
+        # Load a store.
+        def self.store(name)
+            name = symbolize(name)
+            unless @stores.include? name
+                if @loader.load(name)
+                    unless @stores.include? name
+                        Puppet.warning(
+                            "Loaded report file for %s but report was not defined" %
+                            name
+                        )
+                        return nil
+                    end
+                else
+                    return nil
+                end
+            end
+            @stores[name]
+        end
+
+        # Retrieve the facts for a node.
+        def get(node)
+            raise Puppet::DevError, "%s has not overridden get" % self.class.name
+        end
+
+        # Set the facts for a node.
+        def set(node, facts)
+            raise Puppet::DevError, "%s has not overridden set" % self.class.name
+        end
+    end
+end
+
+# $Id$
diff --git a/test/network/handler/facts.rb b/test/network/handler/facts.rb
new file mode 100755
index 000000000..03327b8c4
--- /dev/null
+++ b/test/network/handler/facts.rb
@@ -0,0 +1,112 @@
+#!/usr/bin/env ruby
+
+$:.unshift("../../lib") if __FILE__ =~ /\.rb$/
+
+require 'puppettest'
+require 'mocha'
+require 'puppet/network/handler/facts'
+
+class TestFactsHandler < Test::Unit::TestCase
+    include PuppetTest::ServerTest
+
+    def setup
+        super
+
+        @class = Puppet::Network::Handler.handler(:facts)
+
+        @@client_facts = {}
+
+        unless Puppet::Util::FactStore.store(:testing)
+            Puppet::Util::FactStore.newstore(:testing) do
+                def get(node)
+                    @@client_facts[node]
+                end
+
+                def set(node, facts)
+                    @@client_facts[node] = facts
+                end
+            end
+        end
+
+        Puppet[:factstore] = :testing
+
+        @handler = @class.new
+
+        @facts = {:a => :b, :c => :d}
+        @name = "foo"
+
+        @backend = @handler.instance_variable_get("@backend")
+    end
+
+    def teardown
+        @@client_facts.clear
+    end
+
+    def test_strip_internal
+        @facts[:_puppet_one] = "yay"
+        @facts[:_puppet_two] = "boo"
+        @facts[:_puppetthree] = "foo"
+
+        newfacts = nil
+        assert_nothing_raised("Could not call strip_internal") do
+            newfacts = @handler.send(:strip_internal, @facts)
+        end
+
+        [:_puppet_one, :_puppet_two, :_puppetthree].each do |name|
+            assert(@facts.include?(name), "%s was removed in strip_internal from original hash" % name)
+        end
+        [:_puppet_one, :_puppet_two].each do |name|
+            assert(! newfacts.include?(name), "%s was not removed in strip_internal" % name)
+        end
+        assert_equal("foo", newfacts[:_puppetthree], "_puppetthree was removed in strip_internal")
+    end
+
+    def test_add_internal
+        newfacts = nil
+        assert_nothing_raised("Could not call strip_internal") do
+            newfacts = @handler.send(:add_internal, @facts)
+        end
+
+        assert_instance_of(Time, newfacts[:_puppet_timestamp], "Did not set timestamp in add_internal")
+        assert(! @facts.include?(:_puppet_timestamp), "Modified original hash in add_internal")
+    end
+
+    def test_set
+        newfacts = @facts.dup
+        newfacts[:_puppet_timestamp] = Time.now
+        @handler.expects(:add_internal).with(@facts).returns(newfacts)
+        @backend.expects(:set).with(@name, newfacts).returns(nil)
+
+        assert_nothing_raised("Could not set facts") do
+            assert_nil(@handler.set(@name, @facts), "handler.set did not return nil")
+        end
+    end
+
+    def test_get
+        prefacts = @facts.dup
+        prefacts[:_puppet_timestamp] = Time.now
+        @@client_facts[@name] = prefacts
+        @handler.expects(:strip_internal).with(prefacts).returns(@facts)
+        @backend.expects(:get).with(@name).returns(prefacts)
+
+        assert_nothing_raised("Could not retrieve facts") do
+            assert_equal(@facts,  @handler.get(@name), "did not get correct answer from handler.get")
+        end
+
+        @handler = @class.new
+        assert_nothing_raised("Failed to call 'get' with no stored facts") do
+            @handler.get("nosuchname")
+        end
+    end
+
+    def test_store_date
+        time = Time.now
+        @facts[:_puppet_timestamp] = time
+
+        @handler.expects(:get).with(@name).returns(@facts)
+
+        assert_equal(time.to_i, @handler.store_date(@name), "Did not retrieve timestamp correctly")
+    end
+end
+
+# $Id$
diff --git a/test/util/fact_store.rb b/test/util/fact_store.rb
new file mode 100755
index 000000000..5b04d1374
--- /dev/null
+++ b/test/util/fact_store.rb
@@ -0,0 +1,67 @@
+#!/usr/bin/env ruby
+#
+#  Created by Luke Kanies on 2007-05-02.
+#  Copyright (c) 2007. All rights reserved.
+
+$:.unshift("../lib").unshift("../../lib") if __FILE__ =~ /\.rb$/
+
+require 'puppettest'
+require 'puppet/util/fact_store'
+
+class TestFactStore < Test::Unit::TestCase
+	include PuppetTest
+	
+	def test_new_fact_store
+        klass = nil
+        assert_nothing_raised("Could not create fact store") do
+            klass = Puppet::Util::FactStore.newstore(:yay) do
+            end
+        end
+
+        assert_equal(klass, Puppet::Util::FactStore.store(:yay), "Did not get created store back by name")
+    end
+
+    def test_yaml_store
+        yaml = Puppet::Util::FactStore.store(:yaml)
+        assert(yaml, "Could not retrieve yaml store")
+
+        name = "node"
+        facts = {"a" => :b, :c => "d", :e => :f, "g" => "h"}
+
+        store = nil
+        assert_nothing_raised("Could not create YAML store instance") do
+            store = yaml.new 
+        end
+
+        assert_nothing_raised("Could not store host facts") do
+            store.set(name, facts)
+        end
+
+        dir = Puppet[:yamlfactdir]
+
+        file = File.join(dir, name + ".yaml")
+        assert(FileTest.exists?(file), "Did not create yaml file for node")
+
+        text = File.read(file)
+        newfacts = nil
+        assert_nothing_raised("Could not deserialize yaml") do
+            newfacts = YAML::load(text)
+        end
+
+        # Don't directly compare the hashes, because there might be extra
+        # data stored in the client hash
+        facts.each do |var, value|
+            assert_equal(value, newfacts[var], "Value for %s changed during storage" % var)
+        end
+
+        # Now make sure the facts get retrieved correctly
+        assert_nothing_raised("Could not retrieve facts") do
+            newfacts = store.get(name)
+        end
+
+        # Now make sure the hashes are equal, since internal facts should not be returned.
+        assert_equal(facts, newfacts, "Retrieved facts are not equal")
+    end
+end
+
+# $Id$